Single step processor.

executor/src/witgen/data_structures/mutable_state.rs

@@ -42,20 +42,6 @@ impl<'a, T: FieldElement, Q: QueryCallback<T>> MutableState<'a, T, Q> {
         }
     }
 
-    #[cfg(test)]
-    pub fn get_machine(&self, substring: &str) -> &RefCell<KnownMachine<'a, T>> {
-        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<String, Vec<T>> {

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]

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;

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<CanProcess: CanProcessCall<T> + Clone>(
+        &self,
+        can_process: CanProcess,
+    ) -> Result<Vec<Effect<T, Variable>>, 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<T: FieldElement> FixedEvaluator<T> for &SingleStepProcessor<'_, T> {
+    fn evaluate(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option<T> {
+        // 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<Vec<Effect<GoldilocksField, Variable>>, 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."
+        );
+    }
+}

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
+    }
 }