1extern crate num;
  2extern crate num_bigint;
  3extern crate num_traits;
  4
  5use blake2::digest::{Input, VariableOutput};
  6use blake2::VarBlake2b;
  7
  8use num_bigint::{BigInt, Sign};
  9use num_traits::{One, Zero};
 10
 11const SEED: &str = "poseidon";
 12const NROUNDSF: usize = 8;
 13const NROUNDSP: usize = 57;
 14const T: usize = 6;
 15
 16pub struct Constants {
 17    r: BigInt,
 18    c: Vec<BigInt>,
 19    m: Vec<Vec<BigInt>>,
 20}
 21
 22pub fn generate_constants() -> Constants {
 23    let r: BigInt = BigInt::parse_bytes(
 24        b"21888242871839275222246405745257275088548364400416034343698204186575808495617",
 25        10,
 26    )
 27    .unwrap();
 28    let c = get_pseudo_random(&r, format!("{}{}", SEED, "_constants"), NROUNDSF + NROUNDSP);
 29    let m = get_mds(&r);
 30    Constants { r: r, c: c, m: m }
 31}
 32
 33pub fn get_pseudo_random(r: &BigInt, seed: String, n: usize) -> Vec<BigInt> {
 34    let mut hasher = VarBlake2b::new(32).unwrap();
 35    hasher.input(seed.as_bytes());
 36    let mut h = hasher.vec_result();
 37
 38    let mut res: Vec<BigInt> = Vec::new();
 39    while res.len() < n {
 40        let new_n: BigInt = modulus(&BigInt::from_bytes_le(Sign::Plus, &h), &r);
 41        res.push(new_n);
 42        let mut hasher = VarBlake2b::new(32).unwrap();
 43        hasher.input(h);
 44        h = hasher.vec_result();
 45    }
 46    res
 47}
 48
 49pub fn nonce_to_string(n: usize) -> String {
 50    let mut r = format!("{}", n);
 51    while r.len() < 4 {
 52        r = format!("0{}", r);
 53    }
 54    r
 55}
 56pub fn get_mds(r: &BigInt) -> Vec<Vec<BigInt>> {
 57    let mut nonce = 0;
 58    let mut cauchy_matrix = get_pseudo_random(
 59        r,
 60        format!("{}_matrix_{}", SEED, nonce_to_string(nonce)),
 61        T * 2,
 62    );
 63    while !check_all_different(&cauchy_matrix) {
 64        nonce = nonce + 1;
 65        cauchy_matrix = get_pseudo_random(
 66            r,
 67            format!("{}_matrix_{}", SEED, nonce_to_string(nonce)),
 68            T * 2,
 69        );
 70    }
 71    let mut m: Vec<Vec<BigInt>> = Vec::new();
 72    for i in 0..T {
 73        let mut mi: Vec<BigInt> = Vec::new();
 74        for j in 0..T {
 75            mi.push(modinv(
 76                &modulus(&(&cauchy_matrix[i] - &cauchy_matrix[T + j]), &r),
 77                &r,
 78            ));
 79        }
 80        m.push(mi);
 81    }
 82    m
 83}
 84
 85pub fn check_all_different(v: &Vec<BigInt>) -> bool {
 86    let zero: BigInt = Zero::zero();
 87    for i in 0..v.len() {
 88        if v[i] == zero {
 89            return false;
 90        }
 91        for j in i + 1..v.len() {
 92            if v[i] == v[j] {
 93                return false;
 94            }
 95        }
 96    }
 97    true
 98}
 99
