Add MockProver for developing circuits

src/dev.rs

@@ -0,0 +1,199 @@
+//! Tools for developing circuits.
+
+use ff::Field;
+use std::collections::HashMap;
+
+use crate::{
+    arithmetic::{FieldExt, Group},
+    plonk::{Any, Assignment, Circuit, Column, ConstraintSystem, Error},
+    poly::{EvaluationDomain, LagrangeCoeff, Polynomial},
+};
+
+#[derive(Debug)]
+struct Cell(usize, usize);
+
+/// The reasons why a particular circuit is not satisfied.
+#[derive(Debug, PartialEq)]
+pub enum VerifyFailure {
+    /// A gate was not satisfied for a particular row.
+    Gate { gate_index: usize, row: usize },
+    /// A lookup input did not exist in its corresponding table.
+    Lookup { lookup_index: usize, row: usize },
+}
+
+/// A test
+pub struct MockProver<F: Group> {
+    n: u32,
+    domain: EvaluationDomain<F>,
+    cs: ConstraintSystem<F>,
+
+    // The fixed cells in the circuit, arranged as [column][row].
+    fixed: Vec<Polynomial<F, LagrangeCoeff>>,
+    // The advice cells in the circuit, arranged as [column][row].
+    advice: Vec<Polynomial<F, LagrangeCoeff>>,
+    // The aux cells in the circuit, arranged as [column][row].
+    aux: Vec<Polynomial<F, LagrangeCoeff>>,
+
+    permutations: HashMap<usize, Vec<(Cell, Cell)>>,
+}
+
+impl<F: Field + Group> Assignment<F> for MockProver<F> {
+    fn assign_advice(
+        &mut self,
+        column: crate::plonk::Column<crate::plonk::Advice>,
+        row: usize,
+        to: impl FnOnce() -> Result<F, crate::plonk::Error>,
+    ) -> Result<(), crate::plonk::Error> {
+        *self
+            .advice
+            .get_mut(column.index())
+            .and_then(|v| v.get_mut(row))
+            .ok_or(Error::BoundsFailure)? = to()?;
+
+        Ok(())
+    }
+
+    fn assign_fixed(
+        &mut self,
+        column: crate::plonk::Column<crate::plonk::Fixed>,
+        row: usize,
+        to: impl FnOnce() -> Result<F, crate::plonk::Error>,
+    ) -> Result<(), crate::plonk::Error> {
+        *self
+            .fixed
+            .get_mut(column.index())
+            .and_then(|v| v.get_mut(row))
+            .ok_or(Error::BoundsFailure)? = to()?;
+
+        Ok(())
+    }
+
+    fn copy(
+        &mut self,
+        permutation: usize,
+        left_column: usize,
+        left_row: usize,
+        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(())
+    }
+}
+
+impl<F: FieldExt> MockProver<F> {
+    pub fn run<ConcreteCircuit: Circuit<F>>(
+        k: u32,
+        circuit: &ConcreteCircuit,
+        aux: Vec<Polynomial<F, LagrangeCoeff>>,
+    ) -> Result<Self, Error> {
+        let mut cs = ConstraintSystem::default();
+        let config = ConcreteCircuit::configure(&mut cs);
+
+        // The permutation argument will serve alongside the gates, so must be
+        // accounted for.
+        let mut degree = cs
+            .permutations
+            .iter()
+            .map(|p| p.required_degree())
+            .max()
+            .unwrap_or(1);
+
+        // Account for each gate to ensure our quotient polynomial is the
+        // correct degree and that our extended domain is the right size.
+        for poly in cs.gates.iter() {
+            degree = std::cmp::max(degree, poly.degree());
+        }
+
+        let domain = EvaluationDomain::new(degree as u32, k);
+
+        let fixed = vec![domain.empty_lagrange(); cs.num_fixed_columns];
+        let advice = vec![domain.empty_lagrange(); cs.num_advice_columns];
+
+        let mut prover = MockProver {
+            n: 1 << k,
+            domain,
+            cs,
+            fixed,
+            advice,
+            aux,
+            permutations: HashMap::default(),
+        };
+
+        circuit.synthesize(&mut prover, config)?;
+
+        Ok(prover)
+    }
+
+    /// Returns `Ok(())` if this `MockProver` is satisfied, or an error indicating the
+    /// reason that the circuit is not satisfied.
+    pub fn verify(&self) -> Result<(), VerifyFailure> {
+        let n = self.n as i32;
+
+        // Check that all gates are satisfied for all rows.
+        for (gate_index, gate) in self.cs.gates.iter().enumerate() {
+            // We iterate from n..2n so we can just reduce to handle wrapping.
+            for row in n..(2 * n) {
+                if gate.evaluate(
+                    &|index| {
+                        let (column, at) = self.cs.fixed_queries[index];
+                        let resolved_row = (row + at.0) % n;
+                        self.fixed[column.index()][resolved_row as usize].clone()
+                    },
+                    &|index| {
+                        let (column, at) = self.cs.advice_queries[index];
+                        let resolved_row = (row + at.0) % n;
+                        self.advice[column.index()][resolved_row as usize].clone()
+                    },
+                    &|index| {
+                        let (column, at) = self.cs.aux_queries[index];
+                        let resolved_row = (row + at.0) % n;
+                        self.aux[column.index()][resolved_row as usize].clone()
+                    },
+                    &|a, b| a + &b,
+                    &|a, b| a * &b,
+                    &|a, scalar| a * scalar,
+                ) != F::zero()
+                {
+                    return Err(VerifyFailure::Gate {
+                        gate_index,
+                        row: row as usize,
+                    });
+                }
+            }
+        }
+
+        // Check that all lookups exist in their respective tables.
+        for (lookup_index, lookup) in self.cs.lookups.iter().enumerate() {
+            for input_row in 0..n {
+                let load = |column: &Column<Any>, row| match column.column_type() {
+                    Any::Fixed => self.fixed[column.index()][row as usize],
+                    Any::Advice => self.advice[column.index()][row as usize],
+                    Any::Aux => self.aux[column.index()][row as usize],
+                };
+
+                let inputs: Vec<_> = lookup
+                    .input_columns
+                    .iter()
+                    .map(|c| load(c, input_row))
+                    .collect();
+                if !(0..n)
+                    .map(|table_row| lookup.table_columns.iter().map(move |c| load(c, table_row)))
+                    .any(|table_row| table_row.eq(inputs.iter().cloned()))
+                {
+                    return Err(VerifyFailure::Lookup {
+                        lookup_index,
+                        row: input_row as usize,
+                    });
+                }
+            }
+        }
+
+        // TODO: Implement the rest of the verification checks.
+
+        Ok(())
+    }
+}

src/lib.rs

@@ -20,3 +20,6 @@ pub mod poly;
 pub mod transcript;
 
 pub mod model;
+
+#[cfg(test)]
+mod dev;

src/plonk.rs

@@ -119,6 +119,7 @@ type ChallengeX<F> = ChallengeScalar<F, X>;
 #[test]
 fn test_proving() {
     use crate::arithmetic::{Curve, FieldExt};
+    use crate::dev::MockProver;
     use crate::pasta::{EqAffine, Fp, Fq};
     use crate::poly::commitment::{Blind, Params};
     use crate::transcript::DummyHash;
@@ -462,6 +463,13 @@ fn test_proving() {
         .commit_lagrange(&pubinputs, Blind::default())
         .to_affine();
 
+    // Check this circuit is satisfied.
+    let prover = match MockProver::run(K, &circuit, vec![pubinputs.clone()]) {
+        Ok(prover) => prover,
+        Err(e) => panic!("{:?}", e),
+    };
+    assert_eq!(prover.verify(), Ok(()));
+
     for _ in 0..100 {
         // Create a proof
         let proof = Proof::create::<DummyHash<Fq>, DummyHash<Fp>, _>(