Simplify permutations field of ConstraintSystem

Co-authored-by: therealyingtong <yingtong@electriccoin.co>

src/plonk/circuit.rs

@@ -170,13 +170,8 @@ pub struct ConstraintSystem<F> {
     pub(crate) rotations: BTreeMap<Rotation, PointIndex>,
 
     // Vector of permutation arguments, where each corresponds to a set of wires
-    // that are involved in a permutation argument, as well as the corresponding
-    // query index for each wire. As an example, we could have a permutation
-    // argument between wires (A, B, C) which allows copy constraints to be
-    // enforced between advice wire values in A, B and C, and another
-    // permutation between wires (B, C, D) which allows the same with D instead
-    // of A.
-    pub(crate) permutations: Vec<Vec<(AdviceWire, usize)>>,
+    // that are involved in a permutation argument.
+    pub(crate) permutations: Vec<Vec<AdviceWire>>,
 }
 
 impl<F: Field> Default for ConstraintSystem<F> {
@@ -202,16 +197,16 @@ impl<F: Field> ConstraintSystem<F> {
     /// Add a permutation argument for some advice wires
     pub fn permutation(&mut self, wires: &[AdviceWire]) -> usize {
         let index = self.permutations.len();
-        if index == 0 {
+        if self.permutations.is_empty() {
             let at = Rotation(-1);
             let len = self.rotations.len();
             self.rotations.entry(at).or_insert(PointIndex(len));
         }
-        let wires = wires
-            .iter()
-            .map(|&wire| (wire, self.query_advice_index(wire, 0)))
-            .collect();
-        self.permutations.push(wires);
+
+        for wire in wires {
+            self.query_advice_index(*wire, 0);
+        }
+        self.permutations.push(wires.to_vec());
 
         index
     }
@@ -242,7 +237,18 @@ impl<F: Field> ConstraintSystem<F> {
         Expression::Fixed(self.query_fixed_index(wire, at))
     }
 
-    fn query_advice_index(&mut self, wire: AdviceWire, at: i32) -> usize {
+    pub(crate) fn get_advice_query_index(&self, wire: AdviceWire, at: i32) -> usize {
+        let at = Rotation(at);
+        for (index, advice_query) in self.advice_queries.iter().enumerate() {
+            if advice_query == &(wire, at) {
+                return index;
+            }
+        }
+
+        panic!("get_advice_query_index called for non-existant query");
+    }
+
+    pub(crate) fn query_advice_index(&mut self, wire: AdviceWire, at: i32) -> usize {
         let at = Rotation(at);
         {
             let len = self.rotations.len();

src/plonk/prover.rs

@@ -187,7 +187,7 @@ impl<C: CurveAffine> Proof<C> {
             let mut modified_advice = vec![C::Scalar::one(); params.n as usize];
 
             // Iterate over each wire of the permutation
-            for (&(wire, _), permuted_wire_values) in wires.iter().zip(permuted_values.iter()) {
+            for (&wire, permuted_wire_values) in wires.iter().zip(permuted_values.iter()) {
                 parallelize(&mut modified_advice, |modified_advice, start| {
                     for ((modified_advice, advice_value), permuted_advice_value) in modified_advice
                         .iter_mut()
@@ -219,7 +219,7 @@ impl<C: CurveAffine> Proof<C> {
             // Iterate over each wire again, this time finishing the computation
             // of the entire fraction by computing the numerators
             let mut deltaomega = C::Scalar::one();
-            for &(wire, _) in wires.iter() {
+            for &wire in wires.iter() {
                 let omega = domain.get_omega();
                 parallelize(&mut modified_advice, |modified_advice, start| {
                     let mut deltaomega = deltaomega * &omega.pow_vartime(&[start as u64, 0, 0, 0]);
@@ -320,7 +320,7 @@ impl<C: CurveAffine> Proof<C> {
             let mut left = permutation_product_cosets[permutation_index].clone();
             for (advice, permutation) in wires
                 .iter()
-                .map(|&(_, index)| &advice_cosets[index])
+                .map(|&wire| &advice_cosets[pk.vk.cs.get_advice_query_index(wire, 0)])
                 .zip(pk.permutation_cosets[permutation_index].iter())
             {
                 parallelize(&mut left, |left, start| {
@@ -337,7 +337,10 @@ impl<C: CurveAffine> Proof<C> {
             let mut right = permutation_product_cosets_inv[permutation_index].clone();
             let mut current_delta = x_0 * &C::Scalar::ZETA;
             let step = domain.get_extended_omega();
-            for advice in wires.iter().map(|&(_, index)| &advice_cosets[index]) {
+            for advice in wires
+                .iter()
+                .map(|&wire| &advice_cosets[pk.vk.cs.get_advice_query_index(wire, 0)])
+            {
                 parallelize(&mut right, move |right, start| {
                     let mut beta_term = current_delta * &step.pow_vartime(&[start as u64, 0, 0, 0]);
                     for (right, advice) in right.iter_mut().zip(advice[start..].iter()) {

src/plonk/verifier.rs

@@ -128,7 +128,7 @@ impl<'a, C: CurveAffine> Proof<C> {
             let mut left = self.permutation_product_evals[permutation_index];
             for (advice_eval, permutation_eval) in wires
                 .iter()
-                .map(|&(_, query_index)| self.advice_evals[query_index])
+                .map(|&wire| self.advice_evals[vk.cs.get_advice_query_index(wire, 0)])
                 .zip(self.permutation_evals[permutation_index].iter())
             {
                 left *= &(advice_eval + &(x_0 * permutation_eval) + &x_1);
@@ -138,7 +138,7 @@ impl<'a, C: CurveAffine> Proof<C> {
             let mut current_delta = x_0 * &x_3;
             for advice_eval in wires
                 .iter()
-                .map(|&(_, query_index)| self.advice_evals[query_index])
+                .map(|&wire| self.advice_evals[vk.cs.get_advice_query_index(wire, 0)])
             {
                 right *= &(advice_eval + &current_delta + &x_1);
                 current_delta *= &C::Scalar::DELTA;