Poseidon optimization starknet (#855)

* Poseidon * Add docs * Clearer docs * Remove unneeded pub crate * Clippy * Revert clippy * Allow op ref * Fix wasm target * Fix wasm target * Fmt
lambdaclass · Apr 4, 2024 · a70ec1d · a70ec1d
1 parent 699c12d
commit a70ec1d
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 18 deletions.
diff --git a/crypto/src/hash/poseidon/mod.rs b/crypto/src/hash/poseidon/mod.rs
@@ -18,7 +18,6 @@ pub trait Poseidon: PermutationParameters + self::private::Sealed {
     fn hades_permutation(state: &mut [FE<Self::F>]);
     fn full_round(state: &mut [FE<Self::F>], round_number: usize);
     fn partial_round(state: &mut [FE<Self::F>], round_number: usize);
-    fn mix(state: &mut [FE<Self::F>]);
     fn hash(x: &FE<Self::F>, y: &FE<Self::F>) -> FE<Self::F>;
     fn hash_single(x: &FE<Self::F>) -> FE<Self::F>;
     fn hash_many(inputs: &[FE<Self::F>]) -> FE<Self::F>;
@@ -56,21 +55,7 @@ impl<P: PermutationParameters> Poseidon for P {
 
         state[P::STATE_SIZE - 1] = state[P::STATE_SIZE - 1].pow(P::ALPHA);
 
-        Self::mix(state);
-    }
-
-    fn mix(state: &mut [FE<Self::F>]) {
-        let mut new_state: Vec<FE<Self::F>> = Vec::with_capacity(P::STATE_SIZE);
-        for i in 0..P::STATE_SIZE {
-            let mut new_e = FE::zero();
-            for (j, current_state) in state.iter().enumerate() {
-                let mut mij = P::MDS_MATRIX[i * P::N_MDS_MATRIX_COLS + j].clone();
-                mij *= current_state;
-                new_e += mij;
-            }
-            new_state.push(new_e);
-        }
-        state.clone_from_slice(&new_state[0..P::STATE_SIZE]);
+        P::mix(state);
     }
 
     fn hash(x: &FE<Self::F>, y: &FE<Self::F>) -> FE<Self::F> {

diff --git a/crypto/src/hash/poseidon/parameters.rs b/crypto/src/hash/poseidon/parameters.rs
@@ -1,3 +1,4 @@
+use alloc::vec::Vec;
 use lambdaworks_math::field::{element::FieldElement as FE, traits::IsPrimeField};
 
 /// Parameters for Poseidon
@@ -24,4 +25,21 @@ pub trait PermutationParameters {
     const ROUND_CONSTANTS: &'static [FE<Self::F>];
     const N_ROUND_CONSTANTS_ROWS: usize;
     const N_ROUND_CONSTANTS_COLS: usize;
+
+    /// This is the mix function that operates with the MDS matrix
+    /// Round Constants are sometimes picked to simplify this function,
+    /// so it can be redefined by each set of permutation parameters if a simplification can be made to make it faster. Notice in that case, MDS constants may not be used.
+    fn mix(state: &mut [FE<Self::F>]) {
+        let mut new_state: Vec<FE<Self::F>> = Vec::with_capacity(Self::STATE_SIZE);
+        for i in 0..Self::STATE_SIZE {
+            let mut new_e = FE::zero();
+            for (j, current_state) in state.iter().enumerate() {
+                let mut mij = Self::MDS_MATRIX[i * Self::N_MDS_MATRIX_COLS + j].clone();
+                mij *= current_state;
+                new_e += mij;
+            }
+            new_state.push(new_e);
+        }
+        state.clone_from_slice(&new_state[0..Self::STATE_SIZE]);
+    }
 }
diff --git a/crypto/src/hash/poseidon/starknet/parameters.rs b/crypto/src/hash/poseidon/starknet/parameters.rs
@@ -19,12 +19,22 @@ impl PermutationParameters for PoseidonCairoStark252 {
         &PoseidonCairoStark252::ROUND_CONSTANTS;
     const N_ROUND_CONSTANTS_ROWS: usize = 91;
     const N_ROUND_CONSTANTS_COLS: usize = 3;
+
+    /// Redefined mix function for optimization purposes
+    fn mix(state: &mut [FE<Self::F>]) {
+        let t = &state[0] + &state[1] + &state[2];
+        state[0] = &t + &state[0].double();
+        state[1] = &t - &state[1].double();
+        let minus_state_2 = -&state[2];
+        state[2] = &t + &minus_state_2 + &minus_state_2 + &minus_state_2;
+    }
 }
 
 #[derive(Clone, Default)]
 pub struct PoseidonCairoStark252;
 
 impl PoseidonCairoStark252 {
+    /// Since the mix function was redefined, these constants are not used, but are still valid.
     const MDS_MATRIX: [FE<Stark252PrimeField>; 3 * 3] = [
         FE::from_hex_unchecked("3"),
         FE::from_hex_unchecked("1"),
@@ -41,7 +51,7 @@ impl PoseidonCairoStark252 {
     // And the round 0 ones matches the one used
     // in Cairo Lang
     // https://github.com/starkware-libs/cairo-lang/blob/c98fc0b50529185b7018208cb3460191eeb53e0d/src/starkware/cairo/stark_verifier/air/layouts/starknet/autogenerated.cairo#L1574-L1596
-    const ROUND_CONSTANTS: [FE<Stark252PrimeField>; 3 * 91] = [
+    pub const ROUND_CONSTANTS: [FE<Stark252PrimeField>; 3 * 91] = [
         FE::from_hex_unchecked("6861759ea556a2339dd92f9562a30b9e58e2ad98109ae4780b7fd8eac77fe6f"),
         FE::from_hex_unchecked("3827681995d5af9ffc8397a3d00425a3da43f76abf28a64e4ab1a22f27508c4"),
         FE::from_hex_unchecked("3a3956d2fad44d0e7f760a2277dc7cb2cac75dc279b2d687a0dbe17704a8309"),

diff --git a/crypto/src/lib.rs b/crypto/src/lib.rs
@@ -1,5 +1,5 @@
+#![allow(clippy::op_ref)]
 #![cfg_attr(not(feature = "std"), no_std)]
-
 #[macro_use]
 extern crate alloc;