Refactor permutation keygen to reflect the separate permutations

src/plonk/keygen.rs

@@ -21,7 +21,7 @@ where
 {
     struct Assembly<F: Field> {
         fixed: Vec<Polynomial<F, LagrangeCoeff>>,
-        permutations: permutation::keygen::Assembly,
+        permutations: Vec<permutation::keygen::Assembly>,
         _marker: std::marker::PhantomData<F>,
     }
 
@@ -59,8 +59,12 @@ where
             right_column: usize,
             right_row: usize,
         ) -> Result<(), Error> {
-            self.permutations
+            // Check bounds first
-                .copy(permutation, left_column, left_row, right_column, right_row)
+            if permutation >= self.permutations.len() {
+                return Err(Error::BoundsFailure);
+            }
+
+            self.permutations[permutation].copy(left_column, left_row, right_column, right_row)
         }
     }
 
@@ -97,14 +101,25 @@ where
 
     let mut assembly: Assembly<C::Scalar> = Assembly {
         fixed: vec![domain.empty_lagrange(); cs.num_fixed_columns],
-        permutations: permutation::keygen::Assembly::new(params, &cs),
+        permutations: cs
+            .permutations
+            .iter()
+            .map(|p| permutation::keygen::Assembly::new(params, p))
+            .collect(),
         _marker: std::marker::PhantomData,
     };
 
     // Synthesize the circuit to obtain SRS
     circuit.synthesize(&mut assembly, config)?;
 
-    let (permutation_pks, permutation_vks) = assembly.permutations.build_keys(params, &cs, &domain);
+    let permutation_helper = permutation::keygen::Assembly::build_helper(params, &cs, &domain);
+
+    let (permutation_pks, permutation_vks) = cs
+        .permutations
+        .iter()
+        .zip(assembly.permutations.into_iter())
+        .map(|(p, assembly)| assembly.build_keys(params, &domain, &permutation_helper, p))
+        .unzip();
 
     let fixed_commitments = assembly
         .fixed

src/plonk/permutation/keygen.rs

@@ -1,6 +1,6 @@
 use ff::Field;
 
-use super::{ProvingKey, VerifyingKey};
+use super::{Argument, ProvingKey, VerifyingKey};
 use crate::{
     arithmetic::{Curve, CurveAffine, FieldExt},
     plonk::{circuit::ConstraintSystem, Error},
@@ -10,97 +10,82 @@ use crate::{
     },
 };
 
