Introduce direct lookup. (#2063)

This PR introduces a more "direct" way to perform a lookup during
witness generation. It removes the concept of `EvalValue`, which is a
list of "updates", and instead requests the called machine to directly
fill in mutable pointers to field elements.

The goal is to use this (hopefully) faster interface if
 - the lookup can be fully solved in a single call
 - no cell-based range constraints are needed

If the LHS of the lookup consists of direct polynomial references (or
next references), the caller can pass a pointer to the final table and
does not need to move data around any further.

Some numbers:

For the "keccak test", and only looking at the PC lookup, we get:

Inside `process_plookup_internal`:

40 ns: preparing the `data` and `values` arrays
290 ns: call to process_lookup_direct
1300 ns: computing the result EvalValue.

executor/src/witgen/eval_result.rs

@@ -161,6 +161,7 @@ impl<K, T: FieldElement> EvalValue<K, T> {
 /// New assignments or constraints for witness columns identified by an ID.
 pub type EvalResult<'a, T, K = AlgebraicVariable<'a>> = Result<EvalValue<K, T>, EvalError<T>>;
 
+/// A fatal error for witness generation.
 #[derive(Clone, PartialEq)]
 pub enum EvalError<T: FieldElement> {
     /// We ran out of rows

executor/src/witgen/machines/fixed_lookup_machine.rs

@@ -18,7 +18,7 @@ use crate::witgen::util::try_to_simple_poly_ref;
 use crate::witgen::{EvalError, EvalValue, IncompleteCause, MutableState, QueryCallback};
 use crate::witgen::{EvalResult, FixedData};
 
-use super::{Connection, Machine};
+use super::{Connection, LookupCell, Machine};
 
 /// An Application specifies a lookup cache.
 #[derive(Hash, Eq, PartialEq, Ord, PartialOrd, Clone)]
@@ -197,8 +197,9 @@ impl<'a, T: FieldElement> FixedLookup<'a, T> {
         }
     }
 
-    fn process_plookup_internal(
+    fn process_plookup_internal<'b, Q: QueryCallback<T>>(
         &mut self,
+        mutable_state: &'b mut MutableState<'a, 'b, T, Q>,
         identity_id: u64,
         rows: &RowPair<'_, '_, T>,
         left: &[AffineExpression<AlgebraicVariable<'a>, T>],
@@ -217,62 +218,42 @@ impl<'a, T: FieldElement> FixedLookup<'a, T> {
         }
 
         // Split the left-hand-side into known input values and unknown output expressions.
-        let mut input_values = vec![];
-        let mut known_inputs: BitVec = Default::default();
-        let mut output_expressions = vec![];
-
-        for l in left {
-            if let Some(value) = l.constant_value() {
-                input_values.push(value);
-                known_inputs.push(true);
-            } else {
-                output_expressions.push(l);
-                known_inputs.push(false);
-            }
-        }
-
-        let application = Application {
-            identity_id,
-            inputs: known_inputs,
-        };
-
-        let index = self
-            .indices
-            .entry(application)
-            .or_insert_with_key(|application| {
-                create_index(self.fixed_data, application, &self.connections)
-            });
-        let index_value = index.get(&input_values).ok_or_else(|| {
-            let input_assignment = left
-                .iter()
-                .zip(right)
-                .filter_map(|(l, r)| l.constant_value().map(|v| (r.name.clone(), v)))
-                .collect();
-            EvalError::FixedLookupFailed(input_assignment)
-        })?;
+        let mut data = vec![T::zero(); left.len()];
+        let values = left
+            .iter()
+            .zip(&mut data)
+            .map(|(l, d)| {
+                if let Some(value) = l.constant_value() {
+                    *d = value;
+                    LookupCell::Input(d)
+                } else {
+                    LookupCell::Output(d)
+                }
+            })
+            .collect::<Vec<_>>();
 
-        let Some((row, output)) = index_value.get() else {
+        if !self.process_lookup_direct(mutable_state, identity_id, values)? {
             // multiple matches, we stop and learnt nothing
             return Ok(EvalValue::incomplete(
                 IncompleteCause::MultipleLookupMatches,
             ));
         };
 
-        self.multiplicity_counter.increment_at_row(identity_id, row);
-
         let mut result = EvalValue::complete(vec![]);
-        for (l, r) in output_expressions.into_iter().zip(output) {
-            let evaluated = l.clone() - (*r).into();
-            // TODO we could use bit constraints here
-            match evaluated.solve() {
-                Ok(constraints) => {
-                    result.combine(constraints);
-                }
-                Err(_) => {
-                    // Fail the whole lookup
-                    return Err(EvalError::ConstraintUnsatisfiable(format!(
-                        "Constraint is invalid ({l} != {r}).",
-                    )));
+        for (l, v) in left.iter().zip(data) {
+            if !l.is_constant() {
+                let evaluated = l.clone() - v.into();
+                // TODO we could use bit constraints here
+                match evaluated.solve() {
+                    Ok(constraints) => {
+                        result.combine(constraints);
+                    }
+                    Err(_) => {
+                        // Fail the whole lookup
+                        return Err(EvalError::ConstraintUnsatisfiable(format!(
+                            "Constraint is invalid ({l} != {v}).",
+                        )));
+                    }
                 }
             }
         }
@@ -347,7 +328,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for FixedLookup<'a, T> {
 
     fn process_plookup<'b, Q: crate::witgen::QueryCallback<T>>(
         &mut self,
-        _mutable_state: &'b mut crate::witgen::MutableState<'a, 'b, T, Q>,
+        mutable_state: &'b mut crate::witgen::MutableState<'a, 'b, T, Q>,
         identity_id: u64,
         caller_rows: &'b RowPair<'b, 'a, T>,
     ) -> EvalResult<'a, T> {
@@ -362,7 +343,79 @@ impl<'a, T: FieldElement> Machine<'a, T> for FixedLookup<'a, T> {
             .peekable();
 
         let outer_query = OuterQuery::new(caller_rows, identity);
-        self.process_plookup_internal(identity_id, caller_rows, &outer_query.left, right)
+        self.process_plookup_internal(
+            mutable_state,
+            identity_id,
+            caller_rows,
+            &outer_query.left,
+            right,
+        )
+    }
+
+    fn process_lookup_direct<'b, 'c, Q: QueryCallback<T>>(
+        &mut self,
+        _mutable_state: &'b mut MutableState<'a, 'b, T, Q>,
+        identity_id: u64,
+        values: Vec<LookupCell<'c, T>>,
+    ) -> Result<bool, EvalError<T>> {
+        let mut input_values = vec![];
+
+        let known_inputs = values
+            .iter()
+            .map(|v| match v {
+                LookupCell::Input(value) => {
+                    input_values.push(**value);
+                    true
+                }
+                LookupCell::Output(_) => false,
+            })
+            .collect();
+
+        let application = Application {
+            identity_id,
+            inputs: known_inputs,
+        };
+
+        let index = self
+            .indices
+            .entry(application)
+            .or_insert_with_key(|application| {
+                create_index(self.fixed_data, application, &self.connections)
+            });
+        let index_value = index.get(&input_values).ok_or_else(|| {
+            let right = self.connections[&identity_id].right;
+            let input_assignment = values
+                .iter()
+                .zip(&right.expressions)
+                .filter_map(|(l, r)| match l {
+                    LookupCell::Input(v) => {
+                        let name = try_to_simple_poly_ref(r).unwrap().name.clone();
+                        Some((name, **v))
+                    }
+                    _ => None,
+                })
+                .collect();
+            EvalError::FixedLookupFailed(input_assignment)
+        })?;
+
+        let Some((row, output)) = index_value.get() else {
+            // multiple matches, we stop and learnt nothing
+            return Ok(false);
+        };
+
+        self.multiplicity_counter.increment_at_row(identity_id, row);
+
+        values
+            .into_iter()
+            .filter_map(|v| match v {
+                LookupCell::Output(e) => Some(e),
+                _ => None,
+            })
+            .zip(output)
+            .for_each(|(e, v)| {
+                *e = *v;
+            });
+        Ok(true)
     }
 
     fn take_witness_col_values<'b, Q: QueryCallback<T>>(

executor/src/witgen/machines/mod.rs

@@ -20,7 +20,7 @@ use self::write_once_memory::WriteOnceMemory;
 
 use super::generator::Generator;
 use super::rows::RowPair;
-use super::{EvalResult, FixedData, MutableState, QueryCallback};
+use super::{EvalError, EvalResult, FixedData, MutableState, QueryCallback};
 
 mod block_machine;
 mod double_sorted_witness_machine_16;
@@ -60,6 +60,25 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync {
         caller_rows: &'b RowPair<'b, 'a, T>,
     ) -> EvalResult<'a, T>;
 
+    /// Process a connection of a given ID (which must be known to the callee).
+    /// This is a more direct version of `process_plookup`, where the caller
+    /// provides values or targets to where to write the results directly.
+    /// The length of `values` needs to be the same as the number of expressions
+    /// in the LHS / RHS of the connection.
+    /// It does not allow to return range constraints or complex expressions.
+    /// The boolean return value indicates whether the lookup was successful.
+    /// If it returns true, all output values in `values` need to have been set.
+    /// If it returns false, none of them should be changed.
+    /// An error is always unrecoverable.
+    fn process_lookup_direct<'b, 'c, Q: QueryCallback<T>>(
+        &mut self,
+        _mutable_state: &'b mut MutableState<'a, 'b, T, Q>,
+        _identity_id: u64,
+        _values: Vec<LookupCell<'c, T>>,
+    ) -> Result<bool, EvalError<T>> {
+        unimplemented!("Direct lookup is not supported for this machine.");
+    }
+
     /// Returns the final values of the witness columns.
     fn take_witness_col_values<'b, Q: QueryCallback<T>>(
         &mut self,
@@ -70,6 +89,13 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync {
     fn identity_ids(&self) -> Vec<u64>;
 }
 
+pub enum LookupCell<'a, T> {
+    /// Value is known (i.e. an input)
+    Input(&'a T),
+    /// Value is not known (i.e. an output)
+    Output(&'a mut T),
+}
+
 /// All known implementations of [Machine].
 /// This allows us to treat machines uniformly without putting them into a `Box`,
 /// which requires that all lifetime parameters are 'static.