diff --git a/backend/src/mock/connection_constraint_checker.rs b/backend/src/mock/connection_constraint_checker.rs index 81b3813960..f71e9b2c30 100644 --- a/backend/src/mock/connection_constraint_checker.rs +++ b/backend/src/mock/connection_constraint_checker.rs @@ -155,6 +155,7 @@ impl Connection { pub struct ConnectionConstraintChecker<'a, F: FieldElement> { pub connections: &'a [Connection], pub machines: BTreeMap>, + pub challenges: &'a BTreeMap, } 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(); diff --git a/backend/src/mock/evaluator.rs b/backend/src/mock/evaluator.rs index e852c0f228..853aef961d 100644 --- a/backend/src/mock/evaluator.rs +++ b/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, } impl<'a, F: FieldElement> Variables<'a, F> { @@ -31,4 +34,8 @@ impl<'a, F: FieldElement> SymbolicVariables for &Variables<'a, F> { fn value<'b>(&self, var: AlgebraicVariable<'b>) -> AffineResult, F> { Ok(self.constant_value(var).into()) } + + fn challenge<'b>(&self, challenge: &'b Challenge) -> AffineResult, F> { + Ok(self.challenges[&challenge.id].into()) + } } diff --git a/backend/src/mock/machine.rs b/backend/src/mock/machine.rs index 2bcf0a67f1..1e634224b0 100644 --- a/backend/src/mock/machine.rs +++ b/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)], fixed: &'a [(String, VariablySizedColumn)], pil: &'a Analyzed, + witgen_callback: &WitgenCallback, + challenges: &BTreeMap, ) -> Option { - 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(); diff --git a/backend/src/mock/mod.rs b/backend/src/mock/mod.rs index c0c5e8478d..6b37fd1dd6 100644 --- a/backend/src/mock/mod.rs +++ b/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 Backend for MockBackend { &self, witness: &[(String, Vec)], prev_proof: Option, - _witgen_callback: WitgenCallback, + witgen_callback: WitgenCallback, ) -> Result { 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::>(); + 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::>(); - 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())), diff --git a/backend/src/mock/polynomial_constraint_checker.rs b/backend/src/mock/polynomial_constraint_checker.rs index f6ee3b47d8..194f246c9b 100644 --- a/backend/src/mock/polynomial_constraint_checker.rs +++ b/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, } 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) -> 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); diff --git a/pipeline/tests/asm.rs b/pipeline/tests/asm.rs index 04c1e327a8..99ecf47326 100644 --- a/pipeline/tests/asm.rs +++ b/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 = 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 = make_simple_prepared_pipeline(f); - // TODO Mock prover doesn't support this test yet. + test_mock_backend(pipeline.clone()); test_plonky3_pipeline(pipeline); } @@ -179,7 +179,7 @@ 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 = make_simple_prepared_pipeline(f); - // TODO Mock prover doesn't support this test yet. + test_mock_backend(pipeline.clone()); test_plonky3_pipeline(pipeline); } @@ -187,7 +187,7 @@ fn block_to_block_with_bus_monolithic() { fn block_to_block_with_bus_different_sizes() { let f = "asm/block_to_block_with_bus_different_sizes.asm"; let pipeline: Pipeline = 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); } diff --git a/pipeline/tests/powdr_std.rs b/pipeline/tests/powdr_std.rs index 92fb207462..dd385488ed 100644 --- a/pipeline/tests/powdr_std.rs +++ b/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 = 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 = 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 = 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 = 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 = make_simple_prepared_pipeline(f); + test_mock_backend(pipeline.clone()); test_plonky3_pipeline(pipeline); - // TODO Mock prover doesn't support this test yet. } #[test]