Single step with branching

executor/src/witgen/jit/compiler.rs

@@ -225,7 +225,7 @@ extern "C" fn witgen(
 }
 
 /// Returns an iterator over all variables written to in the effect.
-/// The flag indicates if the variable is used in a lookup and thus needs
+/// The flag indicates if the variable is the return value of a machine call and thus needs
 /// to be declared mutable.
 fn written_vars_in_effect<T: FieldElement>(
     effect: &Effect<T, Variable>,
@@ -337,7 +337,7 @@ fn format_effect<T: FieldElement>(effect: &Effect<T, Variable>, is_top_level: bo
                 "".to_string()
             };
             format!(
-                "{var_decls}if {} {{\n{}}} else {{\n{}}}",
+                "{var_decls}if {} {{\n{}\n}} else {{\n{}\n}}",
                 format_condition(condition),
                 format_effects_inner(first, false),
                 format_effects_inner(second, false)
@@ -893,4 +893,26 @@ extern \"C\" fn witgen(
         (f.function)(params);
         assert_eq!(y_val, GoldilocksField::from(4));
     }
+
+    #[test]
+    fn branches_codegen() {
+        let x = param(0);
+        let y = param(1);
+        let branch_effect = Effect::Branch(
+            BranchCondition {
+                variable: x.clone(),
+                first_branch: RangeConstraint::from_range(7.into(), 20.into()),
+                second_branch: RangeConstraint::from_range(21.into(), 6.into()),
+            },
+            vec![assignment(&y, symbol(&x) + number(1))],
+            vec![assignment(&y, symbol(&x) + number(2))],
+        );
+        let expectation = "    let p_1;
+    if 7 <= IntType::from(p_0) && IntType::from(p_0) <= 20 {
+        p_1 = (p_0 + FieldElement::from(1));
+    } else {
+        p_1 = (p_0 + FieldElement::from(2));
+    }";
+        assert_eq!(format_effects(&[branch_effect]), expectation);
+    }
 }

executor/src/witgen/jit/single_step_processor.rs

@@ -10,7 +10,7 @@ use crate::witgen::{machines::MachineParts, FixedData};
 use super::{
     effect::Effect,
     variable::{Cell, Variable},
-    witgen_inference::{CanProcessCall, FixedEvaluator, WitgenInference},
+    witgen_inference::{BranchResult, CanProcessCall, FixedEvaluator, WitgenInference},
 };
 
 /// A processor for generating JIT code that computes the next row from the previous row.
