arnaucube
arnaucube/poseidon-rsPublic Self-hosted repository
poseidon-rs/src/lib.rs
src/lib.rsRaw
  1extern crate rand;
  2#[macro_use]
  3extern crate ff;
  4use ff::*;
  5
  6#[derive(PrimeField)]
  7#[PrimeFieldModulus = "21888242871839275222246405745257275088548364400416034343698204186575808495617"]
  8#[PrimeFieldGenerator = "7"]
  9pub struct Fr(FrRepr);
 10
 11mod constants;
 12
 13#[derive(Debug)]
 14pub struct Constants {
 15    pub c: Vec<Vec<Fr>>,
 16    pub m: Vec<Vec<Vec<Fr>>>,
 17    pub n_rounds_f: usize,
 18    pub n_rounds_p: Vec<usize>,
 19}
 20pub fn load_constants() -> Constants {
 21    let (c_str, m_str) = constants::constants();
 22    let mut c: Vec<Vec<Fr>> = Vec::new();
 23    for i in 0..c_str.len() {
 24        let mut cci: Vec<Fr> = Vec::new();
 25        for j in 0..c_str[i].len() {
 26            let b: Fr = Fr::from_str(c_str[i][j]).unwrap();
 27            cci.push(b);
 28        }
 29        c.push(cci);
 30    }
 31    let mut m: Vec<Vec<Vec<Fr>>> = Vec::new();
 32    for i in 0..m_str.len() {
 33        let mut mi: Vec<Vec<Fr>> = Vec::new();
 34        for j in 0..m_str[i].len() {
 35            let mut mij: Vec<Fr> = Vec::new();
 36            for k in 0..m_str[i][j].len() {
 37                let b: Fr = Fr::from_str(m_str[i][j][k]).unwrap();
 38                mij.push(b);
 39            }
 40            mi.push(mij);
 41        }
 42        m.push(mi);
 43    }
 44    Constants {
 45        c: c,
 46        m: m,
 47        n_rounds_f: 8,
 48        n_rounds_p: vec![56, 57, 56, 60, 60, 63, 64, 63],
 49    }
 50}
 51
 52pub struct Poseidon {
 53    constants: Constants,
 54}
 55impl Poseidon {
 56    pub fn new() -> Poseidon {
 57        Poseidon {
 58            constants: load_constants(),
 59        }
 60    }
 61    pub fn ark(&self, state: &mut Vec<Fr>, c: &Vec<Fr>, it: usize) {
 62        for i in 0..state.len() {
 63            state[i].add_assign(&c[it + i]);
 64        }
 65    }
 66
 67    pub fn sbox(&self, n_rounds_f: usize, n_rounds_p: usize, state: &mut Vec<Fr>, i: usize) {
 68        if i < n_rounds_f / 2 || i >= n_rounds_f / 2 + n_rounds_p {
 69            for j in 0..state.len() {
 70                let aux = state[j];
 71                state[j].square();
 72                state[j].square();
 73                state[j].mul_assign(&aux);
 74            }
 75        } else {
 76            let aux = state[0];
 77            state[0].square();
 78            state[0].square();
 79            state[0].mul_assign(&aux);
 80        }
 81    }
 82
 83    pub fn mix(&self, state: &Vec<Fr>, m: &Vec<Vec<Fr>>) -> Vec<Fr> {
 84        let mut new_state: Vec<Fr> = Vec::new();
 85        for i in 0..state.len() {
 86            new_state.push(Fr::zero());
 87            for j in 0..state.len() {
 88                let mut mij = m[i][j];
 89                mij.mul_assign(&state[j]);
 90                new_state[i].add_assign(&mij);
 91            }
 92        }
 93        new_state.clone()
 94    }
 95
 96    pub fn hash(&self, inp: Vec<Fr>) -> Result<Fr, String> {
 97        let t = inp.len() + 1;
 98        if inp.len() == 0 || inp.len() >= self.constants.n_rounds_p.len() - 1 {
 99            return Err("Wrong inputs length".to_string());
100        }
101        let n_rounds_f = self.constants.n_rounds_f.clone();
102        let n_rounds_p = self.constants.n_rounds_p[t - 2].clone();
103
104        let mut state = vec![Fr::zero(); t];
105        state[1..].clone_from_slice(&inp);
106
107        for i in 0..(n_rounds_f + n_rounds_p) {
108            self.ark(&mut state, &self.constants.c[t - 2], i * t);
109            self.sbox(n_rounds_f, n_rounds_p, &mut state, i);
110            state = self.mix(&state, &self.constants.m[t - 2]);
111        }
112
113        Ok(state[0])
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_ff() {
123        let a = Fr::from_repr(FrRepr::from(2)).unwrap();
124        assert_eq!(
125            "0000000000000000000000000000000000000000000000000000000000000002",
126            to_hex(&a)
127        );
128
129        let b: Fr = Fr::from_str(
130            "21888242871839275222246405745257275088548364400416034343698204186575808495619",
131        )
132        .unwrap();
133        assert_eq!(
134            "0000000000000000000000000000000000000000000000000000000000000002",
135            to_hex(&b)
136        );
137        assert_eq!(&a, &b);
138    }
139
140    #[test]
141    fn test_load_constants() {
142        let cons = load_constants();
143        assert_eq!(
144            cons.c[0][0].to_string(),
145            "Fr(0x09c46e9ec68e9bd4fe1faaba294cba38a71aa177534cdd1b6c7dc0dbd0abd7a7)"
146        );
147        assert_eq!(
148            cons.c[cons.c.len() - 1][0].to_string(),
149            "Fr(0x2088ce9534577bf38be7bc457f2756d558d66e0c07b9cc001a580bd42cda0e77)"
150        );
151        assert_eq!(
152            cons.m[0][0][0].to_string(),
153            "Fr(0x066f6f85d6f68a85ec10345351a23a3aaf07f38af8c952a7bceca70bd2af7ad5)"
154        );
155        assert_eq!(
156            cons.m[cons.m.len() - 1][0][0].to_string(),
157            "Fr(0x0190f922d97c8a7dcf0a142a3be27749d1c64bc22f1c556aaa24925d158cac56)"
158        );
159    }
160
161    #[test]
162    fn test_hash() {
163        let b0: Fr = Fr::from_str("0").unwrap();
164        let b1: Fr = Fr::from_str("1").unwrap();
165        let b2: Fr = Fr::from_str("2").unwrap();
166        let b3: Fr = Fr::from_str("3").unwrap();
167        let b4: Fr = Fr::from_str("4").unwrap();
168        let b5: Fr = Fr::from_str("5").unwrap();
169        let b6: Fr = Fr::from_str("6").unwrap();
170
171        let mut big_arr: Vec<Fr> = Vec::new();
172        big_arr.push(b1.clone());
173        let poseidon = Poseidon::new();
174        let h = poseidon.hash(big_arr.clone()).unwrap();
175        assert_eq!(
176            h.to_string(),
177            "Fr(0x29176100eaa962bdc1fe6c654d6a3c130e96a4d1168b33848b897dc502820133)" // "18586133768512220936620570745912940619677854269274689475585506675881198879027"
178        );
179
180        let mut big_arr: Vec<Fr> = Vec::new();
181        big_arr.push(b1.clone());
182        big_arr.push(b2.clone());
183        let poseidon = Poseidon::new();
184        let h = poseidon.hash(big_arr.clone()).unwrap();
185        assert_eq!(
186            h.to_string(),
187            "Fr(0x115cc0f5e7d690413df64c6b9662e9cf2a3617f2743245519e19607a4417189a)" // "7853200120776062878684798364095072458815029376092732009249414926327459813530"
188        );
189
190        let mut big_arr: Vec<Fr> = Vec::new();
191        big_arr.push(b1.clone());
192        big_arr.push(b2.clone());
193        big_arr.push(b0.clone());
194        big_arr.push(b0.clone());
195        big_arr.push(b0.clone());
196        let poseidon = Poseidon::new();
197        let h = poseidon.hash(big_arr.clone()).unwrap();
198        assert_eq!(
199            h.to_string(),
200            "Fr(0x024058dd1e168f34bac462b6fffe58fd69982807e9884c1c6148182319cee427)" // "1018317224307729531995786483840663576608797660851238720571059489595066344487"
201        );
202
203        let mut big_arr: Vec<Fr> = Vec::new();
204        big_arr.push(b1.clone());
205        big_arr.push(b2.clone());
206        big_arr.push(b0.clone());
207        big_arr.push(b0.clone());
208        big_arr.push(b0.clone());
209        big_arr.push(b0.clone());
210        let poseidon = Poseidon::new();
211        let h = poseidon.hash(big_arr.clone()).unwrap();
212        assert_eq!(
213            h.to_string(),
214            "Fr(0x21e82f465e00a15965e97a44fe3c30f3bf5279d8bf37d4e65765b6c2550f42a1)" // "15336558801450556532856248569924170992202208561737609669134139141992924267169"
215        );
216
217        let mut big_arr: Vec<Fr> = Vec::new();
218        big_arr.push(b3.clone());
219        big_arr.push(b4.clone());
220        big_arr.push(b0.clone());
221        big_arr.push(b0.clone());
222        big_arr.push(b0.clone());
223        let poseidon = Poseidon::new();
224        let h = poseidon.hash(big_arr.clone()).unwrap();
225        assert_eq!(
226            h.to_string(),
227            "Fr(0x0cd93f1bab9e8c9166ef00f2a1b0e1d66d6a4145e596abe0526247747cc71214)" // "5811595552068139067952687508729883632420015185677766880877743348592482390548"
228        );
229
230        let mut big_arr: Vec<Fr> = Vec::new();
231        big_arr.push(b3.clone());
232        big_arr.push(b4.clone());
233        big_arr.push(b0.clone());
234        big_arr.push(b0.clone());
235        big_arr.push(b0.clone());
236        big_arr.push(b0.clone());
237        let h = poseidon.hash(big_arr.clone()).unwrap();
238        assert_eq!(
239            h.to_string(),
240            "Fr(0x1b1caddfc5ea47e09bb445a7447eb9694b8d1b75a97fff58e884398c6b22825a)" // "12263118664590987767234828103155242843640892839966517009184493198782366909018"
241        );
242
243        let mut big_arr: Vec<Fr> = Vec::new();
244        big_arr.push(b1.clone());
245        big_arr.push(b2.clone());
246        big_arr.push(b3.clone());
247        big_arr.push(b4.clone());
248        big_arr.push(b5.clone());
249        big_arr.push(b6.clone());
250        let h = poseidon.hash(big_arr.clone()).unwrap();
251        assert_eq!(
252            h.to_string(),
253            "Fr(0x2d1a03850084442813c8ebf094dea47538490a68b05f2239134a4cca2f6302e1)" // "20400040500897583745843009878988256314335038853985262692600694741116813247201"
254        );
255    }
256}