Refactorings around process_lookup_direct (#2209)

This PR refactors a few things:
- `process_lookup_direct` no longer has a default implementation.
Eventually, we want all machines to implement it, so I figured it would
be better to explicitly panic in each machine.
- Refactored the implementation of
`FixedLookupMachine::process_plookup`, pulling some stuff out into a new
`CallerData` struct. This is similar to what @chriseth has done on
[`call_jit_from_block`](main...call_jit_from_block),
see the comment below.
- As a first test, I implemented `process_lookup_direct` for the
"large"-field memory machine (and `process_plookup` by wrapping
`process_lookup_direct`)

executor/src/witgen/data_structures/caller_data.rs

@@ -0,0 +1,73 @@
+use itertools::Itertools;
+use powdr_number::FieldElement;
+
+use crate::witgen::{
+    machines::LookupCell,
+    processor::{Left, OuterQuery},
+    EvalError, EvalResult, EvalValue,
+};
+
+/// A representation of the caller's data.
+///
+/// Useful for implementing [Machine::process_plookup] in terms of [Machine::process_lookup_direct].
+pub struct CallerData<'a, 'b, T> {
+    /// The raw data of the caller. Unknown values should be ignored.
+    data: Vec<T>,
+    /// The affine expressions of the caller.
+    left: &'b Left<'a, T>,
+}
+
+impl<'a, 'b, T: FieldElement> From<&'b OuterQuery<'a, '_, T>> for CallerData<'a, 'b, T> {
+    /// Builds a `CallerData` from an `OuterQuery`.
+    fn from(outer_query: &'b OuterQuery<'a, '_, T>) -> Self {
+        let data = outer_query
+            .left
+            .iter()
+            .map(|l| l.constant_value().unwrap_or_default())
+            .collect();
+        Self {
+            data,
+            left: &outer_query.left,
+        }
+    }
+}
+
+impl<'a, 'b, T: FieldElement> CallerData<'a, 'b, T> {
+    /// Returns the data as a list of `LookupCell`s, as expected by `Machine::process_lookup_direct`.
+    pub fn as_lookup_cells(&mut self) -> Vec<LookupCell<'_, T>> {
+        self.data
+            .iter_mut()
+            .zip_eq(self.left.iter())
+            .map(|(value, left)| match left.constant_value().is_some() {
+                true => LookupCell::Input(value),
+                false => LookupCell::Output(value),
+            })
+            .collect()
+    }
+}
+
+impl<'a, 'b, T: FieldElement> From<CallerData<'a, 'b, T>> for EvalResult<'a, T> {
+    /// Turns the result of a direct lookup into an `EvalResult`, as used by `Machine::process_plookup`.
+    ///
+    /// Note that this function assumes that the lookup was successful and complete.
+    fn from(data: CallerData<'a, 'b, T>) -> EvalResult<'a, T> {
+        let mut result = EvalValue::complete(vec![]);
+        for (l, v) in data.left.iter().zip_eq(data.data.iter()) {
+            if !l.is_constant() {
+                let evaluated = l.clone() - (*v).into();
+                match evaluated.solve() {
+                    Ok(constraints) => {
+                        result.combine(constraints);
+                    }
+                    Err(_) => {
+                        // Fail the whole lookup
+                        return Err(EvalError::ConstraintUnsatisfiable(format!(
+                            "Constraint is invalid ({l} != {v}).",
+                        )));
+                    }
+                }
+            }
+        }
+        Ok(result)
+    }
+}

executor/src/witgen/data_structures/mod.rs

@@ -1,3 +1,4 @@
+pub mod caller_data;
 pub mod column_map;
 pub mod copy_constraints;
 pub mod finalizable_data;

executor/src/witgen/machines/block_machine.rs

@@ -2,7 +2,9 @@ use std::collections::{BTreeMap, HashMap};
 use std::fmt::Display;
 use std::iter::{self};
 
-use super::{compute_size_and_log, ConnectionKind, EvalResult, FixedData, MachineParts};
+use super::{
+    compute_size_and_log, ConnectionKind, EvalResult, FixedData, LookupCell, MachineParts,
+};
 
 use crate::witgen::affine_expression::AlgebraicVariable;
 use crate::witgen::analysis::detect_connection_type_and_block_size;
