Add permutation check to MockProver

src/dev.rs

@@ -1,11 +1,10 @@
 //! Tools for developing circuits.
 
 use ff::Field;
-use std::collections::HashMap;
 
 use crate::{
     arithmetic::{FieldExt, Group},
-    plonk::{Any, Assignment, Circuit, Column, ConstraintSystem, Error},
+    plonk::{permutation, Any, Assignment, Circuit, Column, ConstraintSystem, Error},
     poly::{EvaluationDomain, LagrangeCoeff, Polynomial},
 };
 
@@ -19,6 +18,12 @@ pub enum VerifyFailure {
     Gate { gate_index: usize, row: usize },
     /// A lookup input did not exist in its corresponding table.
     Lookup { lookup_index: usize, row: usize },
+    /// A permutation did not preserve the original value of a cell.
+    Permutation {
+        perm_index: usize,
+        column: usize,
+        row: usize,
+    },
 }
 
 /// A test
@@ -34,7 +39,7 @@ pub struct MockProver<F: Group> {
     // The aux cells in the circuit, arranged as [column][row].
     aux: Vec<Polynomial<F, LagrangeCoeff>>,
 
-    permutations: HashMap<usize, Vec<(Cell, Cell)>>,
+    permutations: Vec<permutation::keygen::Assembly>,
 }
 
 impl<F: Field + Group> Assignment<F> for MockProver<F> {
@@ -76,11 +81,12 @@ impl<F: Field + Group> Assignment<F> for MockProver<F> {
         right_column: usize,
         right_row: usize,
     ) -> Result<(), crate::plonk::Error> {
-        self.permutations
-            .entry(permutation)
-            .or_default()
-            .push((Cell(left_column, left_row), Cell(right_column, right_row)));
-        Ok(())
+        // Check bounds first
+        if permutation >= self.permutations.len() {
+            return Err(Error::BoundsFailure);
+        }
+
+        self.permutations[permutation].copy(left_column, left_row, right_column, right_row)
     }
 }
 
@@ -90,6 +96,8 @@ impl<F: FieldExt> MockProver<F> {
         circuit: &ConcreteCircuit,
         aux: Vec<Polynomial<F, LagrangeCoeff>>,
     ) -> Result<Self, Error> {
+        let n = 1 << k;
+
         let mut cs = ConstraintSystem::default();
         let config = ConcreteCircuit::configure(&mut cs);
 
@@ -112,15 +120,20 @@ impl<F: FieldExt> MockProver<F> {
 
         let fixed = vec![domain.empty_lagrange(); cs.num_fixed_columns];
         let advice = vec![domain.empty_lagrange(); cs.num_advice_columns];
+        let permutations = cs
+            .permutations
+            .iter()
+            .map(|p| permutation::keygen::Assembly::new(n as usize, p))
+            .collect();
 
         let mut prover = MockProver {
-            n: 1 << k,
+            n,
             domain,
             cs,
             fixed,
             advice,
             aux,
-            permutations: HashMap::default(),
+            permutations,
         };
 
         circuit.synthesize(&mut prover, config)?;
@@ -192,6 +205,33 @@ impl<F: FieldExt> MockProver<F> {
             }
         }
 
+        // Check that permutations preserve the original values of the cells.
+        for (perm_index, assembly) in self.permutations.iter().enumerate() {
+            // Original values of columns involved in the permutation
+            let original = self.cs.permutations[perm_index]
+                .get_columns()
+                .iter()
+                .map(|c| self.advice[c.index()].clone())
+                .collect::<Vec<_>>();
+
+            // Iterate over each column of the permutation
+            for (column, values) in assembly.mapping.iter().enumerate() {
+                // Iterate over each row of the column to check that the cell's
+                // value is preserved by the mapping.
+                for (row, cell) in values.iter().enumerate() {
+                    let original_cell = original[column][row];
+                    let permuted_cell = original[cell.0][cell.1];
+                    if original_cell != permuted_cell {
+                        return Err(VerifyFailure::Permutation {
+                            perm_index,
+                            column,
+                            row,
+                        });
+                    }
+                }
+            }
+        }
+
         // TODO: Implement the rest of the verification checks.
 
         Ok(())

src/plonk.rs

@@ -14,7 +14,7 @@ use crate::transcript::ChallengeScalar;
 mod circuit;
 mod keygen;
 mod lookup;
-mod permutation;
+pub(crate) mod permutation;
 mod vanishing;
 
 mod prover;

src/plonk/keygen.rs

@@ -104,7 +104,7 @@ where
         permutations: cs
             .permutations
             .iter()
-            .map(|p| permutation::keygen::Assembly::new(params, p))
+            .map(|p| permutation::keygen::Assembly::new(params.n as usize, p))
             .collect(),
         _marker: std::marker::PhantomData,
     };

src/plonk/permutation.rs

@@ -37,6 +37,10 @@ impl Argument {
         // - z(omega^{-1} X) \prod (p(X) + \delta^i \beta X + \gamma)
         std::cmp::max(self.columns.len() + 1, 2)
     }
+
+    pub(crate) fn get_columns(&self) -> Vec<Column<Advice>> {
+        self.columns.clone()
+    }
 }
 
 /// The verifying key for a single permutation argument.

src/plonk/permutation/keygen.rs

@@ -15,19 +15,19 @@ pub(crate) struct AssemblyHelper<C: CurveAffine> {
 }
 
 pub(crate) struct Assembly {
-    mapping: Vec<Vec<(usize, usize)>>,
+    pub(crate) mapping: Vec<Vec<(usize, usize)>>,
     aux: Vec<Vec<(usize, usize)>>,
     sizes: Vec<Vec<usize>>,
 }
 
 impl Assembly {
-    pub(crate) fn new<C: CurveAffine>(params: &Params<C>, p: &Argument) -> Self {
+    pub(crate) fn new(n: usize, p: &Argument) -> Self {
         // Initialize the copy vector to keep track of copy constraints in all
         // the permutation arguments.
         let mut columns = vec![];
         for i in 0..p.columns.len() {
             // Computes [(i, 0), (i, 1), ..., (i, n - 1)]
-            columns.push((0..params.n).map(|j| (i, j as usize)).collect());
+            columns.push((0..n).map(|j| (i, j)).collect());
         }
 
         // Before any equality constraints are applied, every cell in the permutation is
@@ -36,7 +36,7 @@ impl Assembly {
         Assembly {
             mapping: columns.clone(),
             aux: columns,
-            sizes: vec![vec![1usize; params.n as usize]; p.columns.len()],
+            sizes: vec![vec![1usize; n]; p.columns.len()],
         }
     }