Merge pull request #5 from mir-protocol/validate_cosets

Validate that the cosets involved in Plonk's permutation argument are disjoint

src/circuit_builder.rs

@@ -10,7 +10,7 @@ use crate::gates::gate::{GateInstance, GateRef};
 use crate::gates::noop::NoopGate;
 use crate::generator::{CopyGenerator, WitnessGenerator};
 use crate::hash::merkle_root_bit_rev_order;
-use crate::partition::get_subgroup_shift;
+use crate::field::cosets::get_unique_coset_shifts;
 use crate::polynomial::polynomial::PolynomialValues;
 use crate::target::Target;
 use crate::util::{log2_strict, transpose, transpose_poly_values};
@@ -188,13 +188,12 @@ impl<F: Field> CircuitBuilder<F> {
             .max()
             .expect("No gates?");
 
-        let k_is = (0..self.config.num_routed_wires)
-            .map(get_subgroup_shift)
-            .collect();
+        let degree_bits = log2_strict(degree);
+        let k_is = get_unique_coset_shifts(degree, self.config.num_routed_wires);
 
         let common = CommonCircuitData {
             config: self.config,
-            degree_bits: log2_strict(degree),
+            degree_bits,
             gates,
             num_gate_constraints,
             constants_root,

src/circuit_data.rs

@@ -1,10 +1,8 @@
-use crate::circuit_builder::CircuitBuilder;
 use crate::field::field::Field;
 use crate::gates::gate::GateRef;
 use crate::generator::WitnessGenerator;
-use crate::proof::{Hash, Proof, HashTarget};
+use crate::proof::{Hash, HashTarget, Proof};
 use crate::prover::prove;
-use crate::target::Target;
 use crate::verifier::verify;
 use crate::witness::PartialWitness;
 
@@ -112,7 +110,7 @@ pub(crate) struct CommonCircuitData<F: Field> {
     /// A commitment to each permutation polynomial.
     pub(crate) sigmas_root: Hash<F>,
 
-    /// {k_i}. See `get_subgroup_shift`.
+    /// The `{k_i}` valued used in `S_ID_i` in Plonk's permutation argument.
     pub(crate) k_is: Vec<F>,
 }
 

src/field/cosets.rs

@@ -0,0 +1,51 @@
+use crate::field::field::Field;
+
+/// Finds a set of shifts that result in unique cosets for the multiplicative subgroup of size
+/// `2^subgroup_bits`.
+pub(crate) fn get_unique_coset_shifts<F: Field>(
+    subgroup_size: usize,
+    num_shifts: usize,
+) -> Vec<F> {
+    // From Lagrange's theorem.
+    let num_cosets = (F::ORDER - 1) / (subgroup_size as u64);
+    assert!(num_shifts as u64 <= num_cosets,
+            "The subgroup does not have enough distinct cosets");
+
+    // Let g be a generator of the entire multiplicative group. Let n be the order of the subgroup.
+    // The subgroup can be written as <g^(|F*| / n)>. We can use g^0, ..., g^(num_shifts - 1) as our
+    // shifts, since g^i <g^(|F*| / n)> are distinct cosets provided i < |F*| / n, which we checked.
+    F::MULTIPLICATIVE_GROUP_GENERATOR.powers()
+        .take(num_shifts)
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+
+    use crate::field::cosets::get_unique_coset_shifts;
+    use crate::field::crandall_field::CrandallField;
+    use crate::field::field::Field;
+
+    #[test]
+    fn distinct_cosets() {
+        // TODO: Switch to a smaller test field so that collision rejection is likely to occur.
+
+        type F = CrandallField;
+        const SUBGROUP_BITS: usize = 5;
+        const NUM_SHIFTS: usize = 50;
+
+        let generator = F::primitive_root_of_unity(SUBGROUP_BITS);
+        let subgroup_size = 1 << SUBGROUP_BITS;
+
+        let shifts = get_unique_coset_shifts::<F>(SUBGROUP_BITS, NUM_SHIFTS);
+
+        let mut union = HashSet::new();
+        for shift in shifts {
+            let coset = F::cyclic_subgroup_coset_known_order(generator, shift, subgroup_size);
+            assert!(
+                coset.into_iter().all(|x| union.insert(x)),
+                "Duplicate element!");
+        }
+    }
+}

src/field/crandall_field.rs

@@ -5,6 +5,7 @@ use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssi
 use num::Integer;
 
 use crate::field::field::Field;
+use std::hash::{Hash, Hasher};
 
 /// EPSILON = 9 * 2**28 - 1
 const EPSILON: u64 = 2415919103;
@@ -17,7 +18,6 @@ const EPSILON: u64 = 2415919103;
 ///   = 2**64 - 9 * 2**28 + 1
 ///   = 2**28 * (2**36 - 9) + 1
 /// ```
-// TODO: [Partial]Eq should compare canonical representations.
 #[derive(Copy, Clone)]
 pub struct CrandallField(pub u64);
 
@@ -29,6 +29,12 @@ impl PartialEq for CrandallField {
 
 impl Eq for CrandallField {}
 
+impl Hash for CrandallField {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        state.write_u64(self.to_canonical_u64())
+    }
+}
+
 impl Display for CrandallField {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         Display::fmt(&self.0, f)
@@ -64,59 +70,69 @@ impl Field for CrandallField {
     }
 
     fn try_inverse(&self) -> Option<Self> {
-        if *self == Self::ZERO {
+        if self.is_zero() {
             return None;
         }
 
         // Based on Algorithm 16 of "Efficient Software-Implementation of Finite Fields with
         // Applications to Cryptography".
 
-        let mut u = self.0;
-        let mut v = Self::ORDER;
-        let mut b = 1;
-        let mut c = 0;
+        let p = Self::ORDER;
+        let mut u = self.to_canonical_u64();
+        let mut v = p;
+        let mut b = 1u64;
+        let mut c = 0u64;
 
         while u != 1 && v != 1 {
             while u.is_even() {
-                u >>= 1;
+                u /= 2;
                 if b.is_even() {
-                    b >>= 1;
+                    b /= 2;
                 } else {
                     // b = (b + p)/2, avoiding overflow
-                    b = (b >> 1) + (Self::ORDER >> 1) + 1;
+                    b = (b / 2) + (p / 2) + 1;
                 }
             }
 
             while v.is_even() {
-                v >>= 1;
+                v /= 2;
                 if c.is_even() {
-                    c >>= 1;
+                    c /= 2;
                 } else {
                     // c = (c + p)/2, avoiding overflow
-                    c = (c >> 1) + (Self::ORDER >> 1) + 1;
+                    c = (c / 2) + (p / 2) + 1;
                 }
             }
 
-            if u < v {
-                v -= u;
-                if c < b {
-                    c += Self::ORDER;
-                }
-                c -= b;
-            } else {
+            if u >= v {
                 u -= v;
-                if b < c {
-                    b += Self::ORDER;
+                // b -= c
+                let (mut diff, under) = b.overflowing_sub(c);
+                if under {
+                    diff = diff.overflowing_add(p).0;
                 }
-                b -= c;
+                b = diff;
+            } else {
+                v -= u;
+                // c -= b
+                let (mut diff, under) = c.overflowing_sub(b);
+                if under {
+                    diff = diff.overflowing_add(p).0;
+                }
+                c = diff;
             }
         }
 
-        Some(Self(if u == 1 {
+        let inverse = Self(if u == 1 {
             b
         } else {
             c
-        }))
+        });
+
+        // Should change to debug_assert_eq; using assert_eq as an extra precaution for now until
+        // we're more confident the impl is correct.
+        assert_eq!(*self * inverse, Self::ONE);
+        Some(inverse)
     }
 
     #[inline]

src/field/field.rs

@@ -3,11 +3,13 @@ use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssi
 use rand::Rng;
 use rand::rngs::OsRng;
 use crate::util::bits_u64;
+use std::hash::Hash;
 
 /// A finite field with prime order less than 2^64.
 pub trait Field: 'static
 + Copy
 + Eq
++ Hash
 + Neg<Output=Self>
 + Add<Self, Output=Self>
 + AddAssign<Self>
@@ -89,11 +91,13 @@ pub trait Field: 'static
     fn primitive_root_of_unity(n_power: usize) -> Self {
         assert!(n_power <= Self::TWO_ADICITY);
         let base = Self::POWER_OF_TWO_GENERATOR;
+        // TODO: Just repeated squaring should be a bit faster, to avoid conditionals.
         base.exp(Self::from_canonical_u64(1u64 << (Self::TWO_ADICITY - n_power)))
     }
 
+    /// Computes a multiplicative subgroup whose order is known in advance.
     fn cyclic_subgroup_known_order(generator: Self, order: usize) -> Vec<Self> {
-        let mut subgroup = Vec::new();
+        let mut subgroup = Vec::with_capacity(order);
         let mut current = Self::ONE;
         for _i in 0..order {
             subgroup.push(current);
@@ -102,6 +106,14 @@ pub trait Field: 'static
         subgroup
     }
 
+    /// Computes a coset of a multiplicative subgroup whose order is known in advance.
+    fn cyclic_subgroup_coset_known_order(generator: Self, shift: Self, order: usize) -> Vec<Self> {
+        let subgroup = Self::cyclic_subgroup_known_order(generator, order);
+        subgroup.into_iter()
+            .map(|x| x * shift)
+            .collect()
+    }
+
     fn to_canonical_u64(&self) -> u64;
 
     fn from_canonical_u64(n: u64) -> Self;
@@ -131,6 +143,10 @@ pub trait Field: 'static
         self.exp(Self::from_canonical_usize(power))
     }
 
+    fn powers(&self) -> Powers<Self> {
+        Powers { base: *self, current: Self::ONE }
+    }
+
     fn rand_from_rng<R: Rng>(rng: &mut R) -> Self {
         Self::from_canonical_u64(rng.gen_range(0, Self::ORDER))
     }
@@ -139,3 +155,19 @@ pub trait Field: 'static
         Self::rand_from_rng(&mut OsRng)
     }
 }
+
+/// An iterator over the powers of a certain base element `b`: `b^0, b^1, b^2, ...`.
+pub struct Powers<F: Field> {
+    base: F,
+    current: F,
+}
+
+impl<F: Field> Iterator for Powers<F> {
+    type Item = F;
+
+    fn next(&mut self) -> Option<F> {
+        let result = self.current;
+        self.current *= self.base;
+        Some(result)
+    }
+}

src/field/mod.rs

@@ -2,6 +2,7 @@ pub(crate) mod crandall_field;
 pub(crate) mod field;
 pub(crate) mod field_search;
 pub(crate) mod fft;
+pub(crate) mod cosets;
 
 #[cfg(test)]
 mod field_testing;

src/main.rs

@@ -26,7 +26,6 @@ mod gates;
 mod generator;
 mod gmimc;
 mod hash;
-mod partition;
 mod plonk_challenger;
 mod plonk_common;
 mod polynomial;

src/partition.rs

@@ -1,23 +0,0 @@
-use rand::SeedableRng;
-use rand_chacha::ChaCha8Rng;
-
-use crate::field::field::Field;
-
-/// Returns `k_i`, the multiplier used in `S_ID_i` in the context of Plonk's permutation argument.
-// TODO: This is copied from plonky1, but may need revisiting. Is the random approach still OK now
-// that our field is smaller?
-pub(crate) fn get_subgroup_shift<F: Field>(i: usize) -> F {
-    // The optimized variant of Plonk's permutation argument calls for NUM_ROUTED_WIRES shifts,
-    // k_1, ..., k_n, which result in distinct cosets. The paper suggests a method which is
-    // fairly straightforward when only three shifts are needed, but seems a bit complex and
-    // expensive if more are needed.
-
-    // We will "cheat" and just use random field elements. Since our subgroup has |F*|/degree
-    // possible cosets, the probability of a collision is negligible for large fields.
-
-    // Unlike what's shown in the Plonk paper, we do not set k_1=1 to "randomize" the
-    // sigmas polynomials evaluations and making them fit in both fields with high probability.
-    // TODO: Go back to k_1=1 if we change the way we deal with values not fitting in both fields.
-    let mut rng = ChaCha8Rng::seed_from_u64(i as u64);
-    F::rand_from_rng(&mut rng)
-}