Determine range constraints from fixed lookup. (#2300)

Co-authored-by: Georg Wiese <georgwiese@gmail.com>

Author: chriseth

executor/src/witgen/data_structures/mutable_state.rs

@@ -56,15 +56,12 @@ impl<'a, T: FieldElement, Q: QueryCallback<T>> MutableState<'a, T, Q> {
         identity_id: u64,
         known_inputs: &BitVec,
         range_constraints: &[RangeConstraint<T>],
-    ) -> bool {
+    ) -> Option<Vec<RangeConstraint<T>>> {
         // TODO We are currently ignoring bus interaction (also, but not only because there is no
         // unique machine responsible for handling a bus send), so just answer "false" if the identity
         // has no responsible machine.
-        self.responsible_machine(identity_id)
-            .ok()
-            .is_some_and(|mut machine| {
-                machine.can_process_call_fully(identity_id, known_inputs, range_constraints)
-            })
+        let mut machine = self.responsible_machine(identity_id).ok()?;
+        machine.can_process_call_fully(identity_id, known_inputs, range_constraints)
     }
 
     /// Call the machine responsible for the right-hand-side of an identity given its ID

executor/src/witgen/jit/affine_symbolic_expression.rs

@@ -113,6 +113,27 @@ impl<T: FieldElement, V: Ord + Clone + Display> AffineSymbolicExpression<T, V> {
         }
     }
 
+    /// Returns the range constraint of the whole expression.
+    /// This only works for simple expressions since all coefficients
+    /// must be known numbers.
+    pub fn range_constraint(&self) -> RangeConstraint<T> {
+        self.coefficients
+            .iter()
+            .map(|(var, coeff)| {
+                let coeff = coeff.try_to_number()?;
+                let rc = self.range_constraints.get(var)?;
+                Some(rc.multiple(coeff))
+            })
+            .collect::<Option<Vec<_>>>()
+            .and_then(|summands| {
+                summands
+                    .into_iter()
+                    .chain(std::iter::once(self.offset.range_constraint()))
+                    .reduce(|c1, c2| c1.combine_sum(&c2))
+            })
+            .unwrap_or_default()
+    }
+
     /// If this expression contains a single unknown variable, returns it.
     pub fn single_unknown_variable(&self) -> Option<&V> {
         if self.coefficients.len() == 1 {

executor/src/witgen/jit/single_step_processor.rs

@@ -205,4 +205,40 @@ if (VM::instr_add[0] == 1) {
 }"
         );
     }
+
+    #[test]
+    fn range_constraints_from_lookup() {
+        let input = "
+    namespace VM(256);
+        let instr_add: col;
+        let instr_mul: col;
+        let pc: col;
+
+        col fixed LINE = [0, 1] + [2]*;
+        col fixed INSTR_ADD = [0, 1] + [0]*;
+        col fixed INSTR_MUL = [1, 0] + [1]*;
+
+        pc' = pc + 1;
+        instr_add = 0;
+        [ pc, instr_add, instr_mul ] in [ LINE, INSTR_ADD, INSTR_MUL ];
+
+        ";
+        let code = generate_single_step(input, "Main").unwrap();
+        // After the machine call, we should have a direct assignment `VM::instr_mul[1] = 1`,
+        // instead of just an assignment from the call variable.
+        // This is because the fixed lookup machine can already provide a range constraint.
+        // For reasons of processing order, the call variable will also be assigned
+        // right before the call.
+        assert_eq!(
+            format_code(&code),
+            "\
+VM::pc[1] = (VM::pc[0] + 1);
+VM::instr_add[1] = 0;
+call_var(2, 1, 0) = VM::pc[1];
+call_var(2, 1, 1) = 0;
+call_var(2, 1, 2) = 1;
+machine_call(2, [Known(call_var(2, 1, 0)), Known(call_var(2, 1, 1)), Unknown(call_var(2, 1, 2))]);
+VM::instr_mul[1] = 1;"
+        );
+    }
 }

executor/src/witgen/jit/witgen_inference.rs

@@ -233,43 +233,49 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
         }
         let evaluated = arguments
             .iter()
-            .map(|a| {
-                self.evaluate(a, row_offset)
-                    .and_then(|a| a.try_to_known().cloned())
-            })
+            .map(|a| self.evaluate(a, row_offset))
             .collect::<Vec<_>>();
         let range_constraints = evaluated
             .iter()
             .map(|e| e.as_ref().map(|e| e.range_constraint()).unwrap_or_default())
             .collect_vec();
-        let known = evaluated.iter().map(|e| e.is_some()).collect();
+        let known: BitVec = evaluated
+            .iter()
+            .map(|e| e.as_ref().and_then(|e| e.try_to_known()).is_some())
+            .collect();
 
-        if !can_process_call.can_process_call_fully(lookup_id, &known, &range_constraints) {
+        let Some(new_range_constraints) =
+            can_process_call.can_process_call_fully(lookup_id, &known, &range_constraints)
+        else {
             log::trace!(
                 "Sub-machine cannot process call fully (will retry later): {lookup_id}, arguments: {}",
                 arguments.iter().zip(known).map(|(arg, known)| {
                     format!("{arg} [{}]", if known { "known" } else { "unknown" })
                 }).format(", "));
             return ProcessResult::empty();
-        }
+        };
+        let mut effects = vec![];
         let vars = arguments
             .iter()
