diff --git a/executor/src/witgen/data_structures/mutable_state.rs b/executor/src/witgen/data_structures/mutable_state.rs index c8777af7eb..37564084bb 100644 --- a/executor/src/witgen/data_structures/mutable_state.rs +++ b/executor/src/witgen/data_structures/mutable_state.rs @@ -42,20 +42,6 @@ impl<'a, T: FieldElement, Q: QueryCallback> MutableState<'a, T, Q> { } } - #[cfg(test)] - pub fn get_machine(&self, substring: &str) -> &RefCell> { - use itertools::Itertools; - - self.machines - .iter() - .filter(|m| m.borrow().name().contains(substring)) - .exactly_one() - .map_err(|e| { - format!("Expected exactly one machine with substring '{substring}', but found {e}.") - }) - .unwrap() - } - /// Runs the first machine (unless there are no machines) end returns the generated columns. /// The first machine might call other machines, which is handled automatically. pub fn run(self) -> HashMap> { diff --git a/executor/src/witgen/jit/block_machine_processor.rs b/executor/src/witgen/jit/block_machine_processor.rs index fa8ad9effc..c0b1067625 100644 --- a/executor/src/witgen/jit/block_machine_processor.rs +++ b/executor/src/witgen/jit/block_machine_processor.rs @@ -195,31 +195,35 @@ mod test { let (fixed_data, retained_identities) = global_constraints::set_global_constraints(fixed_data, &analyzed.identities); let machines = MachineExtractor::new(&fixed_data).split_out_machines(retained_identities); - let mutable_state = MutableState::new(machines.into_iter(), &|_| { - Err("Query not implemented".to_string()) - }); - - let machine = mutable_state.get_machine(machine_name); - let ((machine_parts, block_size, latch_row), connection_ids) = match *machine.borrow() { - KnownMachine::BlockMachine(ref m) => (m.machine_info(), m.identity_ids()), - _ => panic!("Expected a block machine"), + let [KnownMachine::BlockMachine(machine)] = machines + .iter() + .filter(|m| m.name().contains(machine_name)) + .collect_vec() + .as_slice() + else { + panic!("Expected exactly one matching block machine") }; - assert_eq!(connection_ids.len(), 1); - + let (machine_parts, block_size, latch_row) = machine.machine_info(); + assert_eq!(machine_parts.connections.len(), 1); + let connection_id = *machine_parts.connections.keys().next().unwrap(); let processor = BlockMachineProcessor { fixed_data: &fixed_data, - machine_parts, + machine_parts: machine_parts.clone(), block_size, latch_row, }; + let mutable_state = MutableState::new(machines.into_iter(), &|_| { + Err("Query not implemented".to_string()) + }); + let known_values = BitVec::from_iter( (0..num_inputs) .map(|_| true) .chain((0..num_outputs).map(|_| false)), ); - processor.generate_code(&mutable_state, connection_ids[0], &known_values) + processor.generate_code(&mutable_state, connection_id, &known_values) } #[test] diff --git a/executor/src/witgen/jit/mod.rs b/executor/src/witgen/jit/mod.rs index 91df4a9c6f..c2b33e43f2 100644 --- a/executor/src/witgen/jit/mod.rs +++ b/executor/src/witgen/jit/mod.rs @@ -3,6 +3,7 @@ mod block_machine_processor; mod compiler; mod effect; pub(crate) mod function_cache; +mod single_step_processor; mod symbolic_expression; mod variable; pub(crate) mod witgen_inference; diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs new file mode 100644 index 0000000000..c2780c5241 --- /dev/null +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -0,0 +1,172 @@ +#![allow(dead_code)] +use std::collections::HashSet; + +use itertools::Itertools; +use powdr_ast::analyzed::AlgebraicReference; +use powdr_number::FieldElement; + +use crate::witgen::{machines::MachineParts, FixedData}; + +use super::{ + effect::Effect, + variable::{Cell, Variable}, + witgen_inference::{CanProcessCall, FixedEvaluator, WitgenInference}, +}; + +/// A processor for generating JIT code that computes the next row from the previous row. +pub struct SingleStepProcessor<'a, T: FieldElement> { + fixed_data: &'a FixedData<'a, T>, + machine_parts: MachineParts<'a, T>, +} + +impl<'a, T: FieldElement> SingleStepProcessor<'a, T> { + pub fn new(fixed_data: &'a FixedData<'a, T>, machine_parts: MachineParts<'a, T>) -> Self { + SingleStepProcessor { + fixed_data, + machine_parts, + } + } + + pub fn generate_code + Clone>( + &self, + can_process: CanProcess, + ) -> Result>, String> { + // All witness columns in row 0 are known. + let known_variables = self.machine_parts.witnesses.iter().map(|id| { + Variable::Cell(Cell { + column_name: self.fixed_data.column_name(id).to_string(), + id: id.id, + row_offset: 0, + }) + }); + let mut witgen = WitgenInference::new(self.fixed_data, self, known_variables); + + let mut complete = HashSet::new(); + for iteration in 0.. { + let mut progress = false; + + for id in &self.machine_parts.identities { + if !complete.contains(&id.id()) { + let row_offset = if id.contains_next_ref() { 0 } else { 1 }; + let result = witgen.process_identity(can_process.clone(), id, row_offset); + if result.complete { + complete.insert(id.id()); + } + progress |= result.progress; + } + } + if !progress { + log::debug!("Finishing after {iteration} iterations"); + break; + } + } + + // Check that we could derive all witness values in the next row. + let unknown_witnesses = self + .machine_parts + .witnesses + .iter() + .filter(|wit| { + !witgen.is_known(&Variable::Cell(Cell { + column_name: self.fixed_data.column_name(wit).to_string(), + id: wit.id, + row_offset: 1, + })) + }) + .sorted() + .collect_vec(); + + let missing_identities = self.machine_parts.identities.len() - complete.len(); + if unknown_witnesses.is_empty() && missing_identities == 0 { + Ok(witgen.code()) + } else { + Err(format!( + "Unable to derive algorithm to compute values for witness columns in the next row for the following columns:\n{}\nand {missing_identities} identities are missing.", + unknown_witnesses.iter().map(|wit| self.fixed_data.column_name(wit)).format(", ") + )) + } + } +} + +impl FixedEvaluator for &SingleStepProcessor<'_, T> { + fn evaluate(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option { + // We can only return something here if the fixed column is constant + // in the region we are considering. + // This might be the case if we know we are not evaluating the first or the last + // row, but this is not yet implemented. + None + } +} + +#[cfg(test)] +mod test { + + use itertools::Itertools; + + use powdr_number::GoldilocksField; + + use crate::witgen::{ + data_structures::mutable_state::MutableState, + global_constraints, + jit::{ + effect::{format_code, Effect}, + test_util::read_pil, + }, + machines::{machine_extractor::MachineExtractor, KnownMachine, Machine}, + FixedData, + }; + + use super::{SingleStepProcessor, Variable}; + + fn generate_single_step( + input_pil: &str, + machine_name: &str, + ) -> Result>, String> { + let (analyzed, fixed_col_vals) = read_pil(input_pil); + + let fixed_data = FixedData::new(&analyzed, &fixed_col_vals, &[], Default::default(), 0); + let (fixed_data, retained_identities) = + global_constraints::set_global_constraints(fixed_data, &analyzed.identities); + let machines = MachineExtractor::new(&fixed_data).split_out_machines(retained_identities); + let [KnownMachine::DynamicMachine(machine)] = machines + .iter() + .filter(|m| m.name().contains(machine_name)) + .collect_vec() + .as_slice() + else { + panic!("Expected exactly one matching dynamic machine") + }; + let machine_parts = machine.machine_parts().clone(); + let mutable_state = MutableState::new(machines.into_iter(), &|_| { + Err("Query not implemented".to_string()) + }); + + SingleStepProcessor { + fixed_data: &fixed_data, + machine_parts, + } + .generate_code(&mutable_state) + } + + #[test] + fn fib() { + let input = "namespace M(256); let X; let Y; X' = Y; Y' = X + Y;"; + let code = generate_single_step(input, "M").unwrap(); + assert_eq!( + format_code(&code), + "M::X[1] = M::Y[0];\nM::Y[1] = (M::X[0] + M::Y[0]);" + ); + } + + #[test] + fn no_progress() { + let input = "namespace M(256); let X; let Y; X' = X;"; + let err = generate_single_step(input, "M").err().unwrap(); + assert_eq!( + err.to_string(), + "Unable to derive algorithm to compute values for witness columns in the next row for the following columns:\n\ + M::Y\n\ + and 0 identities are missing." + ); + } +} diff --git a/executor/src/witgen/machines/dynamic_machine.rs b/executor/src/witgen/machines/dynamic_machine.rs index 26f25f1bd8..762d662cd1 100644 --- a/executor/src/witgen/machines/dynamic_machine.rs +++ b/executor/src/witgen/machines/dynamic_machine.rs @@ -278,4 +278,9 @@ impl<'a, T: FieldElement> DynamicMachine<'a, T> { ); } } + + #[cfg(test)] + pub fn machine_parts(&self) -> &MachineParts<'a, T> { + &self.parts + } }