Single step with branching

executor/src/witgen/jit/compiler.rs

@@ -1,8 +1,8 @@
-use std::{ffi::c_void, iter, mem, sync::Arc};
+use std::{cmp::Ordering, ffi::c_void, iter, mem, sync::Arc};
 
-use auto_enums::auto_enum;
 use itertools::Itertools;
 use libloading::Library;
+use powdr_ast::indent;
 use powdr_number::{FieldElement, KnownField};
 
 use crate::witgen::{
@@ -16,7 +16,7 @@ use crate::witgen::{
 };
 
 use super::{
-    effect::{Assertion, Effect},
+    effect::{Assertion, BranchCondition, Effect},
     symbolic_expression::{BinaryOperator, BitOperator, SymbolicExpression, UnaryOperator},
     variable::Variable,
 };
@@ -157,10 +157,11 @@ fn witgen_code<T: FieldElement>(
             format!("    let {var_name} = {value};")
         })
         .format("\n");
-    let main_code = effects.iter().map(format_effect).format("\n");
+    let main_code = format_effects(effects);
     let vars_known = effects
         .iter()
         .flat_map(written_vars_in_effect)
+        .map(|(var, _)| var)
         .collect_vec();
     let store_values = vars_known
         .iter()
@@ -224,26 +225,51 @@ extern "C" fn witgen(
 }
 
 /// Returns an iterator over all variables written to in the effect.
-#[auto_enum(Iterator)]
+/// 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>,
-) -> impl Iterator<Item = &Variable> + '_ {
+) -> Box<dyn Iterator<Item = (&Variable, bool)> + '_> {
     match effect {
-        Effect::Assignment(var, _) => iter::once(var),
+        Effect::Assignment(var, _) => Box::new(iter::once((var, false))),
         Effect::RangeConstraint(..) => unreachable!(),
-        Effect::Assertion(..) => iter::empty(),
-        Effect::MachineCall(_, arguments) => arguments.iter().flat_map(|e| match e {
-            MachineCallArgument::Unknown(v) => Some(v),
+        Effect::Assertion(..) => Box::new(iter::empty()),
+        Effect::MachineCall(_, arguments) => Box::new(arguments.iter().flat_map(|e| match e {
+            MachineCallArgument::Unknown(v) => Some((v, true)),
             MachineCallArgument::Known(_) => None,
-        }),
+        })),
+        Effect::Branch(_, first, second) => Box::new(
+            first
+                .iter()
+                .chain(second)
+                .flat_map(|e| written_vars_in_effect(e)),
+        ),
     }
 }
 