+            .zip_eq(new_range_constraints)
             .enumerate()
-            .map(|(index, arg)| {
+            .map(|(index, (arg, new_rc))| {
                 let var = Variable::MachineCallParam(MachineCallVariable {
                     identity_id: lookup_id,
                     row_offset,
                     index,
                 });
                 self.assign_variable(arg, row_offset, var.clone());
+                effects.push(Effect::RangeConstraint(var.clone(), new_rc.clone()));
                 if known[index] {
                     assert!(self.is_known(&var));
                 }
                 var
             })
             .collect_vec();
+        effects.push(Effect::MachineCall(lookup_id, known, vars.clone()));
         ProcessResult {
-            effects: vec![Effect::MachineCall(lookup_id, known, vars)],
+            effects,
             complete: true,
         }
     }
@@ -514,16 +520,17 @@ pub trait FixedEvaluator<T: FieldElement>: Clone {
 }
 
 pub trait CanProcessCall<T: FieldElement> {
-    /// Returns true if a call to the machine that handles the given identity
+    /// Returns Some(..) if a call to the machine that handles the given identity
     /// can always be processed with the given known inputs and range constraints
     /// on the parameters.
+    /// The value in the Option is a vector of new range constraints.
     /// @see Machine::can_process_call
     fn can_process_call_fully(
         &self,
         _identity_id: u64,
         _known_inputs: &BitVec,
         _range_constraints: &[RangeConstraint<T>],
-    ) -> bool;
+    ) -> Option<Vec<RangeConstraint<T>>>;
 }
 
 impl<T: FieldElement, Q: QueryCallback<T>> CanProcessCall<T> for &MutableState<'_, T, Q> {
@@ -532,7 +539,7 @@ impl<T: FieldElement, Q: QueryCallback<T>> CanProcessCall<T> for &MutableState<'
         identity_id: u64,
         known_inputs: &BitVec,
         range_constraints: &[RangeConstraint<T>],
-    ) -> bool {
+    ) -> Option<Vec<RangeConstraint<T>>> {
         MutableState::can_process_call_fully(self, identity_id, known_inputs, range_constraints)
     }
 }

executor/src/witgen/machines/double_sorted_witness_machine_32.rs

@@ -191,7 +191,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses32<'a, T> {
         identity_id: u64,
         known_arguments: &BitVec,
         range_constraints: &[RangeConstraint<T>],
-    ) -> bool {
+    ) -> Option<Vec<RangeConstraint<T>>> {
         assert!(self.parts.connections.contains_key(&identity_id));
         assert_eq!(known_arguments.len(), 4);
         assert_eq!(range_constraints.len(), 4);
@@ -200,17 +200,18 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses32<'a, T> {
 
         // We need to known operation_id, step and address for all calls.
         if !known_arguments[0] || !known_arguments[1] || !known_arguments[2] {
-            return false;
+            return None;
         }
 
         // For the value, it depends: If we write, we need to know it, if we read we do not need to know it.
         if known_arguments[3] {
             // It is known, so we are good anyway.
-            true
+            Some(vec![RangeConstraint::unconstrained(); 4])
         } else {
             // It is not known, so we can only process if we do not write.
-            !range_constraints[0].allows_value(T::from(OPERATION_ID_BOOTLOADER_WRITE))
-                && !range_constraints[0].allows_value(T::from(OPERATION_ID_WRITE))
+            (!range_constraints[0].allows_value(T::from(OPERATION_ID_BOOTLOADER_WRITE))
+                && !range_constraints[0].allows_value(T::from(OPERATION_ID_WRITE)))
+            .then(|| vec![RangeConstraint::unconstrained(); 4])
         }
     }
 

executor/src/witgen/machines/fixed_lookup_machine.rs

@@ -260,10 +260,10 @@ impl<'a, T: FieldElement> Machine<'a, T> for FixedLookup<'a, T> {
         &mut self,
         identity_id: u64,
         known_arguments: &BitVec,
