Split coset step up so that we can query wires at multiple spots.

src/plonk/circuit.rs

@@ -111,9 +111,9 @@ pub trait Circuit<F: Field> {
 #[derive(Clone, Debug)]
 pub enum Polynomial<F> {
     /// This is a fixed wire queried at a certain relative location
-    Fixed(FixedWire, isize),
+    Fixed(FixedWire, i32),
     /// This is an advice (witness) wire queried at a certain relative location
-    Advice(AdviceWire, isize),
+    Advice(AdviceWire, i32),
     /// This is the sum of two polynomials
     Sum(Box<Polynomial<F>>, Box<Polynomial<F>>),
     /// This is the product of two polynomials
@@ -125,8 +125,8 @@ pub enum Polynomial<F> {
 impl<F: Field> Polynomial<F> {
     fn evaluate<T>(
         &self,
-        fixed_wire: &impl Fn(FixedWire, isize) -> T,
-        advice_wire: &impl Fn(AdviceWire, isize) -> T,
+        fixed_wire: &impl Fn(FixedWire, i32) -> T,
+        advice_wire: &impl Fn(AdviceWire, i32) -> T,
         sum: &impl Fn(T, T) -> T,
         product: &impl Fn(T, T) -> T,
         scaled: &impl Fn(T, F) -> T,
@@ -193,8 +193,8 @@ pub struct MetaCircuit<F> {
     pub(crate) num_advice_wires: usize,
     // permutations: Vec<Vec<Wire>>,
     gates: Vec<Polynomial<F>>,
-    advice_queries: HashMap<(AdviceWire, isize), usize>,
-    fixed_queries: HashMap<(FixedWire, isize), usize>,
+    advice_queries: HashMap<(AdviceWire, i32), usize>,
+    fixed_queries: HashMap<(FixedWire, i32), usize>,
     // num_queries: usize,
 }
 
@@ -212,7 +212,7 @@ impl<F: Field> Default for MetaCircuit<F> {
 
 impl<F: Field> MetaCircuit<F> {
     /// Query a fixed wire at a relative position
-    pub fn query_fixed(&mut self, wire: FixedWire, at: isize) -> Polynomial<F> {
+    pub fn query_fixed(&mut self, wire: FixedWire, at: i32) -> Polynomial<F> {
         let len = self.fixed_queries.len();
         self.fixed_queries.entry((wire, at)).or_insert_with(|| len);
 
@@ -220,7 +220,7 @@ impl<F: Field> MetaCircuit<F> {
     }
 
     /// Query an advice wire at a relative position
-    pub fn query_advice(&mut self, wire: AdviceWire, at: isize) -> Polynomial<F> {
+    pub fn query_advice(&mut self, wire: AdviceWire, at: i32) -> Polynomial<F> {
         let len = self.advice_queries.len();
         self.advice_queries.entry((wire, at)).or_insert_with(|| len);
 

src/plonk/domain.rs

@@ -8,6 +8,7 @@ pub struct EvaluationDomain<G: Group> {
     n: u64,
     k: u32,
     extended_k: u32,
+    omega: G::Scalar,
     omega_inv: G::Scalar,
     extended_omega: G::Scalar,
     extended_omega_inv: G::Scalar,
@@ -91,6 +92,7 @@ impl<G: Group> EvaluationDomain<G> {
             n,
             k,
             extended_k,
+            omega,
             omega_inv,
             extended_omega,
             extended_omega_inv,
@@ -103,32 +105,42 @@ impl<G: Group> EvaluationDomain<G> {
         }
     }
 
-    /// This takes us from an n-length vector into the coset evaluation domain.
-    /// Also returns the polynomial.
+    /// This takes us from an n-length vector into the coefficient form.
     ///
     /// This function will panic if the provided vector is not the correct
     /// length.
-    pub fn obtain_coset(&self, mut a: Vec<G>) -> (Vec<G>, Vec<G>) {
+    pub fn obtain_poly(&self, mut a: Vec<G>) -> Vec<G> {
         assert_eq!(a.len(), 1 << self.k);
 
         // Perform inverse FFT to obtain the polynomial in coefficient form
         Self::ifft(&mut a, self.omega_inv, self.k, self.ifft_divisor);
 
-        // Keep this polynomial around; we'll need to evaluate it at arbitrary
-        // points later.
-        let old = a.clone();
+        a
+    }
 
-        // Distributes powers so that an FFT will move us into the coset
-        // evaluation domain.
-        Self::distribute_powers(&mut a, self.g_coset);
+    /// This takes us from an n-length coefficient vector into the coset
+    /// evaluation domain.
+    ///
+    /// This function will panic if the provided vector is not the correct
+    /// length.
+    pub fn obtain_coset(&self, mut a: Vec<G>, index: i32) -> Vec<G> {
+        assert_eq!(a.len(), 1 << self.k);
 
-        // Resize to account for the quotient polynomial's size
+        assert!(index != i32::MIN);
+        if index == 0 {
+            Self::distribute_powers_zeta(&mut a, self.g_coset);
+        } else {
+            let mut g = G::Scalar::ZETA;
+            if index > 0 {
+                g *= &self.omega.pow_vartime(&[index as u64, 0, 0, 0]);
+            } else {
+                g *= &self.omega_inv.pow_vartime(&[index.abs() as u64, 0, 0, 0]);
+            }
+            Self::distribute_powers(&mut a, g);
+        }
         a.resize(1 << self.extended_k, G::group_zero());
-
-        // Move into coset evaluation domain
         best_fft(&mut a, self.extended_omega, self.extended_k);
-
-        (a, old)
+        a
     }
 
     /// This takes us from the coset evaluation domain and gets us the quotient
@@ -176,7 +188,7 @@ impl<G: Group> EvaluationDomain<G> {
         h_poly
     }
 
-    fn distribute_powers(mut a: &mut [G], g: G::Scalar) {
+    fn distribute_powers_zeta(mut a: &mut [G], g: G::Scalar) {
         let coset_powers = [g, g.square()];
         parallelize(&mut a, |a, mut index| {
             for a in a {
@@ -190,6 +202,16 @@ impl<G: Group> EvaluationDomain<G> {
         });
     }
 
+    fn distribute_powers(mut a: &mut [G], g: G::Scalar) {
+        parallelize(&mut a, |a, index| {
+            let mut cur = g.pow_vartime(&[index as u64, 0, 0, 0]);
+            for a in a {
+                a.group_scale(&cur);
+                cur *= &g;
+            }
+        });
+    }
+
     fn ifft(a: &mut [G], omega_inv: G::Scalar, log_n: u32, divisor: G::Scalar) {
         best_fft(a, omega_inv, log_n);
         parallelize(a, |a, _| {

src/plonk/prover.rs

@@ -160,15 +160,24 @@ impl<C: CurveAffine> Proof<C> {
 
         let domain = &srs.domain;
 
-        let (a_coset, a_poly) = domain.obtain_coset(witness.a);
-        let (b_coset, b_poly) = domain.obtain_coset(witness.b);
-        let (c_coset, c_poly) = domain.obtain_coset(witness.c);
-        let (d_coset, d_poly) = domain.obtain_coset(witness.d);
+        let a_poly = domain.obtain_poly(witness.a);
+        let b_poly = domain.obtain_poly(witness.b);
+        let c_poly = domain.obtain_poly(witness.c);
+        let d_poly = domain.obtain_poly(witness.d);
+
+        let a_coset = domain.obtain_coset(a_poly.clone(), 0);
+        let b_coset = domain.obtain_coset(b_poly.clone(), 0);
+        let c_coset = domain.obtain_coset(c_poly.clone(), 0);
+        let d_coset = domain.obtain_coset(d_poly.clone(), 0);
 
         let advice_polys: Vec<_> = witness
             .advice
             .into_iter()
-            .map(|poly| domain.obtain_coset(poly))
+            .map(|poly| {
+                let poly = domain.obtain_poly(poly);
+                let coset = domain.obtain_coset(poly.clone(), 0);
+                (poly, coset)
+            })
             .collect();
 
         // (a * sa) + (b * sb) + (a * sm * b) + (d * sd) - (c * sc)

src/plonk/srs.rs

@@ -118,24 +118,33 @@ impl<C: CurveAffine> SRS<C> {
 
         let domain = EvaluationDomain::new(GATE_DEGREE, params.k);
 
-        let sa = domain.obtain_coset(assembly.sa);
-        let sb = domain.obtain_coset(assembly.sb);
-        let sc = domain.obtain_coset(assembly.sc);
-        let sd = domain.obtain_coset(assembly.sd);
-        let sm = domain.obtain_coset(assembly.sm);
+        let sa_poly = domain.obtain_poly(assembly.sa);
+        let sb_poly = domain.obtain_poly(assembly.sb);
+        let sc_poly = domain.obtain_poly(assembly.sc);
+        let sd_poly = domain.obtain_poly(assembly.sd);
+        let sm_poly = domain.obtain_poly(assembly.sm);
+        let sa_coset = domain.obtain_coset(sa_poly.clone(), 0);
+        let sb_coset = domain.obtain_coset(sb_poly.clone(), 0);
+        let sc_coset = domain.obtain_coset(sc_poly.clone(), 0);
+        let sd_coset = domain.obtain_coset(sd_poly.clone(), 0);
+        let sm_coset = domain.obtain_coset(sm_poly.clone(), 0);
 
         let fixed_polys = assembly
             .fixed
             .into_iter()
-            .map(|poly| domain.obtain_coset(poly))
+            .map(|poly| {
+                let coeffs = domain.obtain_poly(poly);
+                let coset = domain.obtain_coset(coeffs.clone(), 0);
+                (coeffs, coset)
+            })
             .collect();
 
         Ok(SRS {
-            sa,
-            sb,
-            sc,
-            sd,
-            sm,
+            sa: (sa_coset, sa_poly),
+            sb: (sb_coset, sb_poly),
+            sc: (sc_coset, sc_poly),
+            sd: (sd_coset, sd_poly),
+            sm: (sm_coset, sm_poly),
             sa_commitment,
             sb_commitment,
             sc_commitment,