Merge pull request #411 from kroma-network/feat/prepare-poseidon2

feat(math): prepare poseidon2

tachyon/math/base/BUILD.bazel

@@ -141,6 +141,7 @@ tachyon_cc_unittest(
         ":sign",
         "//tachyon/base/buffer:vector_buffer",
         "//tachyon/base/containers:container_util",
+        "//tachyon/base/strings:string_number_conversions",
         "//tachyon/math/elliptic_curves/msm/test:variable_base_msm_test_set",
         "//tachyon/math/elliptic_curves/short_weierstrass/test:sw_curve_config",
         "//tachyon/math/finite_fields/test:finite_field_test",

tachyon/math/base/field.h

@@ -1,6 +1,7 @@
 #ifndef TACHYON_MATH_BASE_FIELD_H_
 #define TACHYON_MATH_BASE_FIELD_H_
 
+#include <ostream>
 #include <utility>
 
 #include "tachyon/math/base/ring.h"
@@ -32,6 +33,12 @@ class Field : public AdditiveGroup<F>, public MultiplicativeGroup<F> {
   }
 };
 
+template <typename F>
+std::ostream& operator<<(std::ostream& os, const Field<F>& f) {
+  const F& derived = static_cast<const F&>(f);
+  return os << derived.ToString();
+}
+
 }  // namespace tachyon::math
 
 #endif  // TACHYON_MATH_BASE_FIELD_H_

tachyon/math/base/groups.h

@@ -2,6 +2,7 @@
 #define TACHYON_MATH_BASE_GROUPS_H_
 
 #include <limits>
+#include <ostream>
 #include <tuple>
 #include <utility>
 #include <vector>
@@ -158,6 +159,12 @@ class MultiplicativeGroup : public MultiplicativeSemigroup<G> {
   }
 };
 
+template <typename G>
+std::ostream& operator<<(std::ostream& os, const MultiplicativeGroup<G>& g) {
+  const G& derived = static_cast<const G&>(g);
+  return os << derived.ToString();
+}
+
 // AdditiveGroup is a group with the group operation '+'.
 // AdditiveGroup supports subtraction and negation, inheriting the
 // properties of AdditiveSemigroup.
@@ -203,6 +210,12 @@ class AdditiveGroup : public AdditiveSemigroup<G> {
   }
 };
 
+template <typename G>
+std::ostream& operator<<(std::ostream& os, const AdditiveGroup<G>& g) {
+  const G& derived = static_cast<const G&>(g);
+  return os << derived.ToString();
+}
+
 }  // namespace tachyon::math
 
 #endif  // TACHYON_MATH_BASE_GROUPS_H_

tachyon/math/base/groups_unittest.cc

@@ -1,9 +1,12 @@
 #include "tachyon/math/base/groups.h"
 
+#include <string>
+
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 #include "tachyon/base/containers/container_util.h"
+#include "tachyon/base/strings/string_number_conversions.h"
 #include "tachyon/math/finite_fields/test/gf7.h"
 
 namespace tachyon::math {
@@ -20,6 +23,8 @@ TEST(GroupsTest, Div) {
 
     bool operator==(const Int& other) const { return value_ == other.value_; }
 
+    std::string ToString() const { return base::NumberToString(value_); }
+
    private:
     int value_ = 0;
   };
@@ -44,6 +49,8 @@ TEST(GroupsTest, InverseOverride) {
       return denominator_ == other.denominator_;
     }
 
+    std::string ToString() const { return base::NumberToString(denominator_); }
+
    private:
     int denominator_ = 0;
   };
@@ -104,6 +111,8 @@ TEST(GroupsTest, Sub) {
 
     bool operator==(const Int& other) const { return value_ == other.value_; }
 
+    std::string ToString() const { return base::NumberToString(value_); }
+
    private:
     int value_ = 0;
   };
@@ -130,6 +139,8 @@ TEST(GroupsTest, SubOverAdd) {
 
     bool operator==(const Int& other) const { return value_ == other.value_; }
 
+    std::string ToString() const { return base::NumberToString(value_); }
+
    private:
     int value_ = 0;
   };

tachyon/math/matrix/BUILD.bazel

