Range constraints from failing lookups. (#2444)

Identities of the kind
`[ b * clock_0 + c * clock_1 ] in [ byte ];`, which are used in the
arithmetic machine, are currently not process by auto-witjitgen. The
idea is that this results on range constraints for different columns on
different rows.

The mechanism in place previously would only return range constraints in
case a lookup can always be processed, but this is not the case here: We
will never be able to determine a value just from the lookup, the only
information we can get here are range constraints.

TODO:
- [x] it might be that we will never be able to process the lookup (i.e.
emit a call), in that case we should still add the assertion, because
that is what is supposed to be checked. Maybe we should also detect that
it is just a range constraint after all?
Oh I think we will need to make the call in the end, once we know the
value by other means.

Author: chriseth

executor/src/witgen/data_structures/mutable_state.rs

@@ -55,9 +55,9 @@ impl<'a, T: FieldElement, Q: QueryCallback<T>> MutableState<'a, T, Q> {
         &self,
         identity_id: u64,
         known_inputs: &BitVec,
-        range_constraints: &[RangeConstraint<T>],
-    ) -> Option<Vec<RangeConstraint<T>>> {
-        let mut machine = self.responsible_machine(identity_id).unwrap();
+        range_constraints: Vec<RangeConstraint<T>>,
+    ) -> (bool, Vec<RangeConstraint<T>>) {
+        let mut machine = self.responsible_machine(identity_id).ok().unwrap();
         machine.can_process_call_fully(self, identity_id, known_inputs, range_constraints)
     }
 

executor/src/witgen/jit/block_machine_processor.rs

@@ -426,4 +426,42 @@ params[3] = main_binary::C[3];"
 params[1] = Sub::b[0];"
         );
     }
+
+    #[test]
+    fn complex_fixed_lookup_range_constraint() {
+        let input = "
+        namespace main(256);
+            col witness a, b, c;
+            [a, b, c] is SubM.sel $ [SubM.a, SubM.b, SubM.c]; 
+        namespace SubM(256);
+            col witness a, b, c;
+            let sel: col = |i| (i + 1) % 2;
+            let clock_0: col = |i| i % 2;
+            let clock_1: col = |i| (i + 1) % 2;
+            let byte: col = |i| i & 0xff;
+            [ b * clock_0 + c * clock_1 ] in [ byte ];
+            (b' - b) * sel = 0;
+            (c' - c) * sel = 0;
+            a = b * 256 + c;
+        ";
+        let code = generate_for_block_machine(input, "SubM", 1, 2)
+            .unwrap()
+            .code;
+        assert_eq!(
+            format_code(&code),
+            "SubM::a[0] = params[0];
+SubM::b[0] = ((SubM::a[0] & 0xff00) // 256);
+SubM::c[0] = (SubM::a[0] & 0xff);
+assert (SubM::a[0] & 0xffffffffffff0000) == 0;
+params[1] = SubM::b[0];
+params[2] = SubM::c[0];
+call_var(1, 0, 0) = SubM::c[0];
+machine_call(1, [Known(call_var(1, 0, 0))]);
+SubM::b[1] = SubM::b[0];
+call_var(1, 1, 0) = SubM::b[1];
+SubM::c[1] = SubM::c[0];
+machine_call(1, [Known(call_var(1, 1, 0))]);
+SubM::a[1] = ((SubM::b[1] * 256) + SubM::c[1]);"
+        );
+    }
 }

executor/src/witgen/jit/witgen_inference.rs

@@ -333,24 +333,20 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> WitgenInference<'a, T, F
             .collect_vec();
         let known: BitVec = arguments.iter().map(|v| self.is_known(v)).collect();
 
-        let Some(new_range_constraints) =
-            can_process_call.can_process_call_fully(lookup_id, &known, &range_constraints)
-        else {
-            return ProcessResult::empty();
-        };
-        let effects = arguments
+        let (can_process, range_constraints) =
+            can_process_call.can_process_call_fully(lookup_id, &known, range_constraints);
+
+        let mut effects = arguments
             .iter()
-            .zip_eq(new_range_constraints)
+            .zip_eq(range_constraints)
             .map(|(var, new_rc)| Effect::RangeConstraint(var.clone(), new_rc.clone()))
-            .chain(std::iter::once(Effect::MachineCall(
-                lookup_id,
-                known,
-                arguments.to_vec(),
-            )))
-            .collect();
+            .collect_vec();
+        if can_process {
+            effects.push(Effect::MachineCall(lookup_id, known, arguments.to_vec()));
+        }
         ProcessResult {
             effects,
-            complete: true,
+            complete: can_process,
         }
     }
 
@@ -619,26 +615,26 @@ pub trait FixedEvaluator<T: FieldElement>: Clone {
 }
 
 pub trait CanProcessCall<T: FieldElement>: Clone {
-    /// Returns Some(..) if a call to the machine that handles the given identity
+    /// Returns (true, _) 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.
+    /// The second return value 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>],
-    ) -> Option<Vec<RangeConstraint<T>>>;
+        _range_constraints: Vec<RangeConstraint<T>>,
+    ) -> (bool, Vec<RangeConstraint<T>>);
 }
 
 impl<T: FieldElement, Q: QueryCallback<T>> CanProcessCall<T> for &MutableState<'_, T, Q> {
     fn can_process_call_fully(
         &self,
         identity_id: u64,
         known_inputs: &BitVec,
-        range_constraints: &[RangeConstraint<T>],
-    ) -> Option<Vec<RangeConstraint<T>>> {
+        range_constraints: Vec<RangeConstraint<T>>,
+    ) -> (bool, Vec<RangeConstraint<T>>) {
         MutableState::can_process_call_fully(self, identity_id, known_inputs, range_constraints)
     }
 }

