Witgen: Match machine calls by Bus ID (#2488)

With this PR, machines are no longer called by the send's identity ID,
but by the bus ID. We still assume that a send maps to a bus ID
statically, i.e., `BusSend::bus_id()` returns `Some(...)`. This can be
relaxed in a future PR, allowing the receiver to be dynamic.

I recommend starting the review by looking into the [changes of the
`Machine`
trait](https://github.com/powdr-labs/powdr/pull/2488/files#diff-9f434c590cca98f5e5198ee0a62044e930b78487ee751ab211a7ee641501e330).

Author: georgwiese

executor/src/witgen/analysis/mod.rs

@@ -10,8 +10,7 @@ use powdr_ast::{
 use powdr_number::FieldElement;
 
 use super::{
-    machines::{Connection, ConnectionKind},
-    util::try_to_simple_poly,
+    data_structures::identity::BusReceive, machines::ConnectionKind, util::try_to_simple_poly,
     FixedData,
 };
 
@@ -20,14 +19,20 @@ use super::{
 /// On success, return the connection kind, block size and latch row.
 pub fn detect_connection_type_and_block_size<'a, T: FieldElement>(
     fixed_data: &'a FixedData<'a, T>,
-    connections: &BTreeMap<u64, Connection<'a, T>>,
+    receives: &BTreeMap<T, &'a BusReceive<T>>,
 ) -> Option<(ConnectionKind, usize, usize)> {
     // TODO we should check that the other constraints/fixed columns are also periodic.
 
     // Connecting identities should either all be permutations or all lookups.
-    let connection_type = connections
+    let connection_type = receives
         .values()
-        .map(|id| id.kind)
+        .map(|receive| {
+            if receive.has_arbitrary_multiplicity() {
+                ConnectionKind::Lookup
+            } else {
+                ConnectionKind::Permutation
+            }
+        })
         .unique()
         .exactly_one()
         .ok()?;
@@ -36,9 +41,9 @@ pub fn detect_connection_type_and_block_size<'a, T: FieldElement>(
     let (latch_row, block_size) = match connection_type {
         ConnectionKind::Lookup => {
             // We'd expect all RHS selectors to be fixed columns of the same period.
-            connections
+            receives
                 .values()
-                .map(|id| try_to_period(&id.right.selector, fixed_data))
+                .map(|receive| try_to_period(&receive.selected_payload.selector, fixed_data))
                 .unique()
                 .exactly_one()
                 .ok()??
@@ -54,8 +59,8 @@ pub fn detect_connection_type_and_block_size<'a, T: FieldElement>(
                     .max_by_key(|&(_, period)| period)
             };
             let mut latch_candidates = BTreeSet::new();
-            for id in connections.values() {
-                collect_fixed_cols(&id.right.selector, &mut latch_candidates);
+            for receive in receives.values() {
+                collect_fixed_cols(&receive.selected_payload.selector, &mut latch_candidates);
             }
             if latch_candidates.is_empty() {
                 (0, 1)

executor/src/witgen/data_structures/mutable_state.rs

@@ -18,26 +18,21 @@ use crate::witgen::{
 /// This struct uses interior mutability for accessing the machines.
 pub struct MutableState<'a, T: FieldElement, Q: QueryCallback<T>> {
     machines: Vec<RefCell<KnownMachine<'a, T>>>,
-    identity_to_machine_index: BTreeMap<u64, usize>,
+    bus_to_machine_index: BTreeMap<T, usize>,
     query_callback: &'a Q,
 }
 
 impl<'a, T: FieldElement, Q: QueryCallback<T>> MutableState<'a, T, Q> {
     pub fn new(machines: impl Iterator<Item = KnownMachine<'a, T>>, query_callback: &'a Q) -> Self {
         let machines: Vec<_> = machines.map(RefCell::new).collect();
-        let identity_to_machine_index = machines
+        let bus_to_machine_index = machines
             .iter()
             .enumerate()
-            .flat_map(|(index, m)| {
-                m.borrow()
-                    .identity_ids()
-                    .into_iter()
-                    .map(move |id| (id, index))
-            })
+            .flat_map(|(index, m)| m.borrow().bus_ids().into_iter().map(move |id| (id, index)))
             .collect();
         Self {
             machines,
-            identity_to_machine_index,
+            bus_to_machine_index,
             query_callback,
         }
     }
@@ -53,48 +48,49 @@ impl<'a, T: FieldElement, Q: QueryCallback<T>> MutableState<'a, T, Q> {
 
     pub fn can_process_call_fully(
         &self,
-        identity_id: u64,
+        bus_id: T,
         known_inputs: &BitVec,
         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)
+        let mut machine = self.responsible_machine(bus_id).ok().unwrap();
+        machine.can_process_call_fully(self, bus_id, known_inputs, range_constraints)
     }
 
     /// Call the machine responsible for the right-hand-side of an identity given its ID,
     /// the evaluated arguments and the caller's range constraints.
     pub fn call(
         &self,
-        identity_id: u64,
+        bus_id: T,
         arguments: &[AffineExpression<AlgebraicVariable<'a>, T>],
         range_constraints: &dyn RangeConstraintSet<AlgebraicVariable<'a>, T>,
     ) -> EvalResult<'a, T> {
-        self.responsible_machine(identity_id)?
-            .process_plookup_timed(self, identity_id, arguments, range_constraints)
+        self.responsible_machine(bus_id)?.process_plookup_timed(
+            self,
+            bus_id,
+            arguments,
+            range_constraints,
+        )
     }
 
     /// Call the machine responsible for the right-hand-side of an identity given its ID,
     /// use the direct interface.
     pub fn call_direct(
         &self,
-        identity_id: u64,
+        bus_id: T,
         values: &mut [LookupCell<'_, T>],
     ) -> Result<bool, EvalError<T>> {
-        self.responsible_machine(identity_id)?
-            .process_lookup_direct_timed(self, identity_id, values)
+        self.responsible_machine(bus_id)?
+            .process_lookup_direct_timed(self, bus_id, values)
     }
 
-    fn responsible_machine(
-        &self,
-        identity_id: u64,
-    ) -> Result<RefMut<KnownMachine<'a, T>>, EvalError<T>> {
+    fn responsible_machine(&self, bus_id: T) -> Result<RefMut<KnownMachine<'a, T>>, EvalError<T>> {
         let machine_index = *self
-            .identity_to_machine_index
-            .get(&identity_id)
-            .unwrap_or_else(|| panic!("No executor machine matched identity ID: {identity_id}"));
+            .bus_to_machine_index
+            .get(&bus_id)
+            .unwrap_or_else(|| panic!("No executor machine matched identity ID: {bus_id}"));
         self.machines[machine_index].try_borrow_mut().map_err(|_| {
             EvalError::RecursiveMachineCalls(format!(
-                "Detected when processing identity with ID {identity_id}"
+                "Detected when processing machine call with bus ID {bus_id}"
             ))
         })
     }

executor/src/witgen/identity_processor.rs

@@ -37,8 +37,8 @@ impl<'a, 'c, T: FieldElement, Q: QueryCallback<T>> IdentityProcessor<'a, 'c, T,
     ) -> EvalResult<'a, T> {
         let result = match identity {
             Identity::Polynomial(identity) => self.process_polynomial_identity(identity, rows),
-            Identity::BusSend(bus_interaction) => self.process_lookup_or_permutation(
-                bus_interaction.identity_id,
+            Identity::BusSend(bus_interaction) => self.process_machine_call(
+                bus_interaction.bus_id().unwrap(),
                 &bus_interaction.selected_payload,
                 rows,
             ),
@@ -65,9 +65,9 @@ impl<'a, 'c, T: FieldElement, Q: QueryCallback<T>> IdentityProcessor<'a, 'c, T,
         }
     }
 
-    fn process_lookup_or_permutation(
+    fn process_machine_call(
         &mut self,
-        id: u64,
+        bus_id: T,
         left: &'a powdr_ast::analyzed::SelectedExpressions<T>,
         rows: &RowPair<'_, 'a, T>,
     ) -> EvalResult<'a, T> {
@@ -85,7 +85,7 @@ impl<'a, 'c, T: FieldElement, Q: QueryCallback<T>> IdentityProcessor<'a, 'c, T,
             Err(incomplete_cause) => return Ok(EvalValue::incomplete(incomplete_cause)),
         };
 
-        self.mutable_state.call(id, &left, rows)
+        self.mutable_state.call(bus_id, &left, rows)
     }
 
     /// Handles the lookup that connects the current machine to the calling machine.
@@ -102,10 +102,10 @@ impl<'a, 'c, T: FieldElement, Q: QueryCallback<T>> IdentityProcessor<'a, 'c, T,
         outer_query: &OuterQuery<'a, '_, T>,
         current_rows: &RowPair<'_, 'a, T>,
     ) -> EvalResult<'a, T> {
-        let right = outer_query.connection.right;
+        let receive_payload = &outer_query.bus_receive.selected_payload;
         // sanity check that the right hand side selector is active
         current_rows
-            .evaluate(&right.selector)
+            .evaluate(&receive_payload.selector)
             .ok()
             .and_then(|affine_expression| affine_expression.constant_value())
             .and_then(|v| v.is_one().then_some(()))
@@ -116,7 +116,11 @@ impl<'a, 'c, T: FieldElement, Q: QueryCallback<T>> IdentityProcessor<'a, 'c, T,
 
         let mut updates = EvalValue::complete(vec![]);
 
-        for (l, r) in outer_query.arguments.iter().zip(right.expressions.iter()) {
+        for (l, r) in outer_query
+            .arguments
+            .iter()
+            .zip(receive_payload.expressions.iter())
+        {
             match current_rows.evaluate(r) {
                 Ok(r) => {
                     let result = (l.clone() - r).solve_with_range_constraints(&range_constraint)?;

executor/src/witgen/jit/block_machine_processor.rs

@@ -48,17 +48,20 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
         }
     }
 
-    /// Generates the JIT code for a given combination of connection and known arguments.
+    /// Generates the JIT code for a given combination of bus and known arguments.
     /// Fails if it cannot solve for the outputs, or if any sub-machine calls cannot be completed.
     pub fn generate_code(
         &self,
         can_process: impl CanProcessCall<T>,
-        identity_id: u64,
+        bus_id: T,
         known_args: &BitVec,
         known_concrete: Option<(usize, T)>,
     ) -> Result<(ProcessorResult<T>, Vec<ProverFunction<'a, T>>), String> {
-        let connection = self.machine_parts.connections[&identity_id];
-        assert_eq!(connection.right.expressions.len(), known_args.len());
+        let bus_receive = self.machine_parts.bus_receives[&bus_id];
+        assert_eq!(
+            bus_receive.selected_payload.expressions.len(),
+            known_args.len()
+        );
 
         // Set up WitgenInference with known arguments.
         let known_variables = known_args
@@ -73,7 +76,7 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
         let mut queue_items = vec![];
 
         // In the latch row, set the RHS selector to 1.
-        let selector = &connection.right.selector;
+        let selector = &bus_receive.selected_payload.selector;
         queue_items.push(QueueItem::constant_assignment(
             selector,
             T::one(),
@@ -83,15 +86,15 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
         if let Some((index, value)) = known_concrete {
             // Set the known argument to the concrete value.
             queue_items.push(QueueItem::constant_assignment(
-                &connection.right.expressions[index],
+                &bus_receive.selected_payload.expressions[index],
                 value,
                 self.latch_row as i32,
             ));
         }
 
         // Set all other selectors to 0 in the latch row.
-        for other_connection in self.machine_parts.connections.values() {
-            let other_selector = &other_connection.right.selector;
+        for other_receive in self.machine_parts.bus_receives.values() {
+            let other_selector = &other_receive.selected_payload.selector;
             if other_selector != selector {
                 queue_items.push(QueueItem::constant_assignment(
                     other_selector,
@@ -102,7 +105,7 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
         }
 
         // For each argument, connect the expression on the RHS with the formal parameter.
-        for (index, expr) in connection.right.expressions.iter().enumerate() {
+        for (index, expr) in bus_receive.selected_payload.expressions.iter().enumerate() {
             queue_items.push(QueueItem::variable_assignment(
                 expr,
                 Variable::Param(index),
@@ -175,14 +178,14 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
         .generate_code(can_process, witgen)
         .map_err(|e| {
             let err_str = e.to_string_with_variable_formatter(|var| match var {
-                Variable::Param(i) => format!("{} (connection param)", &connection.right.expressions[*i]),
+                Variable::Param(i) => format!("{} (receive param)", &bus_receive.selected_payload.expressions[*i]),
                 _ => var.to_string(),
             });
-            log::trace!("\nCode generation failed for connection:\n  {connection}");
+            log::trace!("\nCode generation failed for bus receive:\n  {bus_receive}");
             let known_args_str = known_args
                 .iter()
                 .enumerate()
-                .filter_map(|(i, b)| b.then_some(connection.right.expressions[i].to_string()))
+                .filter_map(|(i, b)| b.then_some(bus_receive.selected_payload.expressions[i].to_string()))
                 .join("\n  ");
             log::trace!("Known arguments:\n  {known_args_str}");
             log::trace!("Error:\n  {err_str}");
@@ -306,7 +309,7 @@ fn written_rows_per_column<T: FieldElement>(
         })
 }
 
-/// Returns, for each bus send ID, the collection of row offsets that have a machine call
+/// Returns, for each bus send *identity* ID, the collection of row offsets that have a machine call
 /// and if in all the calls or that row, all the arguments are known.
 /// Combines calls from branches.
 fn completed_rows_for_bus_send<T: FieldElement>(
@@ -330,21 +333,18 @@ fn fully_known_call<T: FieldElement>(e: &Effect<T, Variable>) -> bool {
     }
 }
 
-/// Returns all machine calls (bus identity and row offset) found in the effect.
+/// Returns all machine calls (bus send identity ID and row offset) found in the effect.
 /// Recurses into branches.
 fn machine_calls<T: FieldElement>(
     e: &Effect<T, Variable>,
 ) -> Box<dyn Iterator<Item = (u64, i32, &Effect<T, Variable>)> + '_> {
     match e {
-        Effect::MachineCall(id, _, arguments) => match &arguments[0] {
+        Effect::MachineCall(_, _, arguments) => match &arguments[0] {
             Variable::MachineCallParam(MachineCallVariable {
                 identity_id,
                 row_offset,
                 ..
-            }) => {
-                assert_eq!(*id, *identity_id);
-                Box::new(std::iter::once((*identity_id, *row_offset, e)))
-            }
+            }) => Box::new(std::iter::once((*identity_id, *row_offset, e))),
             _ => panic!("Expected machine call variable."),
         },
         Effect::Branch(_, first, second) => {
@@ -420,8 +420,8 @@ mod test {
             panic!("Expected exactly one matching block machine")
         };
         let (machine_parts, block_size, latch_row) = machine.machine_info();
-        assert_eq!(machine_parts.connections.len(), 1);
-        let connection_id = *machine_parts.connections.keys().next().unwrap();
+        assert_eq!(machine_parts.bus_receives.len(), 1);
+        let bus_id = *machine_parts.bus_receives.keys().next().unwrap();
         let processor = BlockMachineProcessor {
             fixed_data: &fixed_data,
             machine_parts: machine_parts.clone(),
@@ -440,7 +440,7 @@ mod test {
         );
 
         processor
-            .generate_code(&mutable_state, connection_id, &known_values, None)
+            .generate_code(&mutable_state, bus_id, &known_values, None)
             .map(|(result, _)| result)
     }
 
@@ -549,18 +549,18 @@ assert (main_binary::B[1] & 0xffffffffffff0000) == 0;
 call_var(9, 0, 2) = main_binary::B_byte[0];
 main_binary::B_byte[-1] = main_binary::B[0];
 call_var(9, -1, 2) = main_binary::B_byte[-1];
-machine_call(9, [Known(call_var(9, -1, 0)), Known(call_var(9, -1, 1)), Known(call_var(9, -1, 2)), Unknown(call_var(9, -1, 3))]);
+machine_call(2, [Known(call_var(9, -1, 0)), Known(call_var(9, -1, 1)), Known(call_var(9, -1, 2)), Unknown(call_var(9, -1, 3))]);
 main_binary::C_byte[-1] = call_var(9, -1, 3);
 main_binary::C[0] = main_binary::C_byte[-1];
-machine_call(9, [Known(call_var(9, 0, 0)), Known(call_var(9, 0, 1)), Known(call_var(9, 0, 2)), Unknown(call_var(9, 0, 3))]);
+machine_call(2, [Known(call_var(9, 0, 0)), Known(call_var(9, 0, 1)), Known(call_var(9, 0, 2)), Unknown(call_var(9, 0, 3))]);
 main_binary::C_byte[0] = call_var(9, 0, 3);
 main_binary::C[1] = (main_binary::C[0] + (main_binary::C_byte[0] * 256));
-machine_call(9, [Known(call_var(9, 1, 0)), Known(call_var(9, 1, 1)), Known(call_var(9, 1, 2)), Unknown(call_var(9, 1, 3))]);
+machine_call(2, [Known(call_var(9, 1, 0)), Known(call_var(9, 1, 1)), Known(call_var(9, 1, 2)), Unknown(call_var(9, 1, 3))]);
 main_binary::C_byte[1] = call_var(9, 1, 3);
 main_binary::C[2] = (main_binary::C[1] + (main_binary::C_byte[1] * 65536));
 main_binary::operation_id_next[2] = main_binary::operation_id[3];
 call_var(9, 2, 0) = main_binary::operation_id_next[2];
-machine_call(9, [Known(call_var(9, 2, 0)), Known(call_var(9, 2, 1)), Known(call_var(9, 2, 2)), Unknown(call_var(9, 2, 3))]);
+machine_call(2, [Known(call_var(9, 2, 0)), Known(call_var(9, 2, 1)), Known(call_var(9, 2, 2)), Unknown(call_var(9, 2, 3))]);
 main_binary::C_byte[2] = call_var(9, 2, 3);
 main_binary::C[3] = (main_binary::C[2] + (main_binary::C_byte[2] * 16777216));
 params[3] = main_binary::C[3];"
@@ -624,11 +624,11 @@ 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))]);
+machine_call(2, [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))]);
+machine_call(2, [Known(call_var(1, 1, 0))]);
 SubM::a[1] = ((SubM::b[1] * 256) + SubM::c[1]);"
         );
     }