@@ -139,6 +141,15 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
         self.parts.connections.keys().copied().collect()
     }
 
+    fn process_lookup_direct<'b, 'c, Q: QueryCallback<T>>(
+        &mut self,
+        _mutable_state: &'b MutableState<'a, T, Q>,
+        _identity_id: u64,
+        _values: &mut [LookupCell<'c, T>],
+    ) -> Result<bool, EvalError<T>> {
+        unimplemented!("Direct lookup not supported by machine {}.", self.name())
+    }
+
     fn process_plookup<'b, Q: QueryCallback<T>>(
         &mut self,
         mutable_state: &'b MutableState<'a, T, Q>,

executor/src/witgen/machines/double_sorted_witness_machine_16.rs

@@ -3,7 +3,7 @@ use std::iter::once;
 
 use itertools::Itertools;
 
-use super::{ConnectionKind, Machine, MachineParts};
+use super::{ConnectionKind, LookupCell, Machine, MachineParts};
 use crate::witgen::data_structures::mutable_state::MutableState;
 use crate::witgen::machines::compute_size_and_log;
 use crate::witgen::rows::RowPair;
@@ -214,6 +214,15 @@ impl<'a, T: FieldElement> DoubleSortedWitnesses16<'a, T> {
 }
 
 impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses16<'a, T> {
+    fn process_lookup_direct<'b, 'c, Q: QueryCallback<T>>(
+        &mut self,
+        _mutable_state: &'b MutableState<'a, T, Q>,
+        _identity_id: u64,
+        _values: &mut [LookupCell<'c, T>],
+    ) -> Result<bool, EvalError<T>> {
+        unimplemented!("Direct lookup not supported by machine {}.", self.name())
+    }
+
     fn identity_ids(&self) -> Vec<u64> {
         self.selector_ids.keys().cloned().collect()
     }

executor/src/witgen/machines/double_sorted_witness_machine_32.rs

@@ -3,13 +3,14 @@ use std::iter::once;
 
 use itertools::Itertools;
 
-use super::{Machine, MachineParts};
+use super::{LookupCell, Machine, MachineParts};
+use crate::witgen::data_structures::caller_data::CallerData;
 use crate::witgen::data_structures::mutable_state::MutableState;
 use crate::witgen::machines::compute_size_and_log;
+use crate::witgen::processor::OuterQuery;
 use crate::witgen::rows::RowPair;
 use crate::witgen::util::try_to_simple_poly;
-use crate::witgen::{EvalError, EvalResult, FixedData, QueryCallback};
-use crate::witgen::{EvalValue, IncompleteCause};
+use crate::witgen::{EvalError, EvalResult, EvalValue, FixedData, IncompleteCause, QueryCallback};
 
 use powdr_number::{DegreeType, FieldElement, LargeInt};
 
@@ -184,6 +185,15 @@ impl<'a, T: FieldElement> DoubleSortedWitnesses32<'a, T> {
 }
 
 impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses32<'a, T> {
+    fn process_lookup_direct<'b, 'c, Q: QueryCallback<T>>(
+        &mut self,
+        _mutable_state: &'b MutableState<'a, T, Q>,
+        identity_id: u64,
+        values: &mut [LookupCell<'c, T>],
+    ) -> Result<bool, EvalError<T>> {
+        self.process_plookup_internal(identity_id, values)
+    }
+
     fn identity_ids(&self) -> Vec<u64> {
         self.selector_ids.keys().cloned().collect()
     }
@@ -194,11 +204,33 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses32<'a, T> {
 
     fn process_plookup<Q: QueryCallback<T>>(
         &mut self,
-        _mutable_state: &MutableState<'a, T, Q>,
+        mutable_state: &MutableState<'a, T, Q>,
         identity_id: u64,
         caller_rows: &RowPair<'_, 'a, T>,
     ) -> EvalResult<'a, T> {
-        self.process_plookup_internal(identity_id, caller_rows)
+        let connection = self.parts.connections[&identity_id];
+        let outer_query = OuterQuery::new(caller_rows, connection);
+        let mut data = CallerData::from(&outer_query);
+        if self.process_lookup_direct(mutable_state, identity_id, &mut data.as_lookup_cells())? {
+            Ok(EvalResult::from(data)?.report_side_effect())
+        } else {
+            // One of the required arguments was not set, find out which:
+            let data = data.as_lookup_cells();
+            Ok(EvalValue::incomplete(
+                IncompleteCause::NonConstantRequiredArgument(
+                    match (&data[0], &data[1], &data[2], &data[3]) {
+                        (LookupCell::Output(_), _, _, _) => "operation_id",
+                        (_, LookupCell::Output(_), _, _) => "m_addr",
+                        (_, _, LookupCell::Output(_), _) => "m_step",
+                        // Note that for the mload operation, we'd expect this to be an output.
+                        // But since process_lookup_direct returned false and all other arguments are set,
+                        // we must have been in the mstore operation, in which case the value is required.
+                        (_, _, _, LookupCell::Output(_)) => "m_value",
+                        _ => unreachable!(),
+                    },
+                ),
+            ))
+        }
     }
 
     fn take_witness_col_values<'b, Q: QueryCallback<T>>(
@@ -343,107 +375,88 @@ impl<'a, T: FieldElement> DoubleSortedWitnesses32<'a, T> {
     fn process_plookup_internal(
         &mut self,
         identity_id: u64,
-        caller_rows: &RowPair<'_, 'a, T>,
-    ) -> EvalResult<'a, T> {
+        values: &mut [LookupCell<'_, T>],
+    ) -> Result<bool, EvalError<T>> {
         // We blindly assume the lookup is of the form
         // OP { operation_id, ADDR, STEP, X } is <selector> { operation_id, m_addr, m_step, m_value }
         // Where:
         // - operation_id == 0: Read
         // - operation_id == 1: Write
         // - operation_id == 2: Bootloader write
 
-        let args = self.parts.connections[&identity_id]
-            .left
-            .expressions
-            .iter()
-            .map(|e| caller_rows.evaluate(e).unwrap())
-            .collect::<Vec<_>>();
-
-        let operation_id = match args[0].constant_value() {
-            Some(v) => v,
-            None => {
-                return Ok(EvalValue::incomplete(
-                    IncompleteCause::NonConstantRequiredArgument("operation_id"),
-                ))
-            }
+        let operation_id = match values[0] {
+            LookupCell::Input(v) => v,
+            LookupCell::Output(_) => return Ok(false),
         };
+        let addr = match values[1] {
+            LookupCell::Input(v) => v,
+            LookupCell::Output(_) => return Ok(false),
+        };
+        let step = match values[2] {
+            LookupCell::Input(v) => v,
+            LookupCell::Output(_) => return Ok(false),
+        };
+        let value_ptr = &mut values[3];
 
         let selector_id = *self.selector_ids.get(&identity_id).unwrap();
 
-        let is_normal_write = operation_id == T::from(OPERATION_ID_WRITE);
-        let is_bootloader_write = operation_id == T::from(OPERATION_ID_BOOTLOADER_WRITE);
+        let is_normal_write = operation_id == &T::from(OPERATION_ID_WRITE);
+        let is_bootloader_write = operation_id == &T::from(OPERATION_ID_BOOTLOADER_WRITE);
         let is_write = is_bootloader_write || is_normal_write;
-        let addr = match args[1].constant_value() {
-            Some(v) => v,
-            None => {
-                return Ok(EvalValue::incomplete(
-                    IncompleteCause::NonConstantRequiredArgument("m_addr"),
-                ))
-            }
-        };
 
         if self.has_bootloader_write_column {
-            let is_initialized = self.is_initialized.get(&addr).cloned().unwrap_or_default();
+            let is_initialized = self.is_initialized.get(addr).cloned().unwrap_or_default();
             if !is_initialized && !is_bootloader_write {
                 panic!("Memory address {addr:x} must be initialized with a bootloader write",);
             }
-            self.is_initialized.insert(addr, true);
+            self.is_initialized.insert(*addr, true);
         }
 
-        let step = args[2]
-            .constant_value()
-            .ok_or_else(|| format!("Step must be known but is: {}", args[2]))?;
-
-        let value_expr = &args[3];
-
-        log::trace!(
-            "Query addr={:x}, step={step}, write: {is_write}, value: {}",
-            addr.to_arbitrary_integer(),
-            value_expr
-        );
-
         // TODO this does not check any of the failure modes
-        let mut assignments = EvalValue::complete(vec![]);
-        let has_side_effect = if is_write {
-            let value = match value_expr.constant_value() {
-                Some(v) => v,
-                None => {
-                    return Ok(EvalValue::incomplete(
-                        IncompleteCause::NonConstantRequiredArgument("m_value"),
-                    ))
-                }
+        let added_memory_access = if is_write {
+            let value = match value_ptr {
+                LookupCell::Input(v) => *v,
+                LookupCell::Output(_) => return Ok(false),
             };
 
             log::trace!(
                 "Memory write: addr={:x}, step={step}, value={:x}",
                 addr,
                 value
             );
-            self.data.write(addr, value);
+            self.data.write(*addr, *value);
             self.trace
                 .insert(
-                    (addr, step),
+                    (*addr, *step),
                     Operation {
                         is_normal_write,
                         is_bootloader_write,
-                        value,
+                        value: *value,
                         selector_id,
                     },
                 )
                 .is_none()
         } else {
-            let value = self.data.read(addr);
+            let value = self.data.read(*addr);
             log::trace!(
                 "Memory read: addr={:x}, step={step}, value={:x}",
                 addr,
                 value
             );
-            let ass =
-                (value_expr.clone() - value.into()).solve_with_range_constraints(caller_rows)?;
-            assignments.combine(ass);
+            match value_ptr {
+                LookupCell::Input(v) => {
+                    if *v != &value {
+                        return Err(EvalError::ConstraintUnsatisfiable(format!(
+                            "{v} != {value} (address 0x{addr:x}, time step)"
+                        )));
+                    }
+                }
+                LookupCell::Output(v) => **v = value,
+            };
+
             self.trace
                 .insert(
-                    (addr, step),
+                    (*addr, *step),
                     Operation {
                         is_normal_write,
                         is_bootloader_write,
@@ -454,16 +467,15 @@ impl<'a, T: FieldElement> DoubleSortedWitnesses32<'a, T> {
                 .is_none()
         };
         assert!(
-            has_side_effect,
+            added_memory_access,
             "Already had a memory access for address 0x{addr:x} and time step {step}!"
         );
-        assignments = assignments.report_side_effect();
 
         if self.trace.len() > (self.degree as usize) {
             return Err(EvalError::RowsExhausted(self.name.clone()));
         }
 
-        Ok(assignments)
+        Ok(true)
     }
 }
 

executor/src/witgen/machines/dynamic_machine.rs

@@ -11,7 +11,11 @@ use crate::witgen::processor::{OuterQuery, SolverState};
 use crate::witgen::rows::{Row, RowIndex, RowPair};
 use crate::witgen::sequence_iterator::{DefaultSequenceIterator, ProcessingSequenceIterator};
 use crate::witgen::vm_processor::VmProcessor;
-use crate::witgen::{AlgebraicVariable, EvalResult, EvalValue, FixedData, QueryCallback};
+use crate::witgen::{
+    AlgebraicVariable, EvalError, EvalResult, EvalValue, FixedData, QueryCallback,
+};
+
+use super::LookupCell;
 
 struct ProcessResult<'a, T: FieldElement> {
     eval_value: EvalValue<AlgebraicVariable<'a>, T>,
@@ -31,6 +35,15 @@ pub struct DynamicMachine<'a, T: FieldElement> {
 }
 
 impl<'a, T: FieldElement> Machine<'a, T> for DynamicMachine<'a, T> {
+    fn process_lookup_direct<'b, 'c, Q: QueryCallback<T>>(
+        &mut self,
+        _mutable_state: &'b MutableState<'a, T, Q>,
+        _identity_id: u64,
+        _values: &mut [LookupCell<'c, T>],
+    ) -> Result<bool, EvalError<T>> {
+        unimplemented!("Direct lookup not supported by machine {}.", self.name())
+    }
+
     fn identity_ids(&self) -> Vec<u64> {
         self.parts.identity_ids()
     }