Witgen inference.

executor/src/witgen/jit/affine_symbolic_expression.rs

@@ -20,6 +20,8 @@ pub enum Effect<T: FieldElement, V> {
     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.
+    MachineCall(u64, Vec<MachineCallArgument<T, V>>),
 }
 
 /// A run-time assertion. If this fails, we have conflicting constraints.
@@ -57,6 +59,32 @@ impl<T: FieldElement, V> Assertion<T, V> {
     }
 }
 
+pub enum MachineCallArgument<T: FieldElement, V> {
+    Known(SymbolicExpression<T, V>),
+    Unknown(AffineSymbolicExpression<T, V>),
+}
+
+#[derive(Default)]
+pub struct ProcessResult<T: FieldElement, V> {
+    pub effects: Vec<Effect<T, V>>,
+    pub complete: bool,
+}
+
+impl<T: FieldElement, V> ProcessResult<T, V> {
+    pub fn empty() -> Self {
+        Self {
+            effects: vec![],
+            complete: false,
+        }
+    }
+    pub fn complete(effects: Vec<Effect<T, V>>) -> Self {
+        Self {
+            effects,
+            complete: true,
+        }
+    }
+}
+
 /// Represents an expression `a_1 * x_1 + ... + a_k * x_k + offset`,
 /// where the `a_i` and `offset` are symbolic expressions, i.e. values known at run-time
 /// (which can still include variables or symbols, which are only known at run-time),
@@ -134,6 +162,15 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
         }
     }
 
+    /// If this expression contains a single unknown variable, returns it.
+    pub fn single_unknown_variable(&self) -> Option<&V> {
+        if self.coefficients.len() == 1 {
+            self.coefficients.keys().next()
+        } else {
+            None
+        }
+    }
+
     /// Tries to multiply this expression with another one.
     /// Returns `None` if the result would be quadratic, i.e.
     /// if both expressions contain unknown variables.
@@ -148,15 +185,16 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
 
     /// Solves the equation `self = 0` and returns how to compute the solution.
     /// The solution can contain assignments to multiple variables.
-    /// If no way to solve the equation has been found, returns the empty vector.
+    /// If no way to solve the equation has been found, but it still contains
+    /// unknon variables, returns an empty, incomplete result.
     /// If the equation is known to be unsolvable, returns an error.
-    pub fn solve(&self) -> Result<Vec<Effect<T, V>>, EvalError<T>> {
+    pub fn solve(&self) -> Result<ProcessResult<T, V>, EvalError<T>> {
         Ok(match self.coefficients.len() {
             0 => {
                 if self.offset.is_known_nonzero() {
                     return Err(EvalError::ConstraintUnsatisfiable(self.to_string()));
                 } else {
-                    vec![]
+                    ProcessResult::complete(vec![])
                 }
             }
             1 => {
@@ -169,43 +207,48 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
                 if coeff.is_known_nonzero() {
                     // In this case, we can always compute a solution.
                     let value = self.offset.field_div(&-coeff);
-                    vec![Effect::Assignment(var.clone(), value)]
+                    ProcessResult::complete(vec![Effect::Assignment(var.clone(), value)])
                 } else if self.offset.is_known_nonzero() {
                     // If the offset is not zero, then the coefficient must be non-zero,
                     // otherwise the constraint is violated.
                     let value = self.offset.field_div(&-coeff);
-                    vec![
+                    ProcessResult::complete(vec![
                         Assertion::assert_is_nonzero(coeff.clone()),
                         Effect::Assignment(var.clone(), value),
-                    ]
+                    ])
                 } else {
                     // If this case, we could have an equation of the form
                     // 0 * X = 0, which is valid and generates no information about X.
-                    vec![]
+                    ProcessResult::empty()
                 }
             }
             _ => {
                 let r = self.solve_bit_decomposition();
-                if !r.is_empty() {
+                if r.complete {
                     r
                 } else {
                     let negated = -self;
                     let r = negated.solve_bit_decomposition();
-                    if !r.is_empty() {
+                    if r.complete {
                         r
                     } else {
-                        self.transfer_constraints()
+                        let effects = self
+                            .transfer_constraints()
                             .into_iter()
                             .chain(negated.transfer_constraints())
-                            .collect()
+                            .collect();
+                        ProcessResult {
+                            effects,
+                            complete: false,
+                        }
                     }
                 }
             }
         })
     }
 
     /// Tries to solve a bit-decomposition equation.
-    fn solve_bit_decomposition(&self) -> Vec<Effect<T, V>> {
+    fn solve_bit_decomposition(&self) -> ProcessResult<T, V> {
         // All the coefficients need to be known numbers and the
         // variables need to be range-constrained.
         let constrained_coefficients = self
@@ -218,7 +261,7 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
             })
             .collect::<Option<Vec<_>>>();
         let Some(constrained_coefficients) = constrained_coefficients else {
-            return vec![];
+            return ProcessResult::empty();
         };
 
         // Check if they are mutually exclusive and compute assignments.
@@ -228,7 +271,7 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
             let mask = *constraint.multiple(coeff).mask();
             if !(mask & covered_bits).is_zero() {
                 // Overlapping range constraints.
-                return vec![];
+                return ProcessResult::empty();
             } else {
                 covered_bits |= mask;
             }
@@ -240,26 +283,26 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
         }
 
         if covered_bits >= T::modulus() {
-            return vec![];
+            return ProcessResult::empty();
         }
 
-        // We need to assert that the masks cover the offset,
+        // We need to assert that the masks cover "-offset",
         // otherwise the equation is not solvable.
-        // We assert offset & !masks == 0 <=> offset == offset | masks.
+        // We assert -offset & !masks == 0 <=> -offset == -offset | masks.
         // We use the latter since we cannot properly bit-negate inside the field.
         effects.push(Assertion::assert_eq(
-            self.offset.clone(),
-            &self.offset | &T::from(covered_bits).into(),
+            -&self.offset,
+            -&self.offset | T::from(covered_bits).into(),
         ));
 
-        effects
+        ProcessResult::complete(effects)
     }
 
-    fn transfer_constraints(&self) -> Vec<Effect<T, V>> {
+    fn transfer_constraints(&self) -> Option<Effect<T, V>> {
         // We are looking for X = a * Y + b * Z + ... or -X = a * Y + b * Z + ...
         // where X is least constrained.
 
-        let Some((solve_for, solve_for_coefficient)) = self
+        let (solve_for, solve_for_coefficient) = self
             .coefficients
             .iter()
             .filter(|(_var, coeff)| coeff.is_known_one() || coeff.is_known_minus_one())
@@ -269,13 +312,10 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
                     .get(var)
                     .map(|c| c.range_width())
                     .unwrap_or_else(|| T::modulus())
-            })
-        else {
-            return vec![];
-        };
+            })?;
 
         // This only works if the coefficients are all known.
