Add GrainLFRS and PrimeField trait to Poseidon (#51)

* refactor(rln): generate poseidon constants

* refactor(rln): use traits in poseidon hash; refactor constants

* fix(rln): fix poseidon hardcoded constants test

* fix(rln): cargo fmt

rln/src/poseidon_hash.rs

@@ -1,69 +1,85 @@
 // This crate implements the Poseidon hash algorithm https://eprint.iacr.org/2019/458.pdf
 
-// The implementation is taken from https://github.com/arnaucube/poseidon-rs/blob/233027d6075a637c29ad84a8a44f5653b81f0410/src/lib.rs
-// and slightly adapted to work over arkworks field data type
+// Implementation partially taken from https://github.com/arnaucube/poseidon-rs/blob/233027d6075a637c29ad84a8a44f5653b81f0410/src/lib.rs
+// and adapted to work over arkworks field traits and custom data structures
 
-use crate::circuit::Fr;
-use crate::poseidon_constants::constants;
-use crate::utils::*;
-use ark_std::Zero;
-use once_cell::sync::Lazy;
+use crate::poseidon_constants::find_poseidon_ark_and_mds;
+use ark_ff::{FpParameters, PrimeField};
 
-#[derive(Debug)]
-pub struct Constants {
-    pub c: Vec<Vec<Fr>>,
-    pub m: Vec<Vec<Vec<Fr>>>,
+// These indexed constants hardcodes the round parameters tuple (t, RF, RN) from the paper for the Bn254 scalar field
+// SKIP_MATRICES is the index of the randomly generated secure MDS matrix. See security note in the poseidon_constants crate on this.
+// TODO: generate in-code such parameters
+pub const ROUND_PARAMS: [(usize, usize, usize, usize); 8] = [
+    (2, 8, 56, 0),
+    (3, 8, 57, 0),
+    (4, 8, 56, 0),
+    (5, 8, 60, 0),
+    (6, 8, 60, 0),
+    (7, 8, 63, 0),
+    (8, 8, 64, 0),
+    (9, 8, 63, 0),
+];
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct RoundParamenters<F: PrimeField> {
+    pub t: usize,
     pub n_rounds_f: usize,
-    pub n_rounds_p: Vec<usize>,
-}
-pub fn load_constants() -> Constants {
-    let (c_str, m_str) = constants();
-    let mut c: Vec<Vec<Fr>> = Vec::new();
-    for i in 0..c_str.len() {
-        let mut cci: Vec<Fr> = Vec::new();
-        for j in 0..c_str[i].len() {
-            let b: Fr = str_to_fr(c_str[i][j], 10);
-            cci.push(b);
-        }
-        c.push(cci);
-    }
-    let mut m: Vec<Vec<Vec<Fr>>> = Vec::new();
-    for i in 0..m_str.len() {
-        let mut mi: Vec<Vec<Fr>> = Vec::new();
-        for j in 0..m_str[i].len() {
-            let mut mij: Vec<Fr> = Vec::new();
-            for k in 0..m_str[i][j].len() {
-                let b: Fr = str_to_fr(m_str[i][j][k], 10);
-                mij.push(b);
-            }
-            mi.push(mij);
-        }
-        m.push(mi);
-    }
-    Constants {
-        c: c,
-        m: m,
-        n_rounds_f: 8,
-        n_rounds_p: vec![56, 57, 56, 60, 60, 63, 64, 63],
-    }
+    pub n_rounds_p: usize,
+    pub skip_matrices: usize,
+    pub c: Vec<F>,
+    pub m: Vec<Vec<F>>,
 }
 
-pub struct Poseidon {
-    constants: Constants,
+pub struct Poseidon<F: PrimeField> {
+    round_params: Vec<RoundParamenters<F>>,
 }
-impl Poseidon {
-    pub fn new() -> Poseidon {
+impl<F: PrimeField> Poseidon<F> {
+    // Loads round parameters and generates round constants
+    // poseidon_params is a vector containing tuples (t, RF, RP, skip_matrices)
+    // where: t is the rate (input lenght + 1), RF is the number of full rounds, RP is the number of partial rounds
+    // and skip_matrices is a (temporary) parameter used to generate secure MDS matrices (see comments in the description of find_poseidon_ark_and_mds)
+    // TODO: implement automatic generation of round parameters
+    pub fn from(poseidon_params: &[(usize, usize, usize, usize)]) -> Self {
+        let mut read_params = Vec::<RoundParamenters<F>>::new();
+
+        for i in 0..poseidon_params.len() {
+            let (t, n_rounds_f, n_rounds_p, skip_matrices) = poseidon_params[i];
+            let (ark, mds) = find_poseidon_ark_and_mds::<F>(
+                1, // is_field = 1
+                0, // is_sbox_inverse = 0
+                F::Params::MODULUS_BITS as u64,
+                t,
+                n_rounds_f as u64,
+                n_rounds_p as u64,
+                skip_matrices,
+            );
+            let rp = RoundParamenters {
+                t: t,
+                n_rounds_p: n_rounds_p,
+                n_rounds_f: n_rounds_f,
+                skip_matrices: skip_matrices,
+                c: ark,
+                m: mds,
+            };
+            read_params.push(rp);
+        }
+
         Poseidon {
-            constants: load_constants(),
+            round_params: read_params,
         }
     }
-    pub fn ark(&self, state: &mut [Fr], c: &[Fr], it: usize) {
+
+    pub fn get_parameters(&self) -> Vec<RoundParamenters<F>> {
+        self.round_params.clone()
+    }
+
+    pub fn ark(&self, state: &mut [F], c: &[F], it: usize) {
         for i in 0..state.len() {
             state[i] += c[it + i];
         }
     }
 
-    pub fn sbox(&self, n_rounds_f: usize, n_rounds_p: usize, state: &mut [Fr], i: usize) {
+    pub fn sbox(&self, n_rounds_f: usize, n_rounds_p: usize, state: &mut [F], i: usize) {
         if (i < n_rounds_f / 2) || (i >= n_rounds_f / 2 + n_rounds_p) {
             for j in 0..state.len() {
                 let aux = state[j];
@@ -79,10 +95,10 @@ impl Poseidon {
         }
     }
 
-    pub fn mix(&self, state: &[Fr], m: &[Vec<Fr>]) -> Vec<Fr> {
-        let mut new_state: Vec<Fr> = Vec::new();
+    pub fn mix(&self, state: &[F], m: &[Vec<F>]) -> Vec<F> {
+        let mut new_state: Vec<F> = Vec::new();
         for i in 0..state.len() {
-            new_state.push(Fr::zero());
+            new_state.push(F::zero());
             for j in 0..state.len() {
                 let mut mij = m[i][j];
                 mij *= state[j];
@@ -92,35 +108,58 @@ impl Poseidon {
         new_state.clone()
     }
 
-    pub fn hash(&self, inp: Vec<Fr>) -> Result<Fr, String> {
+    pub fn hash(&self, inp: Vec<F>) -> Result<F, String> {
+        // Note that the rate t becomes input lenght + 1, hence for lenght N we pick parameters with T = N + 1
         let t = inp.len() + 1;
-        if inp.is_empty() || (inp.len() >= self.constants.n_rounds_p.len() - 1) {
-            return Err("Wrong inputs length".to_string());
-        }
-        let n_rounds_f = self.constants.n_rounds_f;
-        let n_rounds_p = self.constants.n_rounds_p[t - 2];
 
-        let mut state = vec![Fr::zero(); t];
+        // We seek the index (Poseidon's round_params is an ordered vector) for the parameters corresponding to t
+        let param_index = self.round_params.iter().position(|el| el.t == t);
+
+        if inp.is_empty() || param_index.is_none() {
+            return Err("No parameters found for inputs length".to_string());
+        }
+
+        let param_index = param_index.unwrap();
+
+        let mut state = vec![F::zero(); t];
         state[1..].clone_from_slice(&inp);
 
-        for i in 0..(n_rounds_f + n_rounds_p) {
-            self.ark(&mut state, &self.constants.c[t - 2], i * t);
-            self.sbox(n_rounds_f, n_rounds_p, &mut state, i);
-            state = self.mix(&state, &self.constants.m[t - 2]);
+        for i in 0..(self.round_params[param_index].n_rounds_f
+            + self.round_params[param_index].n_rounds_p)
+        {
+            self.ark(
+                &mut state,
+                &self.round_params[param_index].c,
+                (i as usize) * self.round_params[param_index].t,
+            );
+            self.sbox(
+                self.round_params[param_index].n_rounds_f,
+                self.round_params[param_index].n_rounds_p,
+                &mut state,
+                i,
+            );
+            state = self.mix(&state, &self.round_params[param_index].m);
         }
 
         Ok(state[0])
     }
 }
 
-impl Default for Poseidon {
+impl<F> Default for Poseidon<F>
+where
+    F: PrimeField,
+{
+    // Default instantiation has no round constants set. Will return an error when hashing is attempted.
     fn default() -> Self {
-        Self::new()
+        Self::from(&[])
     }
 }
 
+use crate::circuit::Fr;
+use once_cell::sync::Lazy;
 // Poseidon Hash wrapper over above implementation. Adapted from semaphore-rs poseidon hash wrapper.
-static POSEIDON: Lazy<Poseidon> = Lazy::new(Poseidon::new);
+
+static POSEIDON: Lazy<Poseidon<Fr>> = Lazy::new(|| Poseidon::<Fr>::from(&ROUND_PARAMS));
 
 pub fn poseidon_hash(input: &[Fr]) -> Fr {
     POSEIDON

rln/src/utils.rs

@@ -1,25 +1,18 @@
 // This crate provides cross-module useful utilities (mainly type conversions) not necessarily specific to RLN
 
 use crate::circuit::Fr;
-use ark_ff::{BigInteger, FpParameters, PrimeField};
+use ark_ff::PrimeField;
 use num_bigint::{BigInt, BigUint};
 use num_traits::Num;
 use std::iter::Extend;
 
-pub fn modulus_bit_size() -> usize {
-    <Fr as PrimeField>::Params::MODULUS
-        .num_bits()
-        .try_into()
-        .unwrap()
-}
-
 pub fn to_bigint(el: &Fr) -> BigInt {
     let res: BigUint = (*el).try_into().unwrap();
     res.try_into().unwrap()
 }
 
 pub fn fr_byte_size() -> usize {
-    let mbs = modulus_bit_size();
+    let mbs = <Fr as PrimeField>::size_in_bits();
     (mbs + 64 - (mbs % 64)) / 8
 }