anemoi: implemeneted separate mds matrix generation function for each…

… allowed dimension (addresses #77 , #65 (comment)) - fixed conflicts after rebase onto anemoi-hash-r1cs
clearmatics · Jan 12, 2023 · d9f917c · d9f917c
1 parent eb1d516
commit d9f917c
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 28 deletions.
diff --git a/libsnark/gadgetlib1/gadgets/hashes/anemoi/anemoi_components.hpp b/libsnark/gadgetlib1/gadgets/hashes/anemoi/anemoi_components.hpp
@@ -183,10 +183,16 @@ class flystel_prime_field_gadget : public gadget<libff::Fr<ppT>>
     void generate_r1cs_witness();
 };
 
-// get the MDS matrix from the number of columns 2,3 or 4
-template<typename FieldT, size_t NumStateColumns_L>
-std::array<std::array<FieldT, NumStateColumns_L>, NumStateColumns_L>
-anemoi_permutation_mds(const FieldT g);
+// get the MDS matrix for each allowed dimension: 2,3 or 4
+template<typename ppT>
+std::array<std::array<libff::Fr<ppT>, 2>, 2> anemoi_permutation_mds_2x2(
+    const libff::Fr<ppT> g);
+template<typename ppT>
+std::array<std::array<libff::Fr<ppT>, 3>, 3> anemoi_permutation_mds_3x3(
+    const libff::Fr<ppT> g);
+template<typename ppT>
+std::array<std::array<libff::Fr<ppT>, 4>, 4> anemoi_permutation_mds_4x4(
+    const libff::Fr<ppT> g);
 
 } // namespace libsnark
 

diff --git a/libsnark/gadgetlib1/gadgets/hashes/anemoi/anemoi_components.tcc b/libsnark/gadgetlib1/gadgets/hashes/anemoi/anemoi_components.tcc
@@ -320,33 +320,38 @@ void flystel_prime_field_gadget<ppT, parameters>::generate_r1cs_witness()
     this->pb.lc_val(output_y1) = input_x1_value - this->pb.val(a1);
 }
 
-template<typename FieldT, size_t NumStateColumns_L>
-std::array<std::array<FieldT, NumStateColumns_L>, NumStateColumns_L>
-anemoi_permutation_mds(const FieldT g)
+template<typename ppT>
+std::array<std::array<libff::Fr<ppT>, 2>, 2> anemoi_permutation_mds_2x2(
+    const libff::Fr<ppT> g)
+{
+    using FieldT = libff::Fr<ppT>;
+    const FieldT g2 = g * g;
+    std::array<std::array<FieldT, 2>, 2> M = {{{1, g}, {g, g2 + 1}}};
+    return M;
+}
+
+template<typename ppT>
+std::array<std::array<libff::Fr<ppT>, 3>, 3> anemoi_permutation_mds_3x3(
+    const libff::Fr<ppT> g)
 {
-    static_assert(
-        (NumStateColumns_L == 2) || (NumStateColumns_L == 3) ||
-            (NumStateColumns_L == 4),
-        "NumStateColumns_L must be 2,3 or 4");
+    using FieldT = libff::Fr<ppT>;
+    std::array<std::array<FieldT, 3>, 3> M = {
+        {{g + 1, 1, g + 1}, {1, 1, g}, {g, 1, 1}}};
+    return M;
+}
 
-    std::array<std::array<FieldT, NumStateColumns_L>, NumStateColumns_L> M;
+template<typename ppT>
+std::array<std::array<libff::Fr<ppT>, 4>, 4> anemoi_permutation_mds_4x4(
+    const libff::Fr<ppT> g)
+{
+    using FieldT = libff::Fr<ppT>;
     const FieldT g2 = g * g;
-    if (NumStateColumns_L == 2) {
-        M = {{1, g}, {g, g2 + 1}};
-        return M;
-    }
-    if (NumStateColumns_L == 3) {
-        M = {{g + 1, 1, g + 1}, {1, 1, g}, {g, 1, 1}};
-        return M;
-    }
-    if (NumStateColumns_L == 4) {
-        M = {
-            {1, 1 + g, g, g},
-            {g2, g + g2, 1 + g, 1 + 2 * g},
-            {g2, g2, 1, 1 + g},
-            {1 + g, 1 + 2 * g, g, 1 + g}};
-        return M;
-    }
+    std::array<std::array<FieldT, 4>, 4> M = {
+        {{1, g + 1, g, g},
+         {g2, g + g2, g + 1, g + g + 1},
+         {g2, g2, 1, g + 1},
+         {g + 1, g + g + 1, g, g + 1}}};
+    return M;
 }
 
 } // namespace libsnark

diff --git a/libsnark/gadgetlib1/gadgets/hashes/anemoi/tests/test_anemoi_gadget.cpp b/libsnark/gadgetlib1/gadgets/hashes/anemoi/tests/test_anemoi_gadget.cpp
@@ -6,6 +6,7 @@
  * @copyright  MIT license (see LICENSE file)
  *****************************************************************************/
 
+#include <array>
 #include <gtest/gtest.h>
 #include <libff/algebra/curves/bls12_381/bls12_381_init.hpp>
 #include <libff/algebra/curves/bls12_381/bls12_381_pp.hpp>
@@ -224,6 +225,34 @@ void test_flystel_prime_field_gadget()
     libff::print_time("flystel_prime_field_gadget tests successful");
 }
 
+template<typename ppT, class parameters = anemoi_parameters<libff::Fr<ppT>>>
+void test_anemoi_permutation_mds()
+{
+    using FieldT = libff::Fr<ppT>;
+    const FieldT g = anemoi_parameters<ppT>::multiplicative_generator_g;
+    {
+        std::array<std::array<FieldT, 2>, 2> M_expect = {{{1, 7}, {7, 50}}};
+        std::array<std::array<FieldT, 2>, 2> M =
+            anemoi_permutation_mds_2x2<ppT>(g);
+        ASSERT_EQ(M, M_expect);
+    }
+    {
+        std::array<std::array<FieldT, 3>, 3> M_expect = {
+            {{8, 1, 8}, {1, 1, 7}, {7, 1, 1}}};
+        std::array<std::array<FieldT, 3>, 3> M =
+            anemoi_permutation_mds_3x3<ppT>(g);
+        ASSERT_EQ(M, M_expect);
+    }
+    {
+        std::array<std::array<FieldT, 4>, 4> M_expect = {
+            {{1, 8, 7, 7}, {49, 56, 8, 15}, {49, 49, 1, 8}, {8, 15, 7, 8}}};
+        std::array<std::array<FieldT, 4>, 4> M =
+            anemoi_permutation_mds_4x4<ppT>(g);
+        ASSERT_EQ(M, M_expect);
+    }
+    libff::print_time("anemoi_permutation_mds tests successful");
+}
+
 template<typename ppT> void test_for_curve()
 {
     // Execute all tests for the given curve.
@@ -236,6 +265,7 @@ template<typename ppT> void test_for_curve()
     test_flystel_E_power_five_gadget<ppT>();
     test_flystel_E_root_five_gadget<ppT, parameters>();
     test_flystel_prime_field_gadget<ppT, parameters>();
+    test_anemoi_permutation_mds<ppT, parameters>();
 }
 
 TEST(TestAnemoiGadget, BLS12_381) { test_for_curve<libff::bls12_381_pp>(); }