optimize poseidon bn254

field/src/bn254.rs

@@ -4,15 +4,15 @@ use core::fmt::{Debug, Display, Formatter};
 use core::hash::{Hash, Hasher};
 use core::iter::{Product, Sum};
 use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
+use core::str::FromStr;
+
 use itertools::Itertools;
-use serde::{Serialize, Deserialize};
 use num::bigint::BigUint;
 use num::{Integer, One};
 use rand::RngCore;
-use core::str::FromStr;
+use serde::{Deserialize, Serialize};
 
-use crate::types::{Field, Sample};
-use crate::types::PrimeField as native_pf;
+use crate::types::{Field, PrimeField as native_pf, Sample};
 
 #[derive(Copy, Clone, Serialize, Deserialize)]
 pub struct Bn254Field(pub [u64; 4]);
@@ -48,7 +48,6 @@ impl PartialEq for Bn254Field {
     }
 }
 
-
 impl Eq for Bn254Field {}
 
 impl Hash for Bn254Field {
@@ -72,8 +71,8 @@ impl Debug for Bn254Field {
 impl Sample for Bn254Field {
     #[inline]
     fn sample<R>(rng: &mut R) -> Self
-        where
-            R: rand::RngCore + ?Sized,
+    where
+        R: rand::RngCore + ?Sized,
     {
         use num::bigint::RandBigInt;
         Self::from_noncanonical_biguint(rng.gen_biguint_below(&Self::order()))
@@ -91,7 +90,7 @@ impl Field for Bn254Field {
         4891460686036598784 as u64,
         2896914383306846353 as u64,
         13281191951274694749 as u64,
-        3486998266802970665 as u64
+        3486998266802970665 as u64,
     ]);
 
     const TWO_ADICITY: usize = 28;
@@ -100,17 +99,15 @@ impl Field for Bn254Field {
 
     const MULTIPLICATIVE_GROUP_GENERATOR: Self = Self([5, 0, 0, 0]);
 
-    const POWER_OF_TWO_GENERATOR: Self = Self([
-        268435456 as u64,
-        0,
-        0,
-        0
-    ]);
+    const POWER_OF_TWO_GENERATOR: Self = Self([268435456 as u64, 0, 0, 0]);
 
     const BITS: usize = 254;
 
     fn order() -> BigUint {
-        BigUint::from_str("21888242871839275222246405745257275088548364400416034343698204186575808495617").unwrap()
+        BigUint::from_str(
+            "21888242871839275222246405745257275088548364400416034343698204186575808495617",
+        )
+        .unwrap()
     }
 
     fn characteristic() -> BigUint {
@@ -232,7 +229,7 @@ impl Mul for Bn254Field {
     }
 }
 
-impl MulAssign for Bn254Field{
+impl MulAssign for Bn254Field {
     #[inline]
     fn mul_assign(&mut self, rhs: Self) {
         *self = *self * rhs;
@@ -269,4 +266,4 @@ mod tests {
     pub fn test() {
         test_field_arithmetic!(crate::bn254::Bn254Field);
     }
-}
+}

plonky2/src/hash/poseidon_bn254.rs

@@ -1,6 +1,7 @@
 use core::ops::{AddAssign, MulAssign};
 
 use plonky2_field::ops::Square;
+use unroll::unroll_for_loops;
 
 use crate::field::bn254::Bn254Field;
 use crate::field::types::Field;
@@ -14,19 +15,23 @@ pub const GOLDILOCKS_ELEMENTS: usize = 3;
 
 pub type PoseidonState = [Bn254Field; WIDTH];
 
-pub fn permution(state: &mut PoseidonState) {
+#[inline(always)]
+pub fn permutation(state: &mut PoseidonState) {
     ark(state, 0);
     full_rounds(state, true);
     partial_rounds(state);
     full_rounds(state, false);
 }
 
+#[inline(always)]
+#[unroll_for_loops]
 fn ark(state: &mut PoseidonState, it: usize) {
     for i in 0..WIDTH {
         state[i].add_assign(C_CONSTANTS[it + i]);
     }
 }
 
+#[inline(always)]
 fn exp5(mut x: Bn254Field) -> Bn254Field {
     let aux = x;
     x = x.square();
@@ -36,12 +41,16 @@ fn exp5(mut x: Bn254Field) -> Bn254Field {
     x
 }
 
+#[inline(always)]
+#[unroll_for_loops]
 fn exp5_state(state: &mut PoseidonState) {
     for state_element in state.iter_mut().take(WIDTH) {
         *state_element = exp5(*state_element);
     }
 }
 
+#[inline(always)]
+#[unroll_for_loops]
 fn full_rounds(state: &mut PoseidonState, first: bool) {
     for i in 0..FULL_ROUNDS / 2 - 1 {
         exp5_state(state);
@@ -65,6 +74,8 @@ fn full_rounds(state: &mut PoseidonState, first: bool) {
     }
 }
 
+#[inline(always)]
+#[unroll_for_loops]
 fn partial_rounds(state: &mut PoseidonState) {
     for i in 0..PARTIAL_ROUNDS {
         state[0] = exp5(state[0]);
@@ -90,6 +101,8 @@ fn partial_rounds(state: &mut PoseidonState) {
     }
 }
 
+#[inline(always)]
+#[unroll_for_loops]
 fn mix(state: &mut PoseidonState, constant_matrix: &[Vec<Bn254Field>]) {
     let mut result: PoseidonState = [Bn254Field::ZERO; WIDTH];
 
@@ -110,7 +123,7 @@ fn mix(state: &mut PoseidonState, constant_matrix: &[Vec<Bn254Field>]) {
 mod permutation_tests {
     use anyhow::Ok;
 
-    use super::{permution, WIDTH};
+    use super::{permutation, WIDTH};
     use crate::field::bn254::Bn254Field;
     use crate::field::types::Field;
 
@@ -177,7 +190,7 @@ mod permutation_tests {
         ];
 
         for (mut input, expected_output) in test_vectors.into_iter() {
-            permution(&mut input);
+            permutation(&mut input);
             for i in 0..WIDTH {
                 assert_eq!(input[i], expected_output[i]);
             }

plonky2/src/hash/poseidon_bn254_goldilocks_hash.rs

@@ -11,7 +11,7 @@ use crate::field::bn254::Bn254Field;
 use crate::field::types::Field;
 use crate::hash::hash_types::RichField;
 use crate::hash::poseidon::PoseidonPermutation;
-use crate::hash::poseidon_bn254::{permution, GOLDILOCKS_ELEMENTS, RATE};
+use crate::hash::poseidon_bn254::{permutation, GOLDILOCKS_ELEMENTS, RATE};
 use crate::plonk::config::{GenericHashOut, Hasher};
 
 pub const NUM_HASH_OUT_ELTS: usize = 1;
@@ -129,7 +129,7 @@ impl<F: RichField> Hasher<F> for PoseidonBn254Hash {
                 state[j + 1] =
                 Bn254Field::from_noncanonical_biguint(BigUint::from_bytes_le(&sized_bytes));
             }
-            permution(&mut state);
+            permutation(&mut state);
         }
 
         PoseidonBn254HashOut {
@@ -163,8 +163,7 @@ impl<F: RichField> Hasher<F> for PoseidonBn254Hash {
 
     fn two_to_one(left: Self::Hash, right: Self::Hash) -> Self::Hash {
         let mut state = [Bn254Field::ZERO, Bn254Field::ZERO, left.value, right.value];
-        permution(&mut state);
-
+        permutation(&mut state);
         PoseidonBn254HashOut {
             value: state[0],
             _phantom: PhantomData,