feat: replace stdlib poseidon implementation with optimized version

noir_stdlib/src/hash/poseidon.nr

@@ -3,66 +3,134 @@ use crate::field::modulus_num_bits;
 use crate::hash::Hasher;
 use crate::default::Default;
 
-struct PoseidonConfig<M,N> {
-    t: Field, // Width, i.e. state size
-    rf: u8, // Number of full rounds; should be even
-    rp: u8, // Number of partial rounds
-    alpha: Field, // S-box power; depends on the underlying field
-    ark: [Field; M], // Additive round keys
-    mds: [Field; N] // MDS Matrix in row-major order
+// A config struct defining the parameters of the Poseidon instance to use.
+//
+// A thorough writeup of this method (along with an unoptimized method) can be found at: https://spec.filecoin.io/algorithms/crypto/poseidon/
+struct PoseidonConfig<T, N, X> { 
+    // State width, should be equal to `T`
+    t: Field,
+    // Number of full rounds. should be even
+    rf: u8,
+    // Number of partial rounds
+    rp: u8,
+     // S-box power; depends on the underlying field
+    alpha: Field,
+    // The round constants for the
+    round_constants: [Field; N],
+    // The MDS matrix for the Poseidon instance
+    mds: [[Field; T]; T],
+    // An MDS matrix which also applies a basis transformation which allows
+    // sparse matrices to be used for the partial rounds.
+    //
+    // This should be applied instead of `mds` in the final full round.
+    presparse_mds: [[Field; T]; T],
+    // A set of sparse matrices used for MDS mixing for the partial rounds.
+    sparse_mds: [Field; X],
 }
 
-pub fn config<M, N>(
+pub fn config<T, N, X>(
     t: Field,
     rf: u8,
     rp: u8,
     alpha: Field,
-    ark: [Field; M],
-    mds: [Field; N]
-) -> PoseidonConfig<M,N> {
+    round_constants: [Field; N],
+    mds: [[Field; T]; T],
+    presparse_mds: [[Field; T]; T],
+    sparse_mds: [Field; X]
+) -> PoseidonConfig<T, N, X> {
     // Input checks
-    let mul = crate::wrapping_mul(t as u8, (rf + rp));
-    assert(mul == ark.len() as u8);
-    assert(t * t == mds.len() as Field);
+    assert_eq(rf & 1, 0);
+    assert_eq((t as u8) * rf + rp, N);
+    assert_eq(t, T);
     assert(alpha != 0);
 
-    PoseidonConfig { t, rf, rp, alpha, ark, mds }
+    PoseidonConfig { t, rf, rp, alpha, round_constants, mds, presparse_mds, sparse_mds }
 }
-// General Poseidon permutation on elements of type Field
-fn permute<M, N, O>(pos_conf: PoseidonConfig<M, N>, mut state: [Field; O]) -> [Field; O] {
-    let PoseidonConfig {t, rf, rp, alpha, ark, mds} = pos_conf;
 
-    assert(t == state.len() as Field);
+pub fn permute<T, N, X>(pos_conf: PoseidonConfig<T, N, X>, mut state: [Field; T]) -> [Field; T] {
+    let PoseidonConfig {t, rf, rp, alpha, round_constants, mds, presparse_mds, sparse_mds } = pos_conf;
 
-    let mut count = 0;
-    // for r in 0..rf + rp
-    for r in 0..(ark.len() / state.len()) {
-        for i in 0..state.len() {
-            state[i] = state[i] + ark[count + i];
-        } // Shift by round constants
+    for i in 0..state.len() {
+        state[i] += round_constants[i];
+    }
+
+    for _r in 0..rf / 2 - 1 {
+        state = sigma(state);
+        for i in 0..T {
+            state[i] += round_constants[T * (_r + 1) as u64 + i];
+        }
+        state = apply_matrix(mds, state);
+    }
+
+    state = sigma(state);
+    for i in 0..T {
+        state[i] += round_constants[T * (rf / 2) as u64 + i];
+    }
+    state = apply_matrix(presparse_mds, state);
+
+    for i in 0..T {
+        crate::as_witness(state[i]);
+    }
+
+    for _r in 0..rp {
         state[0] = state[0].pow_32(alpha);
-        // Check whether we are in a full round
-        if (r as u8 < rf / 2) | (r as u8 >= rf / 2 + rp) {
-            for i in 1..state.len() {
-                state[i] = state[i].pow_32(alpha);
+        state[0] += round_constants[(rf/2 + 1) as u64 * T + _r as u64];
+        crate::as_witness(state[0]);
+        {
+            let mut newState0 = 0;
+            for j in 0..T {
+                newState0 += sparse_mds[(T * 2 - 1) * _r as u64 + j as u64] * state[j];
+            }
+            for k in 1..T {
+                state[k] += state[0] * sparse_mds[(t * 2 - 1) as u64 * _r as u64 + T + k - 1];
+            }
+            state[0] = newState0;
+
+            if (_r & 1 == 0) {
+                for k in 1..T {
+                    crate::as_witness(state[k]);
+                }
             }
         }
+    }
 
-        state = apply_matrix(mds, state); // Apply MDS matrix
-        count = count + t as u64;
+    for _r in 0..rf / 2 - 1 {
+        state = sigma(state);
+        for i in 0..state.len() {
+            state[i] += round_constants[(rf/2+1) as u64 * T + rp as u64 + (_r as u64) * T + i];
+        }
+        state = apply_matrix(mds, state);
     }
 
+    state = sigma(state);
+    state = apply_matrix(mds, state);
+
     state
 }
-// Absorption. Fully absorbs input message.
-fn absorb<M, N, O, P>(
-    pos_conf: PoseidonConfig<M, N>,
-    mut state: [Field; O], // Initial state; usually [0; O]
-    rate: Field, // Rate
-    capacity: Field, // Capacity; usually 1
-    msg: [Field; P]
-) -> [Field; O] {
-    assert(pos_conf.t == rate + capacity);
+
+// Performs matrix multiplication on a vector
+fn apply_matrix<N>(matrix: [[Field; N]; N], vec: [Field; N]) -> [Field; N] {
+    let mut out = [0; N];
+
+    for i in 0..N {
+        for j in 0..N {
+            out[i] += vec[j] * matrix[j][i];
+        }
+    }
+
+    out
+}
+
+// Corresponding absorption.
+fn absorb<T, N, X, O>(
+    pos_conf: PoseidonConfig<T, N, X>,
+    // Initial state; usually [0; O]
+    mut state: [Field; T],
+    rate: Field,
+    capacity: Field,
+    msg: [Field; O] // Arbitrary length message
+) -> [Field; T] {
+    assert_eq(pos_conf.t, rate + capacity);
 
     let mut i = 0;
 
@@ -83,25 +151,24 @@ fn absorb<M, N, O, P>(
 
     state
 }
+
+fn sigma<O>(x: [Field; O]) -> [Field; O] {
+    let mut y = x;
+    for i in 0..O {
+        let t = y[i];
+        let tt = t * t;
+        let tttt = tt * tt;
+        y[i] *= tttt;
+    }
+    y
+}
+
 // Check security of sponge instantiation
 fn check_security(rate: Field, width: Field, security: Field) -> bool {
     let n = modulus_num_bits();
 
     ((n - 1) as Field * (width - rate) / 2) as u8 > security as u8
 }
-// A*x where A is an n x n matrix in row-major order and x an n-vector
-fn apply_matrix<N, M>(a: [Field; M], x: [Field; N]) -> [Field; N] {
-    let mut y = x;
-
-    for i in 0..x.len() {
-        y[i] = 0;
-        for j in 0..x.len() {
-            y[i] = y[i] + a[x.len()*i + j]* x[j];
-        }
-    }
-
-    y
-}
 
 struct PoseidonHasher{
     _state: [Field],
@@ -174,3 +241,28 @@ impl Default for PoseidonHasher{
         }
     }
 }
+
+mod poseidon_tests {
+    use crate::hash::poseidon;
+
+    #[test]
+    fn reference_impl_test_vectors() {
+        // hardcoded test vectors from https://extgit.iaik.tugraz.at/krypto/hadeshash/-/blob/master/code/test_vectors.txt
+        {
+            let mut state = [0, 1, 2];
+            let mut expected = [
+                0x115cc0f5e7d690413df64c6b9662e9cf2a3617f2743245519e19607a4417189a, 0x0fca49b798923ab0239de1c9e7a4a9a2210312b6a2f616d18b5a87f9b628ae29, 0x0e7ae82e40091e63cbd4f16a6d16310b3729d4b6e138fcf54110e2867045a30c
+            ];
+            assert_eq(expected, poseidon::bn254::perm::x5_3(state), "Failed to reproduce output for [0, 1, 2]");
+        }
+        {
+            let mut state = [0, 1, 2, 3, 4];
+            let mut expected = [
+                0x299c867db6c1fdd79dcefa40e4510b9837e60ebb1ce0663dbaa525df65250465, 0x1148aaef609aa338b27dafd89bb98862d8bb2b429aceac47d86206154ffe053d, 0x24febb87fed7462e23f6665ff9a0111f4044c38ee1672c1ac6b0637d34f24907, 0x0eb08f6d809668a981c186beaf6110060707059576406b248e5d9cf6e78b3d3e, 0x07748bc6877c9b82c8b98666ee9d0626ec7f5be4205f79ee8528ef1c4a376fc7
+            ];
+            assert_eq(
+                expected, poseidon::bn254::perm::x5_5(state), "Failed to reproduce output for [0, 1, 2, 3, 4]"
+            );
+        }
+    }
+}

noir_stdlib/src/hash/poseidon/bn254.nr

@@ -2,93 +2,14 @@
 mod perm;
 mod consts;
 
-use crate::hash::poseidon::PoseidonConfig;
-use crate::hash::poseidon::apply_matrix;
-// Optimised permutation for this particular field; uses hardcoded rf and rp values,
-// which should agree with those in pos_conf.
-#[field(bn254)]
-pub fn permute<M, N, O>(pos_conf: PoseidonConfig<M, N>, mut state: [Field; O]) -> [Field; O] {
-    let PoseidonConfig {t, rf: config_rf, rp: config_rp, alpha, ark, mds} = pos_conf;
-    let rf: u8 = 8;
-    let rp: u8 = [56, 57, 56, 60, 60, 63, 64, 63, 60, 66, 60, 65, 70, 60, 64, 68][state.len() - 2];
-
-    assert(t == state.len() as Field);
-    assert(rf == config_rf);
-    assert(rp == config_rp);
-
-    let mut count = 0;
-    // First half of full rounds
-    for _r in 0..rf / 2 {
-        for i in 0..state.len() {
-            state[i] = state[i] + ark[count + i];
-        } // Shift by round constants
-        for i in 0..state.len() {
-            state[i] = state[i].pow_32(alpha);
-        }
-
-        state = apply_matrix(mds, state); // Apply MDS matrix
-        count = count + t as u64;
-    }
-    // Partial rounds
-    for _r in 0..rp {
-        for i in 0..state.len() {
-            state[i] = state[i] + ark[count + i];
-        } // Shift by round constants
-        state[0] = state[0].pow_32(alpha);
-
-        state = apply_matrix(mds, state); // Apply MDS matrix
-        count = count + t as u64;
-    }
-    // Second half of full rounds
-    for _r in 0..rf / 2 {
-        for i in 0..state.len() {
-            state[i] = state[i] + ark[count + i];
-        } // Shift by round constants
-        for i in 0..state.len() {
-            state[i] = state[i].pow_32(alpha);
-        }
-
-        state = apply_matrix(mds, state); // Apply MDS matrix
-        count = count + t as u64;
-    }
-
-    state
-}
-// Corresponding absorption.
-#[field(bn254)]
-fn absorb<M, N, O, P>(
-    pos_conf: PoseidonConfig<M, N>,
-    mut state: [Field; O], // Initial state; usually [0; O]
-    rate: Field, // Rate
-    capacity: Field, // Capacity; usually 1
-    msg: [Field; P] // Arbitrary length message
-) -> [Field; O] {
-    assert(pos_conf.t == rate + capacity);
-
-    let mut i = 0;
-
-    for k in 0..msg.len() {
-        // Add current block to state
-        state[capacity + i] += msg[k];
-        i = i+1;
-        // Enough to absorb
-        if i == rate {
-            state = permute(pos_conf, state);
-            i = 0;
-        }
-    }
-    // If we have one more block to permute
-    if i != 0 {
-        state = permute(pos_conf, state);
-    }
-
-    state
-}
+use crate::hash::poseidon::{PoseidonConfig, absorb};
+
 // Variable-length Poseidon-128 sponge as suggested in second bullet point of §3 of https://eprint.iacr.org/2019/458.pdf
 #[field(bn254)]
 pub fn sponge<N>(msg: [Field; N]) -> Field {
     absorb(consts::x5_5_config(), [0; 5], 4, 1, msg)[1]
 }
+
 // Various instances of the Poseidon hash function
 // Consistent with Circom's implementation
 pub fn hash_1(input: [Field; 1]) -> Field {