Merge pull request #440 from kroma-network/feat/improve-prime-field-d…

…e-serialization feat!(math): improve prime field (de)serialization
kroma-network · Jun 19, 2024 · f2ede4e · f2ede4e
2 parents 2e2c671 + f76714d
commit f2ede4e
Show file tree

Hide file tree

Showing 15 changed files with 1,144 additions and 2,485 deletions.
diff --git a/tachyon/crypto/commitments/test/bn254_kzg_polynomial_openings.json b/tachyon/crypto/commitments/test/bn254_kzg_polynomial_openings.json
diff --git a/tachyon/math/elliptic_curves/short_weierstrass/affine_point_unittest.cc b/tachyon/math/elliptic_curves/short_weierstrass/affine_point_unittest.cc
@@ -177,7 +177,7 @@ TEST_F(AffinePointTest, Copyable) {
 
 TEST_F(AffinePointTest, JsonValueConverter) {
   test::AffinePoint expected_point(GF7(1), GF7(2));
-  std::string expected_json = R"({"x":{"value":"0x1"},"y":{"value":"0x2"}})";
+  std::string expected_json = R"({"x":1,"y":2})";
 
   test::AffinePoint p;
   std::string error;

diff --git a/tachyon/math/elliptic_curves/short_weierstrass/jacobian_point_unittest.cc b/tachyon/math/elliptic_curves/short_weierstrass/jacobian_point_unittest.cc
@@ -253,8 +253,7 @@ TEST_F(JacobianPointTest, Copyable) {
 
 TEST_F(JacobianPointTest, JsonValueConverter) {
   test::JacobianPoint expected_point(GF7(1), GF7(2), GF7(3));
-  std::string expected_json =
-      R"({"x":{"value":"0x1"},"y":{"value":"0x2"},"z":{"value":"0x3"}})";
+  std::string expected_json = R"({"x":1,"y":2,"z":3})";
 
   test::JacobianPoint p;
   std::string error;

diff --git a/tachyon/math/elliptic_curves/short_weierstrass/point_xyzz_unittest.cc b/tachyon/math/elliptic_curves/short_weierstrass/point_xyzz_unittest.cc
@@ -256,8 +256,7 @@ TEST_F(PointXYZZTest, Copyable) {
 
 TEST_F(PointXYZZTest, JsonValueConverter) {
   test::PointXYZZ expected_point(GF7(1), GF7(2), GF7(3), GF7(4));
-  std::string expected_json =
-      R"({"x":{"value":"0x1"},"y":{"value":"0x2"},"zz":{"value":"0x3"},"zzz":{"value":"0x4"}})";
+  std::string expected_json = R"({"x":1,"y":2,"zz":3,"zzz":4})";
 
   test::PointXYZZ p;
   std::string error;

diff --git a/tachyon/math/elliptic_curves/short_weierstrass/projective_point_unittest.cc b/tachyon/math/elliptic_curves/short_weierstrass/projective_point_unittest.cc
@@ -256,8 +256,7 @@ TEST_F(ProjectivePointTest, Copyable) {
 
 TEST_F(ProjectivePointTest, JsonValueConverter) {
   test::ProjectivePoint expected_point(GF7(1), GF7(2), GF7(3));
-  std::string expected_json =
-      R"({"x":{"value":"0x1"},"y":{"value":"0x2"},"z":{"value":"0x3"}})";
+  std::string expected_json = R"({"x":1,"y":2,"z":3})";
 
   test::ProjectivePoint p;
   std::string error;

diff --git a/tachyon/math/finite_fields/cubic_extension_field_unittest.cc b/tachyon/math/finite_fields/cubic_extension_field_unittest.cc
@@ -181,8 +181,7 @@ TEST_F(CubicExtensionFieldTest, MultiplicativeGroupOperators) {
 
 TEST_F(CubicExtensionFieldTest, JsonValueConverter) {
   GF7_3 expected_point(GF7(1), GF7(2), GF7(3));
-  std::string expected_json =
-      R"({"c0":{"value":"0x1"},"c1":{"value":"0x2"},"c2":{"value":"0x3"}})";
+  std::string expected_json = R"({"c0":1,"c1":2,"c2":3})";
 
   GF7_3 p;
   std::string error;

diff --git a/tachyon/math/finite_fields/prime_field_base.h b/tachyon/math/finite_fields/prime_field_base.h
@@ -169,6 +169,9 @@ template <typename T>
 class Copyable<
     T, std::enable_if_t<std::is_base_of_v<math::PrimeFieldBase<T>, T>>> {
  public:
+  using value_type = typename T::value_type;
+  using BigInt = typename T::BigIntTy;
+
   static bool s_allow_value_greater_than_or_equal_to_modulus;
   static bool s_is_in_montgomery;
 
@@ -178,16 +181,23 @@ class Copyable<
         return buffer->Write(prime_field.value());
       }
     }
-    return buffer->Write(prime_field.ToBigInt());
+    if constexpr (T::Config::kModulusBits <= 32) {
+      if constexpr (T::Config::kUseMontgomery) {
+        return buffer->Write(T::Config::FromMontgomery(prime_field.value()));
+      } else {
+        return buffer->Write(prime_field.value());
+      }
+    } else {
+      return buffer->Write(prime_field.ToBigInt());
+    }
   }
 
   static bool ReadFrom(const ReadOnlyBuffer& buffer, T* prime_field) {
-    using BigInt = typename T::BigIntTy;
-    BigInt v;
+    value_type v;
     if (!buffer.Read(&v)) return false;
     if (s_allow_value_greater_than_or_equal_to_modulus) {
-      if (v >= BigInt(T::Config::kModulus)) {
-        v = v.Mod(BigInt(T::Config::kModulus));
+      if (v >= T::Config::kModulus) {
+        v %= T::Config::kModulus;
       }
     }
     if constexpr (T::Config::kUseMontgomery) {
@@ -196,13 +206,16 @@ class Copyable<
         return true;
       }
     }
-    *prime_field = T::FromBigInt(v);
+    *prime_field = T(v);
     return true;
   }
 
   static size_t EstimateSize(const T& prime_field) {
-    using BigInt = typename T::BigIntTy;
-    return BigInt::kLimbNums * sizeof(uint64_t);
+    if constexpr (T::Config::kModulusBits <= 32) {
+      return sizeof(uint32_t);
+    } else {
+      return BigInt::kLimbNums * sizeof(uint64_t);
+    }
   }
 };
 
@@ -222,30 +235,42 @@ template <typename T>
 class RapidJsonValueConverter<
     T, std::enable_if_t<std::is_base_of_v<math::PrimeFieldBase<T>, T>>> {
  public:
+  using value_type = typename T::value_type;
+  using BigInt = typename T::BigIntTy;
+
   static bool s_allow_value_greater_than_or_equal_to_modulus;
   static bool s_is_in_montgomery;
 
   template <typename Allocator>
   static rapidjson::Value From(const T& value, Allocator& allocator) {
-    rapidjson::Value object(rapidjson::kObjectType);
     if constexpr (T::Config::kUseMontgomery) {
       if (s_is_in_montgomery) {
-        AddJsonElement(object, "value", value.value(), allocator);
-        return object;
+        return RapidJsonValueConverter<value_type>::From(value.value(),
+                                                         allocator);
       }
     }
-    AddJsonElement(object, "value", value.ToBigInt(), allocator);
-    return object;
+    if constexpr (T::Config::kModulusBits <= 32) {
+      if constexpr (T::Config::kUseMontgomery) {
+        return RapidJsonValueConverter<uint32_t>::From(
+            T::Config::FromMontgomery(value.value()), allocator);
+      } else {
+        return RapidJsonValueConverter<uint32_t>::From(value.value(),
+                                                       allocator);
+      }
+    } else {
+      return RapidJsonValueConverter<BigInt>::From(value.ToBigInt(), allocator);
+    }
   }
 
   static bool To(const rapidjson::Value& json_value, std::string_view key,
                  T* value, std::string* error) {
-    using BigInt = typename T::BigIntTy;
-    BigInt v;
-    if (!ParseJsonElement(json_value, "value", &v, error)) return false;
+    value_type v;
+    if (!RapidJsonValueConverter<value_type>::To(json_value, key, &v, error))
+      return false;
+
     if (s_allow_value_greater_than_or_equal_to_modulus) {
-      if (v >= BigInt(T::Config::kModulus)) {
-        v = v.Mod(BigInt(T::Config::kModulus));
+      if (v >= T::Config::kModulus) {
+        v %= T::Config::kModulus;
       }
     }
     if constexpr (T::Config::kUseMontgomery) {
@@ -254,7 +279,7 @@ class RapidJsonValueConverter<
         return true;
       }
     }
-    *value = T::FromBigInt(v);
+    *value = T(v);
     return true;
   }
 };

diff --git a/tachyon/math/finite_fields/prime_field_base_unittest.cc b/tachyon/math/finite_fields/prime_field_base_unittest.cc
@@ -114,11 +114,15 @@ TYPED_TEST(PrimeFieldBaseTest, JsonValueConverter) {
   using F = TypeParam;
 
   F expected_point(3);
-  std::string expected_json = R"({"value":"0x3"})";
+  std::string expected_json = R"("0x3")";
 
   F p;
   std::string error;
-  ASSERT_TRUE(base::ParseJson(expected_json, &p, &error));
+
+  rapidjson::Document document;
+  document.Parse(expected_json.data(), expected_json.length());
+  base::RapidJsonValueConverter<F>::To(document, "", &p, &error);
+
   ASSERT_TRUE(error.empty());
   EXPECT_EQ(p, expected_point);
 

diff --git a/tachyon/math/finite_fields/quadratic_extension_field_unittest.cc b/tachyon/math/finite_fields/quadratic_extension_field_unittest.cc
@@ -193,7 +193,7 @@ TEST(CyclotomicInverseTest, FastCyclotomicInverse) {
 
 TEST_F(QuadraticExtensionFieldTest, JsonValueConverter) {
   GF7_2 expected_point(GF7(1), GF7(2));
-  std::string expected_json = R"({"c0":{"value":"0x1"},"c1":{"value":"0x2"}})";
+  std::string expected_json = R"({"c0":1,"c1":2})";
 
   GF7_2 p;
   std::string error;

diff --git a/tachyon/math/geometry/point2_unittest.cc b/tachyon/math/geometry/point2_unittest.cc
@@ -55,7 +55,7 @@ TEST(Point2Test, Copyable) {
 
 TEST(Point2Test, JsonValueConverter) {
   Point2GF7 expected_point(GF7(1), GF7(2));
-  std::string expected_json = R"({"x":{"value":"0x1"},"y":{"value":"0x2"}})";
+  std::string expected_json = R"({"x":1,"y":2})";
 
   Point2GF7 p;
   std::string error;

diff --git a/tachyon/math/geometry/point3_unittest.cc b/tachyon/math/geometry/point3_unittest.cc
@@ -57,8 +57,7 @@ TEST(Point3Test, Copyable) {
 
 TEST(Point3Test, JsonValueConverter) {
   Point3GF7 expected_point(GF7(1), GF7(2), GF7(3));
-  std::string expected_json =
-      R"({"x":{"value":"0x1"},"y":{"value":"0x2"},"z":{"value":"0x3"}})";
+  std::string expected_json = R"({"x":1,"y":2,"z":3})";
 
   Point3GF7 p;
   std::string error;

diff --git a/tachyon/math/geometry/point4_unittest.cc b/tachyon/math/geometry/point4_unittest.cc
@@ -61,8 +61,7 @@ TEST(Point4Test, Copyable) {
 
 TEST(Point4Test, JsonValueConverter) {
   Point4GF7 expected_point(GF7(1), GF7(2), GF7(3), GF7(4));
-  std::string expected_json =
-      R"({"x":{"value":"0x1"},"y":{"value":"0x2"},"z":{"value":"0x3"},"w":{"value":"0x4"}})";
+  std::string expected_json = R"({"x":1,"y":2,"z":3,"w":4})";
 
   Point4GF7 p;
   std::string error;

diff --git a/tachyon/math/polynomials/univariate/univariate_dense_polynomial_unittest.cc b/tachyon/math/polynomials/univariate/univariate_dense_polynomial_unittest.cc
@@ -399,7 +399,7 @@ TEST_F(UnivariateDensePolynomialTest, Hash) {
 TEST_F(UnivariateDensePolynomialTest, JsonValueConverter) {
   Poly expected_poly(Coeffs({GF7(1), GF7(2), GF7(3), GF7(4), GF7(5)}));
   std::string expected_json =
-      R"({"coefficients":{"coefficients":[{"value":"0x1"},{"value":"0x2"},{"value":"0x3"},{"value":"0x4"},{"value":"0x5"}]}})";
+      R"({"coefficients":{"coefficients":[1,2,3,4,5]}})";
 
   Poly poly;
   std::string error;

diff --git a/tachyon/math/polynomials/univariate/univariate_evaluations_unittest.cc b/tachyon/math/polynomials/univariate/univariate_evaluations_unittest.cc
@@ -285,8 +285,7 @@ TEST_F(UnivariateEvaluationsTest, Hash) {
 
 TEST_F(UnivariateEvaluationsTest, JsonValueConverter) {
   Poly expected_poly({GF7(1), GF7(2), GF7(3), GF7(4), GF7(5)});
-  std::string expected_json =
-      R"({"evaluations":[{"value":"0x1"},{"value":"0x2"},{"value":"0x3"},{"value":"0x4"},{"value":"0x5"}]})";
+  std::string expected_json = R"({"evaluations":[1,2,3,4,5]})";
 
   Poly poly;
   std::string error;

diff --git a/tachyon/math/polynomials/univariate/univariate_sparse_polynomial_unittest.cc b/tachyon/math/polynomials/univariate/univariate_sparse_polynomial_unittest.cc
@@ -451,7 +451,7 @@ TEST_F(UnivariateSparsePolynomialTest, Hash) {
 TEST_F(UnivariateSparsePolynomialTest, JsonValueConverter) {
   Poly expected_poly(Coeffs({{0, GF7(1)}, {1, GF7(2)}, {4, GF7(3)}}));
   std::string expected_json =
-      R"({"coefficients":{"terms":[{"degree":0,"coefficient":{"value":"0x1"}},{"degree":1,"coefficient":{"value":"0x2"}},{"degree":4,"coefficient":{"value":"0x3"}}]}})";
+      R"({"coefficients":{"terms":[{"degree":0,"coefficient":1},{"degree":1,"coefficient":2},{"degree":4,"coefficient":3}]}})";
 
   Poly poly;
   std::string error;