-        let Some(summands) = self
+        let summands = self
             .coefficients
             .iter()
             .filter(|(var, _)| *var != solve_for)
@@ -285,19 +325,14 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
                 Some(rc.multiple(coeff))
             })
             .chain(std::iter::once(self.offset.range_constraint()))
-            .collect::<Option<Vec<_>>>()
-        else {
-            return vec![];
-        };
-        let Some(constraint) = summands.into_iter().reduce(|c1, c2| c1.combine_sum(&c2)) else {
-            return vec![];
-        };
+            .collect::<Option<Vec<_>>>()?;
+        let constraint = summands.into_iter().reduce(|c1, c2| c1.combine_sum(&c2))?;
         let constraint = if solve_for_coefficient.is_known_one() {
             -constraint
         } else {
             constraint
         };
-        vec![Effect::RangeConstraint(solve_for.clone(), constraint)]
+        Some(Effect::RangeConstraint(solve_for.clone(), constraint))
     }
 }
 
@@ -394,6 +429,8 @@ impl<T: FieldElement, V: Clone + Ord> Mul<&SymbolicExpression<T, V>>
 
 #[cfg(test)]
 mod test {
+    use pretty_assertions::assert_eq;
+
     use powdr_number::GoldilocksField;
 
     use super::*;
@@ -419,16 +456,17 @@ mod test {
         let x = &Ase::from_known_symbol("X", None);
         let y = &Ase::from_known_symbol("Y", None);
         let constr = x + y - from_number(10);
-        // We cannot solve it but also cannot know it is unsolvable.
-        assert!(constr.solve().unwrap().is_empty());
+        // We cannot solve it, but we can also not learn anything new from it.
+        let result = constr.solve().unwrap();
+        assert!(result.complete && result.effects.is_empty());
         // But if we know the values, we can be sure there is a conflict.
         assert!(from_number(10).solve().is_err());
     }
 
     #[test]
     fn solvable_without_vars() {
         let constr = &from_number(0);
-        assert!(constr.solve().unwrap().is_empty());
+        assert!(constr.solve().unwrap().effects.is_empty());
     }
 
     #[test]
@@ -440,9 +478,10 @@ mod test {
         let seven = from_number(7);
         let ten = from_number(10);
         let constr = mul(&two, &x) + mul(&seven, &y) - ten;
-        let effects = constr.solve().unwrap();
-        assert_eq!(effects.len(), 1);
-        let Effect::Assignment(var, expr) = &effects[0] else {
+        let result = constr.solve().unwrap();
+        assert!(result.complete);
+        assert_eq!(result.effects.len(), 1);
+        let Effect::Assignment(var, expr) = &result.effects[0] else {
             panic!("Expected assignment");
         };
         assert_eq!(var.to_string(), "X");
@@ -459,12 +498,14 @@ mod test {
         let ten = from_number(10);
         let constr = mul(&z, &x) + mul(&seven, &y) - ten.clone();
         // If we do not range-constrain z, we cannot solve since we don't know if it might be zero.
-        let effects = constr.solve().unwrap();
-        assert_eq!(effects.len(), 0);
+        let result = constr.solve().unwrap();
+        assert!(!result.complete && result.effects.is_empty());
         let z =
             Ase::from_known_symbol("z", Some(RangeConstraint::from_range(10.into(), 20.into())));
         let constr = mul(&z, &x) + mul(&seven, &y) - ten;
-        let effects = constr.solve().unwrap();
+        let result = constr.solve().unwrap();
+        assert!(result.complete);
+        let effects = result.effects;
         let Effect::Assignment(var, expr) = &effects[0] else {
             panic!("Expected assignment");
         };
@@ -488,17 +529,19 @@ mod test {
             + ten.clone()
             + z.clone();
         // Without range constraints, this is not solvable.
-        assert!(constr.solve().unwrap().is_empty());
+        let result = constr.solve().unwrap();
+        assert!(!result.complete && result.effects.is_empty());
         // Now add the range constraint on a, it should be solvable.
         let a = Ase::from_unknown_variable("a", rc.clone());
         let constr = mul(&a, &from_number(0x100))
             + mul(&b, &from_number(0x10000))
             + mul(&c, &from_number(0x1000000))
             + ten.clone()
             + z;
-        let effects = constr
-            .solve()
-            .unwrap()
+        let result = constr.solve().unwrap();
+        assert!(result.complete);
+        let effects = result
+            .effects
             .into_iter()
             .map(|effect| match effect {
                 Effect::Assignment(v, expr) => format!("{v} = {expr};\n"),
@@ -521,7 +564,7 @@ mod test {
             "a = ((-(10 + Z) & 65280) // 256);
 b = ((-(10 + Z) & 16711680) // 65536);
 c = ((-(10 + Z) & 4278190080) // 16777216);
-assert (10 + Z) == ((10 + Z) | 4294967040);
+assert -(10 + Z) == (-(10 + Z) | 4294967040);
 "
         );
     }
@@ -540,9 +583,10 @@ assert (10 + Z) == ((10 + Z) | 4294967040);
             + mul(&c, &from_number(0x1000000))
             + ten
             - z;
-        let effects = constr
-            .solve()
-            .unwrap()
+        let result = constr.solve().unwrap();
+        assert!(!result.complete);
+        let effects = result
+            .effects
             .into_iter()
             .map(|effect| match effect {
                 Effect::RangeConstraint(v, rc) => format!("{v}: {rc};\n"),

executor/src/witgen/jit/cell.rs

@@ -0,0 +1,59 @@
+use std::{
+    fmt::{self, Display, Formatter},
+    hash::{Hash, Hasher},
+};
+
+use powdr_ast::analyzed::AlgebraicReference;
+
+/// The identifier of a witness cell in the trace table.
+/// The `row_offset` is relative to a certain "zero row" defined
+/// by the component that uses this data structure.
+#[derive(Debug, Clone, Eq)]
+pub struct Cell {
+    /// Name of the column, used only for display purposes.
+    pub column_name: String,
+    pub id: u64,
+    pub row_offset: i32,
+}
+
+impl Hash for Cell {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.id.hash(state);
+        self.row_offset.hash(state);
+    }
+}
+
+impl PartialEq for Cell {
+    fn eq(&self, other: &Self) -> bool {
+        self.id == other.id && self.row_offset == other.row_offset
+    }
+}
+
+impl Ord for Cell {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        (self.id, self.row_offset).cmp(&(other.id, other.row_offset))
+    }
+}
+
+impl PartialOrd for Cell {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Cell {
+    pub fn from_reference(r: &AlgebraicReference, row_offset: i32) -> Self {
+        assert!(r.is_witness());
+        Self {
+            column_name: r.name.clone(),
+            id: r.poly_id.id,
+            row_offset: r.next as i32 + row_offset,
+        }
+    }
+}
+
+impl Display for Cell {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        write!(f, "{}[{}]", self.column_name, self.row_offset)
+    }
+}