Optimize poseidon implementation

fawkes-crypto/Cargo.toml

@@ -26,6 +26,7 @@ getrandom = { version = "0.2", optional = true }
 bit-vec = "0.6.3"
 itertools = "0.10.0"
 brotli = "3.3.2"
+serde_json = "1.0.0"
 
 [dependencies.blake2_rfc]
 version = "0.0.1"
@@ -48,5 +49,4 @@ borsh_support = ["borsh", "ff-uint/borsh_support"]
 serde_support = ["serde", "ff-uint/serde_support"]
 rand_support = ["rand", "getrandom", "ff-uint/rand_support"]
 wasm = ["getrandom/js", "bellman/wasm"]
-multicore = ["bellman/multicore"]
-
+multicore = ["bellman/multicore"]

fawkes-crypto/src/native/poseidon/mod.rs

@@ -0,0 +1,69 @@
+use ff_uint::{PrimeField, Num};
+use itertools::Itertools;
+
+use crate::core::sizedvec::SizedVec;
+#[cfg(feature = "serde_support")]
+use crate::serde::{Deserialize, Serialize};
+
+pub mod optimized;
+pub mod unoptimized;
+
+pub type PoseidonParams<Fr> = self::optimized::params::PoseidonParams<Fr>;
+
+pub fn poseidon<Fr: PrimeField>(inputs: &[Num<Fr>], params: &PoseidonParams<Fr>) -> Num<Fr> {
+    self::optimized::poseidon::poseidon(inputs, params)
+}
+
+pub fn perm<Fr: PrimeField>(state: &mut [Num<Fr>], params: &PoseidonParams<Fr>) {
+    self::optimized::poseidon::perm(state, params)
+}
+
+pub fn poseidon_sponge<Fr: PrimeField>(inputs: &[Num<Fr>], params: &PoseidonParams<Fr>) -> Num<Fr> {
+    let mut state = vec![Num::ZERO; params.t];
+    let size = Num::from(inputs.len() as u64);
+    core::iter::once(&size).chain(inputs.iter()).chunks(params.t-1).into_iter().for_each(|c| {
+        state.iter_mut().zip(c.into_iter()).for_each(|(l, r)| *l+=*r);
+        perm(&mut state, params);
+    });
+    state[0]
+}
+
+
+#[derive(Clone, Debug)]
+#[cfg_attr(feature = "serde_support", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "serde_support", serde(bound(serialize = "", deserialize = "")))]
+pub struct MerkleProof<Fr: PrimeField, const L: usize> {
+    pub sibling: SizedVec<Num<Fr>, L>,
+    pub path: SizedVec<bool, L>,
+}
+
+pub fn poseidon_merkle_proof_root<Fr: PrimeField, const L: usize>(
+    leaf: Num<Fr>,
+    proof: &MerkleProof<Fr, L>,
+    params: &PoseidonParams<Fr>,
+) -> Num<Fr> {
+    let mut root = leaf.clone();
+    for (&p, &s) in proof.path.iter().zip(proof.sibling.iter()) {
+        let pair = if p { [s, root] } else { [root, s] };
+        root = poseidon(pair.as_ref(), params);
+    }
+    root
+}
+
+pub fn poseidon_merkle_tree_root<Fr: PrimeField>(
+    leaf: &[Num<Fr>],
+    params: &PoseidonParams<Fr>,
+) -> Num<Fr> {
+    let leaf_sz = leaf.len();
+    assert!(leaf_sz > 0, "should be at least one leaf in the tree");
+    let proof_sz = std::mem::size_of::<usize>() * 8 - (leaf_sz - 1).leading_zeros() as usize;
+    let total_leaf_sz = 1usize << proof_sz;
+    let mut state = leaf.to_vec();
+    state.extend_from_slice(&vec![Num::ZERO; total_leaf_sz - leaf_sz]);
+    for j in 0..proof_sz {
+        for i in 0..total_leaf_sz >> (j + 1) {
+            state[i] = poseidon(&[state[2 * i].clone(), state[2 * i + 1].clone()], params);
+        }
+    }
+    state[0].clone()
+}

fawkes-crypto/src/native/poseidon/optimized/constants.rs

@@ -0,0 +1,38 @@
+use std::iter;
+
+use ff_uint::{Num, PrimeField};
+
+use super::matrix::{apply_matrix, invert, transpose};
+
+// Reference implementation: https://extgit.iaik.tugraz.at/krypto/hadeshash/-/blob/master/code/poseidonperm_x3_64_24_optimized.sage#L43
+pub fn calc_equivalent_constants<Fr: PrimeField>(
+    rc: &Vec<Vec<Num<Fr>>>,
+    m: &Vec<Vec<Num<Fr>>>,
+    f: usize,
+    p: usize,
+    t: usize,
+) -> Vec<Vec<Num<Fr>>> {
+    let num_rounds = f + p;
+    let half_f = f >> 1;
+    let mut constants = rc.clone();
+    let m_transpose_inv = &invert(&transpose(&m)).unwrap();
+
+    let mut i = num_rounds - 2 - half_f;
+    while i > half_f - 1 {
+        let inv_cip1 = apply_matrix(m_transpose_inv, &constants[i + 1]);
+        constants[i] = constants[i]
+            .iter()
+            .zip([Num::ZERO].iter().chain(inv_cip1.iter().skip(1)))
+            .map(|(a, b)| *a + *b)
+            .collect();
+        constants[i + 1] = inv_cip1
+            .into_iter()
+            .take(1)
+            .chain(iter::repeat(Num::ZERO))
+            .take(t)
+            .collect();
+        i -= 1;
+    }
+
+    constants
+}