100pub fn modulus(a: &BigInt, m: &BigInt) -> BigInt {
101    ((a % m) + m) % m
102}
103
104pub fn modinv(a: &BigInt, q: &BigInt) -> BigInt {
105    let mut mn = (q.clone(), a.clone());
106    let mut xy: (BigInt, BigInt) = (Zero::zero(), One::one());
107
108    let big_zero: BigInt = Zero::zero();
109    while mn.1 != big_zero {
110        xy = (xy.1.clone(), xy.0 - (mn.0.clone() / mn.1.clone()) * xy.1);
111        mn = (mn.1.clone(), modulus(&mn.0, &mn.1));
112    }
113
114    while xy.0 < Zero::zero() {
115        xy.0 = modulus(&xy.0, q);
116    }
117    xy.0
118}
119
120pub fn check_bigint_in_field(a: &BigInt, q: &BigInt) -> bool {
121    if a >= q {
122        return false;
123    }
124    true
125}
126
127pub fn check_bigint_array_in_field(arr: &Vec<BigInt>, q: &BigInt) -> bool {
128    for a in arr {
129        if !check_bigint_in_field(a, &q) {
130            return false;
131        }
132    }
133    true
134}
135
136pub struct Poseidon {
137    constants: Constants,
138}
139impl Poseidon {
140    pub fn new() -> Poseidon {
141        Poseidon {
142            constants: generate_constants(),
143        }
144    }
145    pub fn ark(&self, state: &Vec<BigInt>, c: &BigInt) -> Vec<BigInt> {
146        let mut new_state: Vec<BigInt> = state.clone();
147        for i in 0..state.len() {
148            new_state[i] = modulus(&(&state[i] + c), &self.constants.r);
149        }
150
151        new_state
152    }
153
154    pub fn cubic(&self, a: &BigInt) -> BigInt {
155        modulus(&(a * a * a * a * a), &self.constants.r)
156    }
157
158    pub fn sbox(&self, state: &Vec<BigInt>, i: usize) -> Vec<BigInt> {
159        let mut new_state: Vec<BigInt> = state.clone();
160        if i < NROUNDSF / 2 || i >= NROUNDSF / 2 + NROUNDSP {
161            for j in 0..T {
162                new_state[j] = self.cubic(&state[j]);
163            }
164        } else {
165            new_state[0] = self.cubic(&state[0]);
166        }
167        new_state
168    }
169
170    pub fn mix(&self, state: &Vec<BigInt>, m: &Vec<Vec<BigInt>>) -> Vec<BigInt> {
171        let mut new_state: Vec<BigInt> = Vec::new();
172        for i in 0..state.len() {
173            new_state.push(Zero::zero());
174            for j in 0..state.len() {
175                new_state[i] = modulus(
176                    &(&new_state[i] + modulus(&(&m[i][j] * &state[j]), &self.constants.r)),
177                    &self.constants.r,
178                )
179            }
180        }
181        new_state
182    }
183
184    pub fn poseidon_hash(&self, inp: Vec<BigInt>) -> Result<BigInt, String> {
185        if inp.len() == 0 || inp.len() > T {
186            return Err("Wrong inputs length".to_string());
187        }
188        // check if arr elements are inside the finite field over R
189        if !check_bigint_array_in_field(&inp, &self.constants.r) {
190            return Err("elements not inside the finite field over R".to_string());
191        }
192
193        let mut state = inp.clone();
194        for _ in inp.len()..T {
195            state.push(Zero::zero());
196        }
197
198        for i in 0..(NROUNDSF + NROUNDSP) {
199            state = self.ark(&state, &self.constants.c[i]);
200            state = self.sbox(&state, i);
201            state = self.mix(&state, &self.constants.m);
202        }
203
204        Ok(state[0].clone())
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use rustc_hex::ToHex;
212
213    #[test]
214    fn test_blake2_version() {
215        let mut hasher = VarBlake2b::new(32).unwrap();
216        hasher.input(b"poseidon_constants");
217        let h = hasher.vec_result();
218        assert_eq!(
219            h.to_hex(),
220            "e57ba154fb2c47811dc1a2369b27e25a44915b4e4ece4eb8ec74850cb78e01b1"
221        );
222    }
223
224    #[test]
225    fn test_poseidon_hash() {
226        let b1: BigInt = BigInt::parse_bytes(b"1", 10).unwrap();
227        let b2: BigInt = BigInt::parse_bytes(b"2", 10).unwrap();
228        let mut big_arr: Vec<BigInt> = Vec::new();
229        big_arr.push(b1.clone());
230        big_arr.push(b2.clone());
231        let poseidon = Poseidon::new();
232        let h = poseidon.poseidon_hash(big_arr).unwrap();
233        assert_eq!(
234            h.to_string(),
235            "12242166908188651009877250812424843524687801523336557272219921456462821518061"
236        );
237
238        let b3: BigInt = BigInt::parse_bytes(b"3", 10).unwrap();
239        let b4: BigInt = BigInt::parse_bytes(b"4", 10).unwrap();
240        let mut big_arr34: Vec<BigInt> = Vec::new();
241        big_arr34.push(b3.clone());
242        big_arr34.push(b4.clone());
243        let h34 = poseidon.poseidon_hash(big_arr34).unwrap();
244        assert_eq!(
245            h34.to_string(),
246            "17185195740979599334254027721507328033796809509313949281114643312710535000993"
247        );
248    }
249}