Modify constant generation to match reference implementation

src/primitives/poseidon.rs

@@ -25,9 +25,9 @@ pub trait PoseidonSpec<F: FieldExt> {
     fn constants(&self) -> (Vec<Vec<F>>, Vec<Vec<F>>, Vec<Vec<F>>);
 }
 
-/// A little-endian Poseidon specification.
+/// A generic Poseidon specification.
 #[derive(Debug)]
-pub struct LsbPoseidon<F: FieldExt> {
+pub struct Generic<F: FieldExt> {
     sbox: SboxType,
     /// The arity of the Poseidon permutation.
     t: u16,
@@ -41,7 +41,7 @@ pub struct LsbPoseidon<F: FieldExt> {
     _field: PhantomData<F>,
 }
 
-impl<F: FieldExt> LsbPoseidon<F> {
+impl<F: FieldExt> Generic<F> {
     /// Creates a new Poseidon specification for a field, using the `x^\alpha` S-box.
     pub fn with_pow_sbox(
         arity: usize,
@@ -49,7 +49,7 @@ impl<F: FieldExt> LsbPoseidon<F> {
         partial_rounds: usize,
         secure_mds: usize,
     ) -> Self {
-        LsbPoseidon {
+        Generic {
             sbox: SboxType::Pow,
             t: arity as u16,
             r_f: full_rounds as u16,
@@ -60,7 +60,7 @@ impl<F: FieldExt> LsbPoseidon<F> {
     }
 }
 
-impl<F: FieldExt> PoseidonSpec<F> for LsbPoseidon<F> {
+impl<F: FieldExt> PoseidonSpec<F> for Generic<F> {
     fn arity(&self) -> usize {
         self.t as usize
     }

src/primitives/poseidon/grain.rs

@@ -114,9 +114,19 @@ impl<F: FieldExt> Grain<F> {
         loop {
             let mut bytes = F::Repr::default();
 
-            // Fill the repr with bits in little-endian order.
+            // Poseidon reference impl interprets the bits as a repr in MSB order, because
+            // it's easy to do that in Python. Meanwhile, our field elements all use LSB
+            // order. There's little motivation to diverge from the reference impl; these
+            // are all constants, so we aren't introducing big-endianness into the rest of
+            // the circuit (assuming unkeyed Poseidon, but we probably wouldn't want to
+            // implement Grain inside a circuit, so we'd use a different round constant
+            // derivation function there).
             let view = bytes.as_mut();
             for (i, bit) in self.take(F::NUM_BITS as usize).enumerate() {
+                // If we diverged from the reference impl and interpreted the bits in LSB
+                // order, we would remove this line.
+                let i = F::NUM_BITS as usize - 1 - i;
+
                 view[i / 8] |= if bit { 1 << (i % 8) } else { 0 };
             }
 
@@ -125,6 +135,34 @@ impl<F: FieldExt> Grain<F> {
             }
         }
     }
+
+    /// Returns the next field element from this Grain instantiation, without using
+    /// rejection sampling.
+    pub(super) fn next_field_element_without_rejection(&mut self) -> F {
+        let mut bytes = [0u8; 64];
+
+        // Poseidon reference impl interprets the bits as a repr in MSB order, because
+        // it's easy to do that in Python. Additionally, it does not use rejection
+        // sampling in cases where the constants don't specifically need to be uniformly
+        // random for security. We do not provide APIs that take a field-element-sized
+        // array and reduce it modulo the field order, because those are unsafe APIs to
+        // offer generally (accidentally using them can lead to divergence in consensus
+        // systems due to not rejecting canonical forms).
+        //
+        // Given that we don't want to diverge from the reference implementation, we hack
+        // around this restriction by serializing the bits into a 64-byte array and then
+        // calling F::from_bytes_wide. PLEASE DO NOT COPY THIS INTO YOUR OWN CODE!
+        let view = bytes.as_mut();
+        for (i, bit) in self.take(F::NUM_BITS as usize).enumerate() {
+            // If we diverged from the reference impl and interpreted the bits in LSB
+            // order, we would remove this line.
+            let i = F::NUM_BITS as usize - 1 - i;
+
+            view[i / 8] |= if bit { 1 << (i % 8) } else { 0 };
+        }
+
+        F::from_bytes_wide(&bytes)
+    }
 }
 
 impl<F: FieldExt> Iterator for Grain<F> {

src/primitives/poseidon/mds.rs

@@ -10,7 +10,9 @@ pub(super) fn generate_mds<F: FieldExt>(
     let (xs, ys, mds) = loop {
         // Generate two [F; arity] arrays of unique field elements.
         let (xs, ys) = loop {
-            let mut vals: Vec<_> = (0..2 * arity).map(|_| grain.next_field_element()).collect();
+            let mut vals: Vec<_> = (0..2 * arity)
+                .map(|_| grain.next_field_element_without_rejection())
+                .collect();
 
             // Check that we have unique field elements.
             let mut unique = vals.clone();
@@ -33,20 +35,27 @@ pub(super) fn generate_mds<F: FieldExt>(
         }
 
         // Generate a Cauchy matrix, with elements a_ij in the form:
-        //     a_ij = 1/(x_i - y_j); x_i - y_j != 0
-        //
-        // The Poseidon paper uses the alternate definition:
         //     a_ij = 1/(x_i + y_j); x_i + y_j != 0
         //
-        // These are clearly equivalent on `y <= -y`, but it is easier to work with the
+        // It would be much easier to use the alternate definition:
+        //     a_ij = 1/(x_i - y_j); x_i - y_j != 0
+        //
+        // These are clearly equivalent on `y <- -y`, but it is easier to work with the
         // negative formulation, because ensuring that xs ∪ ys is unique implies that
         // x_i - y_j != 0 by construction (whereas the positive case does not hold). It
         // also makes computation of the matrix inverse simpler below (the theorem used
         // was formulated for the negative definition).
+        //
+        // However, the Poseidon paper and reference impl use the positive formulation,
+        // and we want to rely on the reference impl for MDS security, so we use the same
+        // formulation.
         let mut mds = vec![vec![F::zero(); arity]; arity];
         for i in 0..arity {
             for j in 0..arity {
-                mds[i][j] = (xs[i] - ys[j]).invert().unwrap();
+                let sum = xs[i] + ys[j];
+                // We leverage the secure MDS selection counter to also check this.
+                assert!(!sum.is_zero());
+                mds[i][j] = sum.invert().unwrap();
             }
         }
 
@@ -54,11 +63,17 @@ pub(super) fn generate_mds<F: FieldExt>(
     };
 
     // Compute the inverse. All square Cauchy matrices have a non-zero determinant and
-    // thus are invertible. The inverse has elements b_ij given by:
+    // thus are invertible. The inverse for a Cauchy matrix of the form:
+    //
+    //     a_ij = 1/(x_i - y_j); x_i - y_j != 0
+    //
+    // has elements b_ij given by:
     //
     //     b_ij = (x_j - y_i) A_j(y_i) B_i(x_j)    (Schechter 1959, Theorem 1)
     //
     // where A_i(x) and B_i(x) are the Lagrange polynomials for xs and ys respectively.
+    //
+    // We adapt this to the positive Cauchy formulation by negating ys.
     let mut mds_inv = vec![vec![F::zero(); arity]; arity];
     let l = |xs: &[F], j, x: F| {
         let x_j = xs[j];
@@ -71,9 +86,10 @@ pub(super) fn generate_mds<F: FieldExt>(
             }
         })
     };
+    let neg_ys: Vec<_> = ys.iter().map(|y| -*y).collect();
     for i in 0..arity {
         for j in 0..arity {
-            mds_inv[i][j] = (xs[j] - ys[i]) * l(&xs, j, ys[i]) * l(&ys, i, xs[j]);
+            mds_inv[i][j] = (xs[j] - neg_ys[i]) * l(&xs, j, neg_ys[i]) * l(&neg_ys, i, xs[j]);
         }
     }
 

src/primitives/poseidon/test_vectors.rs

@@ -1,7 +1,7 @@
 use halo2::arithmetic::FieldExt;
 use pasta_curves::pallas;
 
-use super::{LsbPoseidon, PoseidonSpec};
+use super::{Generic, PoseidonSpec};
 
 // $ sage generate_parameters_grain.sage 1 0 255 3 8 120 0x40000000000000000000000000000000224698fc094cf91b992d30ed00000001
 // Number of round constants: 384
@@ -424,7 +424,7 @@ const MDS: [[&str; 3]; 3] = [
 
 #[test]
 fn test_vectors() {
-    let poseidon = LsbPoseidon::<pallas::Base>::with_pow_sbox(3, 8, 120, 0);
+    let poseidon = Generic::<pallas::Base>::with_pow_sbox(3, 8, 120, 0);
     let (round_constants, mds, _) = poseidon.constants();
 
     for (actual, expected) in round_constants