executor/src/witgen/machines/block_machine.rs

@@ -170,18 +170,24 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
         can_process: impl CanProcessCall<T>,
         identity_id: u64,
         known_arguments: &BitVec,
-        range_constraints: &[RangeConstraint<T>],
-    ) -> Option<Vec<RangeConstraint<T>>> {
+        range_constraints: Vec<RangeConstraint<T>>,
+    ) -> (bool, Vec<RangeConstraint<T>>) {
         // We use the input range constraints to see if there is a column
         // containing the substring "operation_id" which is constrained to a
         // single value and use that value as part of the cache key.
         let operation_id = self.find_operation_id(identity_id).and_then(|index| {
             let v = range_constraints[index].try_to_single_value()?;
             Some((index, v))
         });
-        self.function_cache
-            .compile_cached(can_process, identity_id, known_arguments, operation_id)
-            .map(|r| r.range_constraints.clone())
+        match self.function_cache.compile_cached(
+            can_process,
+            identity_id,
+            known_arguments,
+            operation_id,
+        ) {
+            Some(entry) => (true, entry.range_constraints.clone()),
+            None => (false, range_constraints),
+        }
     }
 
     fn process_lookup_direct<'b, 'c, Q: QueryCallback<T>>(

executor/src/witgen/machines/double_sorted_witness_machine_32.rs

@@ -197,8 +197,8 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses32<'a, T> {
         _can_process: impl CanProcessCall<T>,
         identity_id: u64,
         known_arguments: &BitVec,
-        range_constraints: &[RangeConstraint<T>],
-    ) -> Option<Vec<RangeConstraint<T>>> {
+        range_constraints: Vec<RangeConstraint<T>>,
+    ) -> (bool, Vec<RangeConstraint<T>>) {
         assert!(self.parts.connections.contains_key(&identity_id));
         assert_eq!(known_arguments.len(), 4);
         assert_eq!(range_constraints.len(), 4);
@@ -207,19 +207,19 @@ 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 None;
+            return (false, range_constraints);
         }
 
         // 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] {
+        let can_answer = if known_arguments[3] {
             // It is known, so we are good anyway.
-            Some(vec![RangeConstraint::unconstrained(); 4])
+            true
         } 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)))
-            .then(|| vec![RangeConstraint::unconstrained(); 4])
-        }
+            !range_constraints[0].allows_value(T::from(OPERATION_ID_BOOTLOADER_WRITE))
+                && !range_constraints[0].allows_value(T::from(OPERATION_ID_WRITE))
+        };
+        (can_answer, range_constraints)
     }
 
     fn process_lookup_direct<'b, 'c, Q: QueryCallback<T>>(