review comments

argumentcomputer · Jun 6, 2024 · 6d58de2 · 6d58de2
1 parent 2329d91
commit 6d58de2
Show file tree

Hide file tree

Showing 4 changed files with 2,416 additions and 2,351 deletions.
diff --git a/scripts/poseidon2_rust_params.sage b/scripts/poseidon2_rust_params.sage
@@ -603,7 +603,7 @@ def to_field(value, p):
     value = hex(int(value))[2:]
     value = "0x" + value.zfill(l - 2)
     # print("from_hex(\"{}\"),".format(value))
-    print(f"f({value}),")
+    print(f"BabyBear::from_canonical_u32({value}),")
 
 
 def generate_constants_file():
@@ -615,15 +615,10 @@ def generate_constants_file():
     """)
     print()
     print("use lazy_static::*;")
+    print("use hybrid_array::{Array, typenum::*};")
     print("use p3_baby_bear::BabyBear;")
     print("use p3_field::AbstractField;")
     print()
-    print("""
-fn f(u: u32) -> BabyBear {
-    BabyBear::from_canonical_u32(u)
-}
-          """
-          )
     for t in range(4, 44, 4):
         p = 2013265921 # BabyBear
         n = len(p.bits())
@@ -653,10 +648,10 @@ fn f(u: u32) -> BabyBear {
         MATRIX_PARTIAL_DIAGONAL_M_1 = [matrix_partial_m_1(MATRIX_PARTIAL, NUM_CELLS, F)[i,i] for i in range(0, NUM_CELLS)]
 
         # Efficient partial matrix (diagonal - 1)
-        print("pub static ref MATRIX_DIAG_{}_BABYBEAR: [BabyBear; {}]= [".format(t, t))
+        print("pub static ref MATRIX_DIAG_{}_BABYBEAR: Array<BabyBear, U{}> = Array::clone_from_slice([".format(t, t))
         for val in MATRIX_PARTIAL_DIAGONAL_M_1:
-            to_field(val, p)
-        print("];")
+            to_field(val - 1, p)
+        print("].as_ref());")
         print()
 
 

diff --git a/src/poseidon/config.rs b/src/poseidon/config.rs
@@ -5,7 +5,8 @@ use std::slice;
 
 use hybrid_array::{typenum::*, Array, ArraySize};
 use p3_baby_bear::BabyBear;
-use p3_field::{AbstractField, Field};
+use p3_field::Field;
+use p3_poseidon2::matmul_internal;
 use p3_poseidon2::{DiffusionPermutation, Poseidon2, Poseidon2ExternalMatrixGeneral};
 use p3_symmetric::Permutation;
 
@@ -51,8 +52,9 @@ macro_rules! impl_poseidon_config {
             const R_F: usize = $r_f;
             type R = $r_t;
 
+            #[inline]
             fn matrix_diag() -> &'static Array<Self::F, Self::WIDTH> {
-                Array::from_slice(&*$diag)
+                $diag.as_ref()
             }
 
             fn round_constants_iter() -> impl IntoIterator<Item = &'static [Self::F]> {
@@ -67,10 +69,7 @@ macro_rules! impl_poseidon_config {
 
         impl Permutation<[$field; $width]> for InternalDiffusion {
             fn permute_mut(&self, input: &mut [$field; $width]) {
-                let sum: $field = input.iter().copied().sum();
-                for i in 0..$width {
-                    input[i] = sum + ($diag[i] + $field::neg_one()) * input[i];
-                }
+                matmul_internal(input, *$diag.as_ref());
             }
         }