@@ -34,6 +34,7 @@ tachyon_cc_unittest(
     srcs = ["matrix_types_unittest.cc"],
     deps = [
         ":matrix_types",
+        ":prime_field_num_traits",
         "//tachyon/base/buffer:vector_buffer",
         "//tachyon/math/finite_fields/test:finite_field_test",
         "//tachyon/math/finite_fields/test:gf7",

tachyon/math/matrix/matrix_types.h

@@ -10,26 +10,30 @@
 namespace tachyon {
 namespace math {
 
-template <typename PrimeField>
-using Matrix = Eigen::Matrix<PrimeField, Eigen::Dynamic, Eigen::Dynamic>;
+template <typename Field, int Rows = Eigen::Dynamic, int Cols = Eigen::Dynamic,
+          int Options = 0, int MaxRows = Rows, int MaxCols = Cols>
+using Matrix = Eigen::Matrix<Field, Rows, Cols, Options, MaxRows, MaxCols>;
 
-template <typename PrimeField>
-using Vector = Eigen::Matrix<PrimeField, Eigen::Dynamic, 1>;
+template <typename Field, int Size = Eigen::Dynamic, int MaxSize = Size>
+using DiagonalMatrix = Eigen::DiagonalMatrix<Field, Size, MaxSize>;
 
-template <typename PrimeField>
-using RowVector = Eigen::Matrix<PrimeField, 1, Eigen::Dynamic>;
+template <typename Field, int Rows = Eigen::Dynamic, int Options = 0,
+          int MaxRows = Rows>
+using Vector = Eigen::Matrix<Field, Rows, 1, Options, MaxRows, 1>;
+
+template <typename Field, int Cols = Eigen::Dynamic, int Options = 0,
+          int MaxCols = Cols>
+using RowVector = Eigen::Matrix<Field, 1, Cols, Options, 1, MaxCols>;
 
 }  // namespace math
 
 namespace base {
 
-template <typename PrimeField, int Rows, int Cols, int Options, int MaxRows,
+template <typename Field, int Rows, int Cols, int Options, int MaxRows,
           int MaxCols>
-class Copyable<
-    Eigen::Matrix<PrimeField, Rows, Cols, Options, MaxRows, MaxCols>> {
+class Copyable<Eigen::Matrix<Field, Rows, Cols, Options, MaxRows, MaxCols>> {
  public:
-  using Matrix =
-      Eigen::Matrix<PrimeField, Rows, Cols, Options, MaxRows, MaxCols>;
+  using Matrix = Eigen::Matrix<Field, Rows, Cols, Options, MaxRows, MaxCols>;
 
   static bool WriteTo(const Matrix& matrix, Buffer* buffer) {
     if (!buffer->WriteMany(matrix.rows(), matrix.cols())) return false;
@@ -62,7 +66,43 @@ class Copyable<
   }
 
   static size_t EstimateSize(const Matrix& matrix) {
-    return matrix.size() * sizeof(PrimeField) + sizeof(Eigen::Index) * 2;
+    return matrix.size() * sizeof(Field) + sizeof(Eigen::Index) * 2;
+  }
+};
+
+template <typename Field, int Size, int MaxSize>
+class Copyable<Eigen::DiagonalMatrix<Field, Size, MaxSize>> {
+ public:
+  using DiagonalMatrix = Eigen::DiagonalMatrix<Field, Size, MaxSize>;
+  using DiagonalVector = typename DiagonalMatrix::DiagonalVectorType;
+
+  static bool WriteTo(const DiagonalMatrix& matrix, Buffer* buffer) {
+    if (!buffer->WriteMany(matrix.rows())) return false;
+    const DiagonalVector& diagonal = matrix.diagonal();
+    for (Eigen::Index i = 0; i < diagonal.size(); ++i) {
+      if (!buffer->Write(diagonal.data()[i])) return false;
+    }
+    return true;
+  }
+
+  static bool ReadFrom(const ReadOnlyBuffer& buffer, DiagonalMatrix* matrix) {
+    Eigen::Index size;
+    DiagonalVector vector_tmp;
+    if (!buffer.ReadMany(&size)) return false;
+    if (Size != Eigen::Dynamic) {
+      if (size != Size) return false;
+    } else {
+      vector_tmp.resize(size);
+    }
+    for (Eigen::Index i = 0; i < vector_tmp.size(); ++i) {
+      if (!buffer.Read(&vector_tmp.data()[i])) return false;
+    }
+    *matrix = DiagonalMatrix(std::move(vector_tmp));
+    return true;
+  }
+
+  static size_t EstimateSize(const DiagonalMatrix& matrix) {
+    return matrix.rows() * sizeof(Field) + sizeof(Eigen::Index);
   }
 };
 

tachyon/math/matrix/matrix_types_unittest.cc

@@ -3,6 +3,7 @@
 #include "tachyon/base/buffer/vector_buffer.h"
 #include "tachyon/math/finite_fields/test/finite_field_test.h"
 #include "tachyon/math/finite_fields/test/gf7.h"
+#include "tachyon/math/matrix/prime_field_num_traits.h"
 
 namespace tachyon::math {
 
@@ -22,30 +23,30 @@ TEST_F(MatrixTypesTest, CopyableDynamicMatrix) {
 
   {
     write_buf.set_buffer_offset(0);
-    Eigen::Matrix<GF7, 2, 3> value;
+    Matrix<GF7, 2, 3> value;
     ASSERT_FALSE(write_buf.Read(&value));
   }
   {
     write_buf.set_buffer_offset(0);
-    Eigen::Matrix<GF7, 3, 2> value;
+    Matrix<GF7, 3, 2> value;
     ASSERT_FALSE(write_buf.Read(&value));
   }
   {
     write_buf.set_buffer_offset(0);
-    Eigen::Matrix<GF7, 3, 3> value;
+    Matrix<GF7, 3, 3> value;
     ASSERT_TRUE(write_buf.Read(&value));
-    EXPECT_TRUE(value == expected);
+    EXPECT_EQ(value, expected);
   }
   {
     write_buf.set_buffer_offset(0);
     Matrix<GF7> value;
     ASSERT_TRUE(write_buf.Read(&value));
-    EXPECT_TRUE(value == expected);
+    EXPECT_EQ(value, expected);
   }
 }
 
 TEST_F(MatrixTypesTest, Copyable3x3Matrix) {
-  Eigen::Matrix<GF7, 3, 3> expected{
+  Matrix<GF7, 3, 3> expected{
       {GF7(0), GF7(1), GF7(2)},
       {GF7(3), GF7(4), GF7(5)},
       {GF7(6), GF7(0), GF7(1)},
@@ -60,25 +61,79 @@ TEST_F(MatrixTypesTest, Copyable3x3Matrix) {
 
   {
     write_buf.set_buffer_offset(0);
-    Eigen::Matrix<GF7, 2, 3> value;
+    Matrix<GF7, 2, 3> value;
     ASSERT_FALSE(write_buf.Read(&value));
   }
   {
     write_buf.set_buffer_offset(0);
-    Eigen::Matrix<GF7, 3, 2> value;
+    Matrix<GF7, 3, 2> value;
     ASSERT_FALSE(write_buf.Read(&value));
   }
   {
     write_buf.set_buffer_offset(0);
-    Eigen::Matrix<GF7, 3, 3> value;
+    Matrix<GF7, 3, 3> value;
     ASSERT_TRUE(write_buf.Read(&value));
-    EXPECT_TRUE(value == expected);
+    EXPECT_EQ(value, expected);
   }
   {
     write_buf.set_buffer_offset(0);
     Matrix<GF7> value;
     ASSERT_TRUE(write_buf.Read(&value));
-    EXPECT_TRUE(value == expected);
+    EXPECT_EQ(value, expected);
+  }
+}
+
+TEST_F(MatrixTypesTest, CopyableDynamicDiagonalMatrix) {
+  DiagonalMatrix<GF7> expected{{GF7(1), GF7(2), GF7(3)}};
+
+  base::Uint8VectorBuffer write_buf;
+  ASSERT_TRUE(write_buf.Grow(base::EstimateSize(expected)));
+  ASSERT_TRUE(write_buf.Write(expected));
+  ASSERT_TRUE(write_buf.Done());
+
+  {
+    write_buf.set_buffer_offset(0);
+    DiagonalMatrix<GF7, 2> value;
+    ASSERT_FALSE(write_buf.Read(&value));
+  }
+  {
+    write_buf.set_buffer_offset(0);
+    DiagonalMatrix<GF7, 3> value;
+    ASSERT_TRUE(write_buf.Read(&value));
+    EXPECT_EQ(value.diagonal(), expected.diagonal());
+  }
+  {
+    write_buf.set_buffer_offset(0);
+    DiagonalMatrix<GF7> value;
+    ASSERT_TRUE(write_buf.Read(&value));
+    EXPECT_EQ(value.diagonal(), expected.diagonal());
+  }
+}
+
+TEST_F(MatrixTypesTest, Copyable3x3DiagonalMatrix) {
+  DiagonalMatrix<GF7, 3> expected{GF7(1), GF7(2), GF7(3)};
+
+  base::Uint8VectorBuffer write_buf;
+  ASSERT_TRUE(write_buf.Grow(base::EstimateSize(expected)));
+  ASSERT_TRUE(write_buf.Write(expected));
+  ASSERT_TRUE(write_buf.Done());
+
+  {
+    write_buf.set_buffer_offset(0);
+    DiagonalMatrix<GF7, 2> value;
+    ASSERT_FALSE(write_buf.Read(&value));
+  }
+  {
+    write_buf.set_buffer_offset(0);
+    DiagonalMatrix<GF7, 3> value;
+    ASSERT_TRUE(write_buf.Read(&value));
+    EXPECT_EQ(value.diagonal(), expected.diagonal());
+  }
+  {
+    write_buf.set_buffer_offset(0);
+    DiagonalMatrix<GF7> value;
+    ASSERT_TRUE(write_buf.Read(&value));
+    EXPECT_EQ(value.diagonal(), expected.diagonal());
   }
 }
 

tachyon/math/matrix/prime_field_num_traits.h

@@ -1,30 +1,53 @@
 #ifndef TACHYON_MATH_MATRIX_PRIME_FIELD_NUM_TRAITS_H_
 #define TACHYON_MATH_MATRIX_PRIME_FIELD_NUM_TRAITS_H_
 
+#include <type_traits>
+
 #include "third_party/eigen3/Eigen/Core"
 
 #include "tachyon/math/finite_fields/finite_field_forwards.h"
 
 namespace Eigen {
 
 template <typename Config>
-struct NumTraits<tachyon::math::PrimeField<Config>>
-    : GenericNumTraits<tachyon::math::PrimeField<Config>> {
+struct CostCalculator;
+
+template <typename Config>
+struct CostCalculator<tachyon::math::PrimeField<Config>> {
   constexpr static size_t kLimbNums =
       tachyon::math::PrimeField<Config>::kLimbNums;
+  using NumTraitsType =
+      std::conditional_t<Config::kModulusBits <= 32, NumTraits<uint32_t>,
+                         NumTraits<uint64_t>>;
+
+  constexpr static int ComputeReadCost() {
+    return static_cast<int>(kLimbNums * NumTraitsType::ReadCost);
+  }
+  constexpr static int ComputeAddCost() {
+    // In general, c = (a + b) % M = (a + b) > M ? (a + b) - M : (a + b)
+    return static_cast<int>(kLimbNums * NumTraitsType::AddCost * 3 / 2);
+  }
+  constexpr static int ComputeMulCost() {
+    // In general, c = (a * b) % M = (a * b) - [(a * b) / M] * M
+    return static_cast<int>(
+        kLimbNums * (4 * NumTraitsType::MulCost + NumTraitsType::AddCost));
+  }
+};
 
+template <typename Config>
+struct NumTraits<tachyon::math::PrimeField<Config>>
+    : GenericNumTraits<tachyon::math::PrimeField<Config>> {
   enum {
     IsInteger = 1,
     IsSigned = 0,
     IsComplex = 0,
     RequireInitialization = 1,
-    ReadCost = static_cast<int>(kLimbNums * NumTraits<uint64_t>::ReadCost),
-    // In general, c = (a + b) % M = (a + b) > M ? (a + b) - M : (a + b)
+    ReadCost =
+        CostCalculator<tachyon::math::PrimeField<Config>>::ComputeReadCost(),
     AddCost =
-        static_cast<int>(kLimbNums * NumTraits<uint64_t>::AddCost * 3 / 2),
-    // In general, c = (a * b) % M = (a * b) - [(a * b) / M] * M
-    MulCost = static_cast<int>(kLimbNums * (4 * NumTraits<uint64_t>::MulCost +
-                                            NumTraits<uint64_t>::AddCost)),
+        CostCalculator<tachyon::math::PrimeField<Config>>::ComputeAddCost(),
+    MulCost =
+        CostCalculator<tachyon::math::PrimeField<Config>>::ComputeMulCost(),
   };
 };