Determine range constraints from fixed lookup.

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,30 @@ 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> {
+        let Some(summands) = self
+            .coefficients
+            .iter()
+            .map(|(var, coeff)| {
+                let coeff = coeff.try_to_number()?;
+                let rc = self.range_constraints.get(var)?;
+                Some(rc.multiple(coeff))
+            })
+            .chain(std::iter::once(Some(self.offset.range_constraint())))
+            .collect::<Option<Vec<_>>>()
+        else {
+            return Default::default();
+        };
+        summands
+            .into_iter()
+            .reduce(|c1, c2| c1.combine_sum(&c2))
+            // We always have at least the offset.
+            .unwrap()
+    }
+
     /// 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/witgen_inference.rs

@@ -242,43 +242,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 = evaluated
+            .iter()
+            .map(|e| e.as_ref().map(|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,
         }
     }
@@ -523,16 +529,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> {
@@ -541,7 +548,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,58 @@ 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 mut values = index
+            .iter()
+            .filter(|(key, _)| matches_range_constraint(key, &input_range_constraints))
+            .map(|(_, value)| {
+                let (_, values) = value.get()?;
+                Some(values)
+            });
+        let Some(first) = values.next() else {
+            // If the iterator is empty, we have no match, but it can always be answered,
+            // so this is successful.
+            // Unfortunately, we cannot express the "empty" range constraint.
+            return Some(vec![Default::default(); range_constraints.len()]);
+        };
+        // If an item (including the first) is None, it means the match was not unique.
+        let mut new_output_range_constraints = first?
+            .iter()
+            // min, max, mask
+            .map(|v| (*v, *v, v.to_integer()))
+            .collect_vec();
+        for v in values {
+            // If any value is None, the match was not unique.
+            let v = v?;
+            for ((min, max, mask), v) in new_output_range_constraints.iter_mut().zip_eq(v) {
+                *min = (*min).min(*v);
+                *max = (*max).max(*v);
+                *mask |= v.to_integer();
+            }
+        }
+        let mut new_output_range_constraints = new_output_range_constraints.into_iter();
+        Some(
+            known_arguments
+                .iter()
+                .map(|is_known| {
+                    if is_known {
+                        // We could also copy the input range constraint.
+                        RangeConstraint::default()
+                    } 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 +448,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)