+pub(crate) struct AssemblyHelper<C: CurveAffine> {
+    deltaomega: Vec<Vec<C::Scalar>>,
+}
+
 pub(crate) struct Assembly {
-    mapping: Vec<Vec<Vec<(usize, usize)>>>,
+    mapping: Vec<Vec<(usize, usize)>>,
-    aux: Vec<Vec<Vec<(usize, usize)>>>,
+    aux: Vec<Vec<(usize, usize)>>,
-    sizes: Vec<Vec<Vec<usize>>>,
+    sizes: Vec<Vec<usize>>,
 }
 
 impl Assembly {
-    pub(crate) fn new<C: CurveAffine>(
+    pub(crate) fn new<C: CurveAffine>(params: &Params<C>, p: &Argument) -> Self {
-        params: &Params<C>,
-        cs: &ConstraintSystem<C::Scalar>,
-    ) -> Self {
-        let mut assembly = Assembly {
-            mapping: vec![],
-            aux: vec![],
-            sizes: vec![],
-        };
-
         // Initialize the copy vector to keep track of copy constraints in all
         // the permutation arguments.
-        for p in &cs.permutations {
+        let mut columns = vec![];
-            let mut columns = vec![];
+        for i in 0..p.columns.len() {
-            for i in 0..p.columns.len() {
+            // Computes [(i, 0), (i, 1), ..., (i, n - 1)]
-                // Computes [(i, 0), (i, 1), ..., (i, n - 1)]
+            columns.push((0..params.n).map(|j| (i, j as usize)).collect());
-                columns.push((0..params.n).map(|j| (i, j as usize)).collect());
-            }
-            assembly.mapping.push(columns.clone());
-            assembly.aux.push(columns);
-            assembly
-                .sizes
-                .push(vec![vec![1usize; params.n as usize]; p.columns.len()]);
         }
 
-        assembly
+        Assembly {
+            mapping: columns.clone(),
+            aux: columns,
+            sizes: vec![vec![1usize; params.n as usize]; p.columns.len()],
+        }
     }
 
     pub(crate) fn copy(
         &mut self,
-        permutation: usize,
         left_column: usize,
         left_row: usize,
         right_column: usize,
         right_row: usize,
     ) -> Result<(), Error> {
         // Check bounds first
-        if permutation >= self.mapping.len()
+        if left_column >= self.mapping.len()
-            || left_column >= self.mapping[permutation].len()
+            || left_row >= self.mapping[left_column].len()
-            || left_row >= self.mapping[permutation][left_column].len()
+            || right_column >= self.mapping.len()
-            || right_column >= self.mapping[permutation].len()
+            || right_row >= self.mapping[right_column].len()
-            || right_row >= self.mapping[permutation][right_column].len()
         {
             return Err(Error::BoundsFailure);
         }
 
-        let mut left_cycle = self.aux[permutation][left_column][left_row];
+        let mut left_cycle = self.aux[left_column][left_row];
-        let mut right_cycle = self.aux[permutation][right_column][right_row];
+        let mut right_cycle = self.aux[right_column][right_row];
 
         if left_cycle == right_cycle {
             return Ok(());
         }
 
-        if self.sizes[permutation][left_cycle.0][left_cycle.1]
+        if self.sizes[left_cycle.0][left_cycle.1] < self.sizes[right_cycle.0][right_cycle.1] {
-            < self.sizes[permutation][right_cycle.0][right_cycle.1]
-        {
             std::mem::swap(&mut left_cycle, &mut right_cycle);
         }
 
-        self.sizes[permutation][left_cycle.0][left_cycle.1] +=
+        self.sizes[left_cycle.0][left_cycle.1] += self.sizes[right_cycle.0][right_cycle.1];
-            self.sizes[permutation][right_cycle.0][right_cycle.1];
         let mut i = right_cycle;
         loop {
-            self.aux[permutation][i.0][i.1] = left_cycle;
+            self.aux[i.0][i.1] = left_cycle;
-            i = self.mapping[permutation][i.0][i.1];
+            i = self.mapping[i.0][i.1];
             if i == right_cycle {
                 break;
             }
         }
 
-        let tmp = self.mapping[permutation][left_column][left_row];
+        let tmp = self.mapping[left_column][left_row];
-        self.mapping[permutation][left_column][left_row] =
+        self.mapping[left_column][left_row] = self.mapping[right_column][right_row];
-            self.mapping[permutation][right_column][right_row];
+        self.mapping[right_column][right_row] = tmp;
-        self.mapping[permutation][right_column][right_row] = tmp;
 
         Ok(())
     }
 
-    pub(crate) fn build_keys<C: CurveAffine>(
+    pub(crate) fn build_helper<C: CurveAffine>(
-        self,
         params: &Params<C>,
         cs: &ConstraintSystem<C::Scalar>,
         domain: &EvaluationDomain<C::Scalar>,
-    ) -> (Vec<ProvingKey<C>>, Vec<VerifyingKey<C>>) {
+    ) -> AssemblyHelper<C> {
         // Get the largest permutation argument length in terms of the number of
         // advice columns involved.
         let largest_permutation_length = cs
@@ -136,44 +121,50 @@ impl Assembly {
             }
         }
 
+        AssemblyHelper { deltaomega }
+    }
+
+    pub(crate) fn build_keys<C: CurveAffine>(
+        self,
+        params: &Params<C>,
+        domain: &EvaluationDomain<C::Scalar>,
+        helper: &AssemblyHelper<C>,
+        p: &Argument,
+    ) -> (ProvingKey<C>, VerifyingKey<C>) {
         // Compute permutation polynomials, convert to coset form and
         // pre-compute commitments for the SRS.
-        let mut pks = vec![];
+        let mut commitments = vec![];
-        let mut vks = vec![];
+        let mut permutations = vec![];
-        for (p, mapping) in cs.permutations.iter().zip(self.mapping.iter()) {
+        let mut polys = vec![];
-            let mut commitments = vec![];
+        let mut cosets = vec![];
-            let mut permutations = vec![];
+        for i in 0..p.columns.len() {
-            let mut polys = vec![];
+            // Computes the permutation polynomial based on the permutation
-            let mut cosets = vec![];
+            // description in the assembly.
-            for i in 0..p.columns.len() {
+            let mut permutation_poly = domain.empty_lagrange();
-                // Computes the permutation polynomial based on the permutation
+            for (j, p) in permutation_poly.iter_mut().enumerate() {
-                // description in the assembly.
+                let (permuted_i, permuted_j) = self.mapping[i][j];
-                let mut permutation_poly = domain.empty_lagrange();
+                *p = helper.deltaomega[permuted_i][permuted_j];
-                for (j, p) in permutation_poly.iter_mut().enumerate() {
-                    let (permuted_i, permuted_j) = mapping[i][j];
-                    *p = deltaomega[permuted_i][permuted_j];
-                }
-
-                // Compute commitment to permutation polynomial
-                commitments.push(
-                    params
-                        .commit_lagrange(&permutation_poly, Blind::default())
-                        .to_affine(),
-                );
-                // Store permutation polynomial and precompute its coset evaluation
-                permutations.push(permutation_poly.clone());
-                let poly = domain.lagrange_to_coeff(permutation_poly);
-                polys.push(poly.clone());
-                cosets.push(domain.coeff_to_extended(poly, Rotation::default()));
             }
-            vks.push(VerifyingKey { commitments });
+
-            pks.push(ProvingKey {
+            // Compute commitment to permutation polynomial
+            commitments.push(
+                params
+                    .commit_lagrange(&permutation_poly, Blind::default())
+                    .to_affine(),
+            );
+            // Store permutation polynomial and precompute its coset evaluation
+            permutations.push(permutation_poly.clone());
+            let poly = domain.lagrange_to_coeff(permutation_poly);
+            polys.push(poly.clone());
+            cosets.push(domain.coeff_to_extended(poly, Rotation::default()));
+        }
+        (
+            ProvingKey {
                 permutations,
                 polys,
                 cosets,
-            });
+            },
-        }
+            VerifyingKey { commitments },
-
+        )
-        (pks, vks)
     }
 }