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[j][i];
 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 = inp.clone();
105        for _ in inp.len()..t {
106            state.push(Fr::zero());
107        }
108        // state[state.len() - 1] = Fr::zero();
109
110        for i in 0..(n_rounds_f + n_rounds_p) {
111            self.ark(&mut state, &self.constants.c[t - 2], i * t);
112            self.sbox(n_rounds_f, n_rounds_p, &mut state, i);
113            if i < n_rounds_f + n_rounds_p - 1 {
114                state = self.mix(&state, &self.constants.m[t - 2]);
115            }
116        }
117
118        Ok(state[0])
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn test_ff() {
128        let a = Fr::from_repr(FrRepr::from(2)).unwrap();
129        assert_eq!(
130            "0000000000000000000000000000000000000000000000000000000000000002",
131            to_hex(&a)
132        );
133
134        let b: Fr = Fr::from_str(
135            "21888242871839275222246405745257275088548364400416034343698204186575808495619",
136        )
137        .unwrap();
138        assert_eq!(
139            "0000000000000000000000000000000000000000000000000000000000000002",
140            to_hex(&b)
141        );
142        assert_eq!(&a, &b);
143    }
144
145    #[test]
146    fn test_load_constants() {
147        let cons = load_constants();
148        assert_eq!(
149            cons.c[0][0].to_string(),
150            "Fr(0x09c46e9ec68e9bd4fe1faaba294cba38a71aa177534cdd1b6c7dc0dbd0abd7a7)"
151        );
152        assert_eq!(
153            cons.c[cons.c.len() - 1][0].to_string(),
154            "Fr(0x2088ce9534577bf38be7bc457f2756d558d66e0c07b9cc001a580bd42cda0e77)"
155        );
156        assert_eq!(
157            cons.m[0][0][0].to_string(),
158            "Fr(0x066f6f85d6f68a85ec10345351a23a3aaf07f38af8c952a7bceca70bd2af7ad5)"
159        );
160        assert_eq!(
161            cons.m[cons.m.len() - 1][0][0].to_string(),
162            "Fr(0x0190f922d97c8a7dcf0a142a3be27749d1c64bc22f1c556aaa24925d158cac56)"
163        );
164    }
165
166    #[test]
167    fn test_hash() {
168        let b0: Fr = Fr::from_str("0").unwrap();
169        let b1: Fr = Fr::from_str("1").unwrap();
170        let b2: Fr = Fr::from_str("2").unwrap();
171        let b3: Fr = Fr::from_str("3").unwrap();
172        let b4: Fr = Fr::from_str("4").unwrap();
173        let b5: Fr = Fr::from_str("5").unwrap();
174        let b6: Fr = Fr::from_str("6").unwrap();
175
176        let mut big_arr: Vec<Fr> = Vec::new();
177        big_arr.push(b1.clone());
178        let poseidon = Poseidon::new();
179        let h = poseidon.hash(big_arr.clone()).unwrap();
180        assert_eq!(
181            h.to_string(),
182            "Fr(0x186a5454a7c47c73dfc74ac32ea40a57d27eeb4e2bfc6551dd7b66686d3fd1ab)" // "11043376183861534927536506085090418075369306574649619885724436265926427398571"
183        );
184
185        let mut big_arr: Vec<Fr> = Vec::new();
186        big_arr.push(b1.clone());
187        big_arr.push(b2.clone());
188        let poseidon = Poseidon::new();
189        let h = poseidon.hash(big_arr.clone()).unwrap();
190        assert_eq!(
191            h.to_string(),
192            "Fr(0x25d86fb7c42fd70a7e800e871f22f2f03a282abb18f86c347a1078a92f713f60)" // "17117985411748610629288516079940078114952304104811071254131751175361957805920"
193        );
194
195        let mut big_arr: Vec<Fr> = Vec::new();
196        big_arr.push(b1.clone());
197        big_arr.push(b2.clone());
198        big_arr.push(b0.clone());
199        big_arr.push(b0.clone());
200        big_arr.push(b0.clone());
201        let poseidon = Poseidon::new();
202        let h = poseidon.hash(big_arr.clone()).unwrap();
203        assert_eq!(
204            h.to_string(),
205            "Fr(0x08ca0a9154fccd6426092b2404e1ceeb80a7849734f1d3fe7952c2075e489566)" // "3975478831357328722254985704342968745327876719981393787143845259590563829094"
206        );
207
208        let mut big_arr: Vec<Fr> = Vec::new();
209        big_arr.push(b1.clone());
210        big_arr.push(b2.clone());
211        big_arr.push(b0.clone());
212        big_arr.push(b0.clone());
213        big_arr.push(b0.clone());
214        big_arr.push(b0.clone());
215        let poseidon = Poseidon::new();
216        let h = poseidon.hash(big_arr.clone()).unwrap();
217        assert_eq!(
218            h.to_string(),
219            "Fr(0x2bb6c270db4ca49d129e315cdad9e0e678c1692c420dbf4667fdabc0f158e4ae)" // "19772360636270345724087386688434825760738403416279047262510528378903625000110"
220        );
221
222        let mut big_arr: Vec<Fr> = Vec::new();
223        big_arr.push(b3.clone());
224        big_arr.push(b4.clone());
225        big_arr.push(b0.clone());
226        big_arr.push(b0.clone());
227        big_arr.push(b0.clone());
228        let poseidon = Poseidon::new();
229        let h = poseidon.hash(big_arr.clone()).unwrap();
230        assert_eq!(
231            h.to_string(),
232            "Fr(0x07087ef123b0fc18a7487a9b3112aec23601e3d2b7ea27a85b35c7ecb595e6f6)" // "3181200837746671699652342497997860344148947482942465819251904554707352676086"
233        );
234
235        let mut big_arr: Vec<Fr> = Vec::new();
236        big_arr.push(b3.clone());
237        big_arr.push(b4.clone());
238        big_arr.push(b0.clone());
239        big_arr.push(b0.clone());
240        big_arr.push(b0.clone());
241        big_arr.push(b0.clone());
242        let h = poseidon.hash(big_arr.clone()).unwrap();
243        assert_eq!(
244            h.to_string(),
245            "Fr(0x128a815839bb66db834533b9c837e5a09df55e90aa9aba7ad46782234e083c20)" // "8386348873272147968934270337233829407378789978142456170950021426339096575008"
246        );
247
248        let mut big_arr: Vec<Fr> = Vec::new();
249        big_arr.push(b1.clone());
250        big_arr.push(b2.clone());
251        big_arr.push(b3.clone());
252        big_arr.push(b4.clone());
253        big_arr.push(b5.clone());
254        big_arr.push(b6.clone());
255        let h = poseidon.hash(big_arr.clone()).unwrap();
256        assert_eq!(
257            h.to_string(),
258            "Fr(0x0b807dafd5ecc62acdf7ae48e3a1dfb14ccc1ce398f865ac85ff0b4afd90ea6c)" // "5202465217520500374834597824465244016759843635092906214933648999760272616044"
259        );
260    }
261}