-        _range_constraints: &[RangeConstraint<T>],
-    ) -> bool {
+        range_constraints: &[RangeConstraint<T>],
+    ) -> Option<Vec<RangeConstraint<T>>> {
         if !Self::is_responsible(&self.connections[&identity_id]) {
-            return false;
+            return None;
         }
         let index = self
             .indices
@@ -274,8 +274,72 @@ impl<'a, T: FieldElement> Machine<'a, T> for FixedLookup<'a, T> {
             .or_insert_with_key(|application| {
                 create_index(self.fixed_data, application, &self.connections)
             });
-        // Check that if there is a match, it is unique.
-        index.values().all(|value| value.0.is_some())
+        let input_range_constraints = known_arguments
+            .iter()
+            .zip_eq(range_constraints)
+            .filter_map(|(is_known, range_constraint)| is_known.then_some(range_constraint.clone()))
+            .collect_vec();
+
+        // Now only consider the index entries that match the input range constraints,
+        // see that the result is unique and determine new output range constraints.
+        let values_matching_input_constraints = index
+            .iter()
+            .filter(|(key, _)| matches_range_constraint(key, &input_range_constraints))
+            .map(|(_, value)| {
+                let (_, values) = value.get()?;
+                Some(values)
+            });
+        let mut new_range_constraints: Option<Vec<(T, T, T::Integer)>> = None;
+        for values in values_matching_input_constraints {
+            // If any value is None, it means the lookup does not have a unique answer,
+            // and thus we cannot process the call.
+            let values = values?;
+            new_range_constraints = Some(match new_range_constraints {
+                // First item, turn each value into (min, max, mask).
+                None => values
+                    .iter()
+                    .map(|v| (*v, *v, v.to_integer()))
+                    .collect_vec(),
+                // Reduce range constraint by disjunction.
+                Some(mut acc) => {
+                    for ((min, max, mask), v) in acc.iter_mut().zip_eq(values) {
+                        *min = (*min).min(*v);
+                        *max = (*max).max(*v);
+                        *mask |= v.to_integer();
+                    }
+                    acc
+                }
+            })
+        }
+        Some(match new_range_constraints {
+            None => {
+                // The iterator was empty, i.e. there are no inputs in the index matching the
+                // range constraints.
+                // This means that every call like this will lead to a fatal error, but there is
+                // enough information in the inputs to hanlde the call. Unfortunately, there is
+                // no way to signal this in the return type, yet.
+                // TODO(#2324): change this.
+                // We just return the input range constraints to signal "everything allright".
+                range_constraints.to_vec()
+            }
+            Some(new_output_range_constraints) => {
+                let mut new_output_range_constraints = new_output_range_constraints.into_iter();
+                known_arguments
+                    .iter()
+                    .enumerate()
+                    .map(|(i, is_known)| {
+                        if is_known {
+                            // Just copy the input range constraint.
+                            range_constraints[i].clone()
+                        } else {
+                            let (min, max, mask) = new_output_range_constraints.next().unwrap();
+                            RangeConstraint::from_range(min, max)
+                                .conjunction(&RangeConstraint::from_mask(mask))
+                        }
+                    })
+                    .collect()
+            }
+        })
     }
 
     fn process_plookup<Q: crate::witgen::QueryCallback<T>>(
@@ -398,3 +462,14 @@ impl<'a, T: FieldElement> RangeConstraintSet<AlgebraicVariable<'a>, T>
         }
     }
 }
+
+/// Checks that an array of values satisfies a set of range constraints.
+fn matches_range_constraint<T: FieldElement>(
+    values: &[T],
+    range_constraints: &[RangeConstraint<T>],
+) -> bool {
+    values
+        .iter()
+        .zip_eq(range_constraints)
+        .all(|(value, range_constraint)| range_constraint.allows_value(*value))
+}

executor/src/witgen/machines/mod.rs

@@ -55,21 +55,23 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync {
         );
     }
 
-    /// Returns true if this machine can alway fully process a call via the given
+    /// Returns Some(..) if this machine can alway fully process a call via the given
     /// identity, the set of known arguments and a list of range constraints
     /// on the parameters. Note that the range constraints can be imposed both
     /// on inputs and on outputs.
-    /// If this returns true, then corresponding calls to `process_lookup_direct`
+    /// If this returns Some(..), then corresponding calls to `process_lookup_direct`
     /// are safe.
+    /// The value returned inside the option is a vector of range constraints on the arguments,
+    /// which again can be imposed both on inputs and on outputs.
     /// The function requires `&mut self` because it usually builds an index structure
     /// or something similar.
     fn can_process_call_fully(
         &mut self,
         _identity_id: u64,
         _known_arguments: &BitVec,
         _range_constraints: &[RangeConstraint<T>],
-    ) -> bool {
-        false
+    ) -> Option<Vec<RangeConstraint<T>>> {
+        None
     }
 
     /// Like `process_plookup`, but also records the time spent in this machine.
@@ -188,7 +190,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
         identity_id: u64,
         known_arguments: &BitVec,
         range_constraints: &[RangeConstraint<T>],
-    ) -> bool {
+    ) -> Option<Vec<RangeConstraint<T>>> {
         match self {
             KnownMachine::SecondStageMachine(m) => {
                 m.can_process_call_fully(identity_id, known_arguments, range_constraints)

pipeline/tests/asm.rs

@@ -467,7 +467,7 @@ fn permutation_to_block() {
 }
 
 #[test]
-#[should_panic = "called `Result::unwrap()` on an `Err` value: Linear constraint is not satisfiable: 18446744069414584320 != 0"]
+#[should_panic = "Column main_bin::pc is not stackable in a 1-row block, conflict in rows 0 and 1"]
 fn permutation_to_vm() {
     // TODO: witgen issue: Machine incorrectly detected as block machine.
     let f = "asm/permutations/vm_to_vm.asm";