@@ -45,43 +45,67 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> {
         // Check that we could derive all witness values in the next row.
         let unknown_witnesses = self
             .unknown_witness_cols_on_next_row(&witgen)
+            // Sort to get deterministic code.
             .sorted()
             .collect_vec();
 
         let missing_identities = self.machine_parts.identities.len() - complete.len();
         let code = if unknown_witnesses.is_empty() && missing_identities == 0 {
             witgen.code()
         } else {
-            let Some((most_constrained_var, _)) = witgen
+            let Some(most_constrained_var) = witgen
                 .known_variables()
                 .iter()
                 .filter_map(|var| witgen.range_constraint(var).map(|rc| (var, rc)))
                 .filter(|(_, rc)| rc.try_to_single_value().is_none())
                 .sorted()
                 .min_by_key(|(_, rc)| rc.range_width())
+                .map(|(var, _)| var.clone())
             else {
                 let incomplete_identities = self
                     .machine_parts
                     .identities
                     .iter()
                     .filter(|id| !complete.contains(&id.id()));
+                let column_errors = if unknown_witnesses.is_empty() {
+                    "".to_string()
+                } else {
+                    format!(
+                        "\nThe following columns are still missing: {}",
+                        unknown_witnesses
+                            .iter()
+                            .map(|wit| self.fixed_data.column_name(wit))
+                            .format(", ")
+                    )
+                };
+                let identity_errors = if missing_identities == 0 {
+                    "".to_string()
+                } else {
+                    format!(
+                        "\nThe following identities have not been fully processed:\n{}",
+                        incomplete_identities
+                            .map(|id| format!("    {id}"))
+                            .join("\n")
+                    )
+                };
                 return Err(format!(
                     "Unable to derive algorithm to compute values for witness columns in the next row and\n\
-                    unable to branch on a variable. The following columns are still missing:\n{}\nThe following identities have not been fully processed:\n{}",
-                    unknown_witnesses.iter().map(|wit| self.fixed_data.column_name(wit)).format(", "),
-                    incomplete_identities.map(|id| format!("    {id}")).join("\n")
+                    unable to branch on a variable.{column_errors}{identity_errors}",
                 ));
             };
 
-            let (common_code, condition, other_branch) =
-                witgen.branch_on(&most_constrained_var.clone());
+            let BranchResult {
+                common_code,
+                condition,
+                branches: [first_branch, second_branch],
+            } = witgen.branch_on(&most_constrained_var.clone());
 
             // TODO Tuning: If this fails (or also if it does not generate progress right away),
             // we could also choose a different variable to branch on.
             let left_branch_code =
-                self.generate_code_for_branch(can_process.clone(), witgen, complete.clone())?;
+                self.generate_code_for_branch(can_process.clone(), first_branch, complete.clone())?;
             let right_branch_code =
-                self.generate_code_for_branch(can_process, other_branch, complete)?;
+                self.generate_code_for_branch(can_process, second_branch, complete)?;
             if left_branch_code == right_branch_code {
                 common_code.into_iter().chain(left_branch_code).collect()
             } else {
@@ -99,6 +123,7 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> {
     }
 
     fn initialize_witgen(&self) -> WitgenInference<'a, T, NoEval> {
+        // 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(),
@@ -119,7 +144,9 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> {
         while progress {
             progress = false;
 
-            // TODO propagate known.
+            // TODO At this point, we should call a function on `witgen`
+            // to propagate known concrete values across the identities
+            // to other known (but not concrete) variables.
 
             for id in &self.machine_parts.identities {
                 if complete.contains(&id.id()) {
@@ -229,9 +256,7 @@ mod test {
         assert_eq!(
             err.to_string(),
             "Unable to derive algorithm to compute values for witness columns in the next row and\n\
-            unable to branch on a variable. The following columns are still missing:\n\
-            M::Y\n\
-            The following identities have not been fully processed:\n"
+            unable to branch on a variable.\nThe following columns are still missing: M::Y"
         );
     }
 
@@ -246,13 +271,13 @@ mod test {
         let pc: col;
 
         col fixed LINE = [0, 1] + [2]*;
-        col fixed INSTR_ADD = [0, 1] +  [0]*;
+        col fixed INSTR_ADD = [0, 1] + [0]*;
         col fixed INSTR_MUL = [1, 0] + [1]*;
 
         pc' = pc + 1;
         [ pc, instr_add, instr_mul ] in [ LINE, INSTR_ADD, INSTR_MUL ];
 
-        instr_add * (A' - (A + B)) + instr_mul * (A' - A * B)  + (1 - instr_add - instr_mul) * (A' - A) = 0;
+        instr_add * (A' - (A + B)) + instr_mul * (A' - A * B) + (1 - instr_add - instr_mul) * (A' - A) = 0;
         B' = B;
         ";
         let code = generate_single_step(input, "Main").unwrap();

executor/src/witgen/jit/witgen_inference.rs

@@ -35,7 +35,7 @@ pub struct ProcessSummary {
 /// This component can generate code that solves identities.
 /// It needs a driver that tells it which identities to process on which rows.
 #[derive(Clone)]
-pub struct WitgenInference<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> {
+pub struct WitgenInference<'a, T: FieldElement, FixedEval> {
     fixed_data: &'a FixedData<'a, T>,
     fixed_evaluator: FixedEval,
     derived_range_constraints: HashMap<Variable, RangeConstraint<T>>,
@@ -62,6 +62,16 @@ impl<T: Display> Display for Value<T> {
     }
 }
 
+/// Return type of the `branch_on` method.
+pub struct BranchResult<'a, T: FieldElement, FixedEval> {
+    /// The code common to both branches.
+    pub common_code: Vec<Effect<T, Variable>>,
+    /// The condition of the branch.
+    pub condition: BranchCondition<T, Variable>,
+    /// The two branches.
+    pub branches: [WitgenInference<'a, T, FixedEval>; 2],
+}
+
 impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, FixedEval> {
     pub fn new(
         fixed_data: &'a FixedData<'a, T>,
@@ -101,10 +111,11 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
         }
     }
 
-    pub fn branch_on(
-        &mut self,
-        variable: &Variable,
-    ) -> (Vec<Effect<T, Variable>>, BranchCondition<T, Variable>, Self) {
+    /// Splits the current inference into two copies - one where the provided variable
+    /// is in the "second half" of its range constraint and one where it is in the
+    /// "first half" of its range constraint (determined by calling the `bisect` method).
+    /// Returns the common code, the branch condition and the two branches.
+    pub fn branch_on(mut self, variable: &Variable) -> BranchResult<'a, T, FixedEval> {
         // The variable needs to be known, we need to have a range constraint but
         // it cannot be a single value.
         assert!(self.known_variables.contains(variable));
@@ -118,21 +129,21 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
 
         let (low_condition, high_condition) = rc.bisect();
 
-        let code = std::mem::take(&mut self.code);
+        let common_code = std::mem::take(&mut self.code);
         let mut low_branch = self.clone();
 
         self.add_range_constraint(variable.clone(), high_condition.clone());
         low_branch.add_range_constraint(variable.clone(), low_condition.clone());
 
-        (
-            code,
-            BranchCondition {
+        BranchResult {
+            common_code,
+            condition: BranchCondition {
                 variable: variable.clone(),
                 first_branch: high_condition,
                 second_branch: low_condition,
             },
-            low_branch,
-        )
+            branches: [self, low_branch],
+        }
     }
 
     /// Process an identity on a certain row.