Skip to content

Commit

Permalink
Range constraints from failing lookups. (#2444)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
chriseth authored Feb 11, 2025
1 parent 40318d5 commit d4f9ca6
Show file tree
Hide file tree
Showing 8 changed files with 187 additions and 133 deletions.
6 changes: 3 additions & 3 deletions executor/src/witgen/data_structures/mutable_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand Down
38 changes: 38 additions & 0 deletions executor/src/witgen/jit/block_machine_processor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]);"
);
}
}
36 changes: 16 additions & 20 deletions executor/src/witgen/jit/witgen_inference.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}

Expand Down Expand Up @@ -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)
}
}
Expand Down
16 changes: 11 additions & 5 deletions executor/src/witgen/machines/block_machine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>>(
Expand Down
18 changes: 9 additions & 9 deletions executor/src/witgen/machines/double_sorted_witness_machine_32.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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>>(
Expand Down
Loading

0 comments on commit d4f9ca6

Please sign in to comment.