Poseidon2 m31

plain_implementations/src/fields/m31.rs

@@ -0,0 +1,8 @@
+use ark_ff::fields::{Fp64, MontBackend, MontConfig};
+use std::convert::TryInto;
+
+#[derive(MontConfig)]
+#[modulus = "2147483647"] // M31 = 2^31 - 1
+#[generator = "7"]
+pub struct FqConfig;
+pub type FpM31 = Fp64<MontBackend<FqConfig, 1>>;

plain_implementations/src/fields/mod.rs

@@ -15,6 +15,7 @@ pub mod babybear;
 pub mod pallas;
 pub mod vesta;
 pub mod utils;
+pub mod m31;
 
 // sage:
 // p = 21888242871839275222246405745257275088548364400416034343698204186575808495617

plain_implementations/src/poseidon2/mod.rs

@@ -6,4 +6,5 @@ pub mod poseidon2_instance_babybear;
 pub mod poseidon2_instance_bls12;
 pub mod poseidon2_instance_bn256;
 pub mod poseidon2_instance_pallas;
-pub mod poseidon2_instance_vesta;
+pub mod poseidon2_instance_vesta;
+pub mod poseidon2_instance_m31;

plain_implementations/src/poseidon2/poseidon2.rs

@@ -40,7 +40,7 @@ impl<F: PrimeField> Poseidon2<F> {
             current_state[0] = self.sbox_p(&current_state[0]);
             self.matmul_internal(&mut current_state, &self.params.mat_internal_diag_m_1);
         }
-        
+
         for r in p_end..self.params.rounds {
             current_state = self.add_rc(&current_state, &self.params.round_constants[r]);
             current_state = self.sbox(&current_state);
@@ -611,4 +611,66 @@ mod poseidon2_tests_vesta {
         assert_eq!(perm[2], from_hex("0x262316c0ce5244838c75873299b59d763ae0849d2dd31bdc95caf7db1c2901bf"));
 
     }
-}
+}
+
+#[allow(unused_imports)]
+#[cfg(test)]
+mod poseidon2_tests_m31 {
+    use super::*;
+    use crate::fields::{m31::FpM31, utils::from_hex, utils::random_scalar};
+    use crate::poseidon2::poseidon2_instance_m31::POSEIDON2_M31_16_PARAMS;
+    use std::convert::TryFrom;
+
+    type Scalar = FpM31;
+
+    static TESTRUNS: usize = 5;
+
+    #[test]
+    fn consistent_perm() {
+        let poseidon2 = Poseidon2::new(&POSEIDON2_M31_16_PARAMS);
+        let t = poseidon2.params.t;
+        for _ in 0..TESTRUNS {
+            let input1: Vec<Scalar> = (0..t).map(|_| random_scalar()).collect();
+
+            let mut input2: Vec<Scalar>;
+            loop {
+                input2 = (0..t).map(|_| random_scalar()).collect();
+                if input1 != input2 {
+                    break;
+                }
+            }
+
+            let perm1 = poseidon2.permutation(&input1);
+            let perm2 = poseidon2.permutation(&input1);
+            let perm3 = poseidon2.permutation(&input2);
+            assert_eq!(perm1, perm2);
+            assert_ne!(perm1, perm3);
+        }
+    }
+
+    #[test]
+    fn kats() {
+        let poseidon2 = Poseidon2::new(&POSEIDON2_M31_16_PARAMS);
+        let mut input: Vec<Scalar> = vec![];
+        for i in 0..poseidon2.params.t {
+            input.push(Scalar::from(i as u64));
+        }
+        let perm = poseidon2.permutation(&input);
+        assert_eq!(perm[0], from_hex("0x505d9689"));
+        assert_eq!(perm[1], from_hex("0x3b64c904"));
+        assert_eq!(perm[2], from_hex("0x79e2fd81"));
+        assert_eq!(perm[3], from_hex("0x4ba8015f"));
+        assert_eq!(perm[4], from_hex("0x24b6d2f5"));
+        assert_eq!(perm[5], from_hex("0x23845add"));
+        assert_eq!(perm[6], from_hex("0x521f4314"));
+        assert_eq!(perm[7], from_hex("0x69dfb019"));
+        assert_eq!(perm[8], from_hex("0x2aaae419"));
+        assert_eq!(perm[9], from_hex("0x6cb4502c"));
+        assert_eq!(perm[10], from_hex("0x6f7fa65a"));
+        assert_eq!(perm[11], from_hex("0x75feff24"));
+        assert_eq!(perm[12], from_hex("0x128d6587"));
+        assert_eq!(perm[13], from_hex("0x515877e4"));
+        assert_eq!(perm[14], from_hex("0x037f4dd7"));
+        assert_eq!(perm[15], from_hex("0x134b427f"));
+    }
+}