-fn format_effect<T: FieldElement>(effect: &Effect<T, Variable>) -> String {
+pub fn format_effects<T: FieldElement>(effects: &[Effect<T, Variable>]) -> String {
+    format_effects_inner(effects, true)
+}
+
+fn format_effects_inner<T: FieldElement>(
+    effects: &[Effect<T, Variable>],
+    is_top_level: bool,
+) -> String {
+    indent(
+        effects
+            .iter()
+            .map(|effect| format_effect(effect, is_top_level))
+            .join("\n"),
+        1,
+    )
+}
+
+fn format_effect<T: FieldElement>(effect: &Effect<T, Variable>, is_top_level: bool) -> String {
     match effect {
         Effect::Assignment(var, e) => {
             format!(
-                "    let {} = {};",
+                "{}{} = {};",
+                if is_top_level { "let " } else { "" },
                 variable_to_string(var),
                 format_expression(e)
             )
@@ -256,7 +282,7 @@ fn format_effect<T: FieldElement>(effect: &Effect<T, Variable>) -> String {
             rhs,
             expected_equal,
         }) => format!(
-            "    assert!({} {} {});",
+            "assert!({} {} {});",
             format_expression(lhs),
             if *expected_equal { "==" } else { "!=" },
             format_expression(rhs)
@@ -268,7 +294,9 @@ fn format_effect<T: FieldElement>(effect: &Effect<T, Variable>) -> String {
                 .map(|a| match a {
                     MachineCallArgument::Unknown(v) => {
                         let var_name = variable_to_string(v);
-                        result_vars.push(var_name.clone());
+                        if is_top_level {
+                            result_vars.push(var_name.clone());
+                        }
                         format!("LookupCell::Output(&mut {var_name})")
                     }
                     MachineCallArgument::Known(v) => {
@@ -279,11 +307,40 @@ fn format_effect<T: FieldElement>(effect: &Effect<T, Variable>) -> String {
                 .to_string();
             let var_decls = result_vars
                 .iter()
-                .map(|var_name| format!("    let mut {var_name} = FieldElement::default();"))
-                .format("\n");
+                .map(|var_name| format!("let mut {var_name} = FieldElement::default();\n"))
+                .format("");
+            format!(
+                "{var_decls}assert!(call_machine(mutable_state, {id}, MutSlice::from((&mut [{args}]).as_mut_slice())));"
+            )
+        }
+        Effect::Branch(condition, first, second) => {
+            let var_decls = if is_top_level {
+                // We need to declare all assigned variables at top level,
+                // so that they are available after the branches.
+                first
+                    .iter()
+                    .chain(second)
+                    .flat_map(|e| written_vars_in_effect(e))
+                    .sorted()
+                    .dedup()
+                    .map(|(v, needs_mut)| {
+                        let v = variable_to_string(v);
+                        if needs_mut {
+                            format!("let mut {v} = FieldElement::default();\n")
+                        } else {
+                            format!("let {v};\n")
+                        }
+                    })
+                    .format("")
+                    .to_string()
+            } else {
+                "".to_string()
+            };
             format!(
-                "{var_decls}
-    assert!(call_machine(mutable_state, {id}, MutSlice::from((&mut [{args}]).as_mut_slice())));"
+                "{var_decls}if {} {{\n{}\n}} else {{\n{}\n}}",
+                format_condition(condition),
+                format_effects_inner(first, false),
+                format_effects_inner(second, false)
             )
         }
     }
@@ -319,6 +376,16 @@ fn format_expression<T: FieldElement>(e: &SymbolicExpression<T, Variable>) -> St
     }
 }
 
+fn format_condition<T: FieldElement>(condition: &BranchCondition<T, Variable>) -> String {
+    let var = format!("IntType::from({})", variable_to_string(&condition.variable));
+    let (min, max) = condition.first_branch.range();
+    match min.cmp(&max) {
+        Ordering::Equal => format!("{var} == {min}",),
+        Ordering::Less => format!("{min} <= {var} && {var} <= {max}"),
+        Ordering::Greater => format!("{var} <= {min} || {var} >= {max}"),
+    }
+}
+
 /// Returns the name of a local (stack) variable for the given expression variable.
 fn variable_to_string(v: &Variable) -> String {
     match v {
@@ -423,6 +490,7 @@ mod tests {
 
     use crate::witgen::jit::variable::Cell;
     use crate::witgen::jit::variable::MachineCallReturnVariable;
+    use crate::witgen::range_constraints::RangeConstraint;
 
     use super::*;
 
@@ -780,4 +848,71 @@ extern \"C\" fn witgen(
         assert_eq!(data[1], GoldilocksField::from(18));
         assert_eq!(data[2], GoldilocksField::from(0));
     }
+
+    #[test]
+    fn branches() {
+        let x = param(0);
+        let y = param(1);
+        let mut x_val: GoldilocksField = 7.into();
+        let mut y_val: GoldilocksField = 9.into();
+        let effects = vec![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 f = compile_effects(0, 1, &[x], &effects).unwrap();
+        let mut data = vec![];
+        let mut known = vec![];
+
+        let mut params = vec![LookupCell::Input(&x_val), LookupCell::Output(&mut y_val)];
+        let params = WitgenFunctionParams {
+            data: data.as_mut_slice().into(),
+            known: known.as_mut_ptr(),
+            row_offset: 0,
+            params: params.as_mut_slice().into(),
+            mutable_state: std::ptr::null(),
+            call_machine: no_call_machine,
+        };
+        (f.function)(params);
+        assert_eq!(y_val, GoldilocksField::from(8));
+
+        x_val = 2.into();
+        let mut params = vec![LookupCell::Input(&x_val), LookupCell::Output(&mut y_val)];
+        let params = WitgenFunctionParams {
+            data: data.as_mut_slice().into(),
+            known: known.as_mut_ptr(),
+            row_offset: 0,
+            params: params.as_mut_slice().into(),
+            mutable_state: std::ptr::null(),
+            call_machine: no_call_machine,
+        };
+        (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/effect.rs

@@ -1,23 +1,30 @@
+use std::cmp::Ordering;
+
 use itertools::Itertools;
+use powdr_ast::indent;
 use powdr_number::FieldElement;
 
 use crate::witgen::range_constraints::RangeConstraint;
 
 use super::{symbolic_expression::SymbolicExpression, variable::Variable};
 
 /// The effect of solving a symbolic equation.
+#[derive(Clone, PartialEq, Eq)]
 pub enum Effect<T: FieldElement, V> {
     /// Variable can be assigned a value.
     Assignment(V, SymbolicExpression<T, V>),
     /// We learnt a new range constraint on variable.
     RangeConstraint(V, RangeConstraint<T>),
     /// A run-time assertion. If this fails, we have conflicting constraints.
     Assertion(Assertion<T, V>),
-    /// a call to a different machine.
+    /// A call to a different machine.
     MachineCall(u64, Vec<MachineCallArgument<T, V>>),
+    /// A branch on a variable.
+    Branch(BranchCondition<T, V>, Vec<Effect<T, V>>, Vec<Effect<T, V>>),
 }
 
 /// A run-time assertion. If this fails, we have conflicting constraints.
+#[derive(Clone, PartialEq, Eq)]
 pub struct Assertion<T: FieldElement, V> {
     pub lhs: SymbolicExpression<T, V>,
     pub rhs: SymbolicExpression<T, V>,
@@ -52,11 +59,19 @@ impl<T: FieldElement, V> Assertion<T, V> {
     }
 }
 
+#[derive(Clone, PartialEq, Eq)]
 pub enum MachineCallArgument<T: FieldElement, V> {
     Known(SymbolicExpression<T, V>),
     Unknown(V),
 }
 
+#[derive(Clone, PartialEq, Eq)]
+pub struct BranchCondition<T: FieldElement, V> {
+    pub variable: V,
+    pub first_branch: RangeConstraint<T>,
+    pub second_branch: RangeConstraint<T>,
+}
+
 /// Helper function to render a list of effects. Used for informational purposes only.
 pub fn format_code<T: FieldElement>(effects: &[Effect<T, Variable>]) -> String {
     effects
@@ -87,6 +102,23 @@ pub fn format_code<T: FieldElement>(effects: &[Effect<T, Variable>]) -> String {
             Effect::RangeConstraint(..) => {
                 panic!("Range constraints should not be part of the code.")
             }
+            Effect::Branch(condition, first, second) => {
+                let first = indent(format_code(first), 1);
+                let second = indent(format_code(second), 1);
+                let condition = format_condition(condition);
+
+                format!("if ({condition}) {{\n{first}\n}} else {{\n{second}\n}}")
+            }
         })
         .join("\n")
 }
+
+fn format_condition<T: FieldElement>(condition: &BranchCondition<T, Variable>) -> String {
+    let var = &condition.variable;
+    let (min, max) = condition.first_branch.range();
+    match min.cmp(&max) {
+        Ordering::Equal => format!("{var} == {min}"),
+        Ordering::Less => format!("{min} <= {var} && {var} <= {max}"),
+        Ordering::Greater => format!("{var} <= {min} || {var} >= {max}"),
+    }
+}

executor/src/witgen/jit/includes/field_generic_up_to_64.rs

@@ -19,6 +19,12 @@ impl From<IntType> for FieldElement {
         Self(i)
     }
 }
+impl From<FieldElement> for IntType {
+    #[inline]
+    fn from(f: FieldElement) -> Self {
+        f.0
+    }
+}
 impl std::ops::Add for FieldElement {
     type Output = Self;
     #[inline]

executor/src/witgen/jit/includes/field_goldilocks.rs

@@ -3,6 +3,7 @@
 struct GoldilocksField(u64);
 
 type FieldElement = GoldilocksField;
+type IntType = u64;
 
 const EPSILON: u64 = (1 << 32) - 1;
 
@@ -292,6 +293,13 @@ impl From<u64> for GoldilocksField {
     }
 }
 
+impl From<FieldElement> for IntType {
+    #[inline]
+    fn from(f: FieldElement) -> Self {
+        f.0
+    }
+}
+
 impl std::fmt::Display for GoldilocksField {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}", self.0)