MockBackend: Implement later-stage witness generation (#2168)

Ticks the third item off #2152

This should make the mock prover feature-complete!

---------

Co-authored-by: Thibaut Schaeffer <[email protected]>

backend/src/mock/connection_constraint_checker.rs

@@ -155,6 +155,7 @@ impl<F: FieldElement> Connection<F> {
 pub struct ConnectionConstraintChecker<'a, F: FieldElement> {
     pub connections: &'a [Connection<F>],
     pub machines: BTreeMap<String, Machine<'a, F>>,
+    pub challenges: &'a BTreeMap<u64, F>,
 }
 
 impl<'a, F: FieldElement> ConnectionConstraintChecker<'a, F> {
@@ -247,7 +248,11 @@ impl<'a, F: FieldElement> ConnectionConstraintChecker<'a, F> {
         (0..machine.size)
             .into_par_iter()
             .filter_map(|row| {
-                let variables = Variables { machine, row };
+                let variables = Variables {
+                    machine,
+                    row,
+                    challenges: self.challenges,
+                };
                 let mut evaluator =
                     ExpressionEvaluator::new(&variables, &machine.intermediate_definitions);
                 let result = evaluator.evaluate(&selected_expressions.selector).unwrap();

backend/src/mock/evaluator.rs

@@ -1,4 +1,6 @@
-use powdr_ast::analyzed::PolynomialType;
+use std::collections::BTreeMap;
+
+use powdr_ast::analyzed::{Challenge, PolynomialType};
 use powdr_executor::witgen::{AffineResult, AlgebraicVariable, SymbolicVariables};
 use powdr_number::FieldElement;
 
@@ -7,6 +9,7 @@ use super::machine::Machine;
 pub struct Variables<'a, F> {
     pub machine: &'a Machine<'a, F>,
     pub row: usize,
+    pub challenges: &'a BTreeMap<u64, F>,
 }
 
 impl<'a, F: FieldElement> Variables<'a, F> {
@@ -31,4 +34,8 @@ impl<'a, F: FieldElement> SymbolicVariables<F> for &Variables<'a, F> {
     fn value<'b>(&self, var: AlgebraicVariable<'b>) -> AffineResult<AlgebraicVariable<'b>, F> {
         Ok(self.constant_value(var).into())
     }
+
+    fn challenge<'b>(&self, challenge: &'b Challenge) -> AffineResult<AlgebraicVariable<'b>, F> {
+        Ok(self.challenges[&challenge.id].into())
+    }
 }

backend/src/mock/machine.rs

@@ -3,7 +3,7 @@ use std::collections::BTreeMap;
 use itertools::Itertools;
 use powdr_ast::analyzed::{AlgebraicExpression, AlgebraicReferenceThin, Analyzed, PolyID};
 use powdr_backend_utils::{machine_fixed_columns, machine_witness_columns};
-use powdr_executor::constant_evaluator::VariablySizedColumn;
+use powdr_executor::{constant_evaluator::VariablySizedColumn, witgen::WitgenCallback};
 use powdr_number::{DegreeType, FieldElement};
 
 /// A collection of columns with self-contained constraints.
@@ -22,8 +22,10 @@ impl<'a, F: FieldElement> Machine<'a, F> {
         witness: &'a [(String, Vec<F>)],
         fixed: &'a [(String, VariablySizedColumn<F>)],
         pil: &'a Analyzed<F>,
+        witgen_callback: &WitgenCallback<F>,
+        challenges: &BTreeMap<u64, F>,
     ) -> Option<Self> {
-        let witness = machine_witness_columns(witness, pil, &machine_name);
+        let mut witness = machine_witness_columns(witness, pil, &machine_name);
         let size = witness
             .iter()
             .map(|(_, v)| v.len())
@@ -36,6 +38,12 @@ impl<'a, F: FieldElement> Machine<'a, F> {
             return None;
         }
 
+        for stage in 1..pil.stage_count() {
+            log::debug!("Generating stage-{stage} witness for machine {machine_name}");
+            witness =
+                witgen_callback.next_stage_witness(pil, &witness, challenges.clone(), stage as u8);
+        }
+
         let fixed = machine_fixed_columns(fixed, pil);
         let fixed = fixed.get(&(size as DegreeType)).unwrap();
 

backend/src/mock/mod.rs

@@ -1,9 +1,19 @@
-use std::{collections::BTreeMap, io, marker::PhantomData, path::PathBuf, sync::Arc};
+use std::{
+    collections::BTreeMap,
+    hash::{DefaultHasher, Hash, Hasher},
+    io,
+    marker::PhantomData,
+    path::PathBuf,
+    sync::Arc,
+};
 
 use connection_constraint_checker::{Connection, ConnectionConstraintChecker};
 use machine::Machine;
 use polynomial_constraint_checker::PolynomialConstraintChecker;
-use powdr_ast::analyzed::Analyzed;
+use powdr_ast::{
+    analyzed::{AlgebraicExpression, Analyzed},
+    parsed::visitor::AllChildren,
+};
 use powdr_executor::{constant_evaluator::VariablySizedColumn, witgen::WitgenCallback};
 use powdr_number::{DegreeType, FieldElement};
 
@@ -70,37 +80,61 @@ impl<F: FieldElement> Backend<F> for MockBackend<F> {
         &self,
         witness: &[(String, Vec<F>)],
         prev_proof: Option<Proof>,
-        _witgen_callback: WitgenCallback<F>,
+        witgen_callback: WitgenCallback<F>,
     ) -> Result<Proof, Error> {
         if prev_proof.is_some() {
             unimplemented!();
         }
 
+        let challenges = self
+            .machine_to_pil
+            .values()
+            .flat_map(|pil| pil.identities.iter())
+            .flat_map(|identity| identity.all_children())
+            .filter_map(|expr| match expr {
+                AlgebraicExpression::Challenge(challenge) => {
+                    // Use the hash of the ID as the challenge.
+                    // This way, if the same challenge is used by different machines, they will
+                    // have the same value.
+                    let mut hasher = DefaultHasher::new();
+                    challenge.id.hash(&mut hasher);
+                    Some((challenge.id, F::from(hasher.finish())))
+                }
+                _ => None,
+            })
+            .collect::<BTreeMap<_, _>>();
+
         let machines = self
             .machine_to_pil
+            // TODO: We should probably iterate in parallel, because Machine::try_new might generate
+            // later-stage witnesses, which is expensive.
+            // However, for now, doing it sequentially simplifies debugging.
             .iter()
-            .filter_map(|(machine, pil)| {
-                Machine::try_new(machine.clone(), witness, &self.fixed, pil)
+            .filter_map(|(machine_name, pil)| {
+                Machine::try_new(
+                    machine_name.clone(),
+                    witness,
+                    &self.fixed,
+                    pil,
+                    &witgen_callback,
+                    &challenges,
+                )
             })
             .map(|machine| (machine.machine_name.clone(), machine))
             .collect::<BTreeMap<_, _>>();
 
-        let mut is_ok = true;
-        for (_, machine) in machines.iter() {
-            let result = PolynomialConstraintChecker::new(machine).check();
-            is_ok &= !result.has_errors();
-        }
-
-        is_ok &= ConnectionConstraintChecker {
+        let is_ok = machines.values().all(|machine| {
+            !PolynomialConstraintChecker::new(machine, &challenges)
+                .check()
+                .has_errors()
+        }) && ConnectionConstraintChecker {
             connections: &self.connections,
             machines,
+            challenges: &challenges,
         }
         .check()
         .is_ok();
 
-        // TODO:
-        // - Check later-stage witness
-
         match is_ok {
             true => Ok(Vec::new()),
             false => Err(Error::BackendError("Constraint check failed".to_string())),

backend/src/mock/polynomial_constraint_checker.rs

@@ -14,11 +14,15 @@ use super::machine::Machine;
 
 pub struct PolynomialConstraintChecker<'a, F> {
     machine: &'a Machine<'a, F>,
+    challenges: &'a BTreeMap<u64, F>,
 }
 
 impl<'a, F: FieldElement> PolynomialConstraintChecker<'a, F> {
-    pub fn new(machine: &'a Machine<'a, F>) -> Self {
-        Self { machine }
+    pub fn new(machine: &'a Machine<'a, F>, challenges: &'a BTreeMap<u64, F>) -> Self {
+        Self {
+            machine,
+            challenges,
+        }
     }
 
     pub fn check(&self) -> MachineResult<'a, F> {
@@ -54,6 +58,7 @@ impl<'a, F: FieldElement> PolynomialConstraintChecker<'a, F> {
         let variables = Variables {
             machine: self.machine,
             row,
+            challenges: self.challenges,
         };
         let mut evaluator =
             ExpressionEvaluator::new(&variables, &self.machine.intermediate_definitions);

pipeline/tests/asm.rs

@@ -39,7 +39,7 @@ fn block_machine_exact_number_of_rows_asm() {
 fn challenges_asm() {
     let f = "asm/challenges.asm";
     let pipeline: Pipeline<GoldilocksField> = make_simple_prepared_pipeline(f);
-    // TODO Mock prover doesn't support this test yet.
+    test_mock_backend(pipeline.clone());
     test_plonky3_pipeline(pipeline);
 }
 
@@ -69,7 +69,7 @@ fn secondary_block_machine_add2() {
 fn second_phase_hint() {
     let f = "asm/second_phase_hint.asm";
     let pipeline: Pipeline<GoldilocksField> = make_simple_prepared_pipeline(f);
-    // TODO Mock prover doesn't support this test yet.
+    test_mock_backend(pipeline.clone());
     test_plonky3_pipeline(pipeline);
 }
 
@@ -179,15 +179,15 @@ fn block_to_block_empty_submachine() {
 fn block_to_block_with_bus_monolithic() {
     let f = "asm/block_to_block_with_bus.asm";
     let pipeline: Pipeline<GoldilocksField> = make_simple_prepared_pipeline(f);
-    // TODO Mock prover doesn't support this test yet.
+    test_mock_backend(pipeline.clone());
     test_plonky3_pipeline(pipeline);
 }
 
 #[test]
 fn block_to_block_with_bus_different_sizes() {
     let f = "asm/block_to_block_with_bus_different_sizes.asm";
     let pipeline: Pipeline<GoldilocksField> = make_simple_prepared_pipeline(f);
-    // TODO Mock prover doesn't support this test yet.
+    test_mock_backend(pipeline.clone());
     test_plonky3_pipeline(pipeline);
 }
 
@@ -205,7 +205,7 @@ fn block_to_block_with_bus_composite() {
     use powdr_pipeline::test_util::test_halo2_with_backend_variant;
     let f = "asm/block_to_block_with_bus.asm";
     let pipeline = make_simple_prepared_pipeline(f);
-    // TODO Mock prover doesn't support this test yet.
+    test_mock_backend(pipeline.clone());
     test_halo2_with_backend_variant(pipeline, BackendVariant::Composite);
 }
 

pipeline/tests/powdr_std.rs

@@ -181,40 +181,40 @@ fn memory_small_test() {
 fn permutation_via_challenges() {
     let f = "std/permutation_via_challenges.asm";
     let pipeline: Pipeline<GoldilocksField> = make_simple_prepared_pipeline(f);
+    test_mock_backend(pipeline.clone());
     test_plonky3_pipeline(pipeline);
-    // TODO Mock prover doesn't support this test yet.
 }
 
 #[test]
 fn lookup_via_challenges() {
     let f = "std/lookup_via_challenges.asm";
     let pipeline: Pipeline<GoldilocksField> = make_simple_prepared_pipeline(f);
+    test_mock_backend(pipeline.clone());
     test_plonky3_pipeline(pipeline);
-    // TODO Mock prover doesn't support this test yet.
 }
 
 #[test]
 fn lookup_via_challenges_range_constraint() {
     let f = "std/lookup_via_challenges_range_constraint.asm";
     let pipeline: Pipeline<GoldilocksField> = make_simple_prepared_pipeline(f);
+    test_mock_backend(pipeline.clone());
     test_plonky3_pipeline(pipeline);
-    // TODO Mock prover doesn't support this test yet.
 }
 
 #[test]
 fn bus_lookup() {
     let f = "std/bus_lookup.asm";
     let pipeline: Pipeline<GoldilocksField> = make_simple_prepared_pipeline(f);
+    test_mock_backend(pipeline.clone());
     test_plonky3_pipeline(pipeline);
-    // TODO Mock prover doesn't support this test yet.
 }
 
 #[test]
 fn bus_permutation() {
     let f = "std/bus_permutation.asm";
     let pipeline: Pipeline<GoldilocksField> = make_simple_prepared_pipeline(f);
+    test_mock_backend(pipeline.clone());
     test_plonky3_pipeline(pipeline);
-    // TODO Mock prover doesn't support this test yet.
 }
 
 #[test]