Implement dynamic bus (runtime witgen)

executor/src/witgen/data_structures/identity.rs

@@ -295,16 +295,15 @@ fn convert_phantom_bus_interaction<T: FieldElement>(
         },
         _ => (false, bus_interaction.multiplicity.clone()),
     };
-    let bus_id = match bus_interaction.bus_id {
-        AlgebraicExpression::Number(id) => id,
-        // TODO: Relax this for sends when implementing dynamic sends
-        _ => panic!("Expected first payload entry to be a static ID"),
-    };
     let selected_payload = SelectedExpressions {
         selector: bus_interaction.latch.clone(),
         expressions: bus_interaction.payload.0.clone(),
     };
     if is_receive {
+        let bus_id = match bus_interaction.bus_id {
+            AlgebraicExpression::Number(id) => id,
+            _ => panic!("Expected first payload entry of a receive to be a static ID"),
+        };
         IdentityOrReceive::Receive(BusReceive {
             bus_id,
             multiplicity: Some(multiplicity),
@@ -314,7 +313,7 @@ fn convert_phantom_bus_interaction<T: FieldElement>(
         assert_eq!(multiplicity, bus_interaction.latch);
         IdentityOrReceive::Identity(Identity::BusSend(BusSend {
             identity_id: bus_interaction.id,
-            bus_id: AlgebraicExpression::Number(bus_id),
+            bus_id: bus_interaction.bus_id.clone(),
             selected_payload,
         }))
     }

executor/src/witgen/eval_result.rs

@@ -33,6 +33,8 @@ pub enum IncompleteCause<K = usize> {
     NonConstantQueryMatchScrutinee,
     /// Query element is not constant.
     NonConstantQueryElement,
+    /// Bus ID is not constant.
+    NonConstantBusID,
     /// A required argument was not provided
     NonConstantRequiredArgument(&'static str),
     /// The left selector in a lookup is not constant. Example: `x * {1} in [{1}]` where `x` is not constant.

executor/src/witgen/global_constraints.rs

@@ -256,7 +256,11 @@ fn propagate_constraints<T: FieldElement>(
             }
         }
         Identity::BusSend(send) => {
-            let receive = send.try_match_static(bus_receives).unwrap();
+            let receive = match send.try_match_static(bus_receives) {
+                Some(r) => r,
+                // For dynamic sends, we can only propagate constraints at runtime.
+                None => return false,
+            };
             if !send.selected_payload.selector.is_one() {
                 return false;
             }

executor/src/witgen/identity_processor.rs

@@ -7,7 +7,7 @@ use powdr_number::FieldElement;
 use crate::witgen::data_structures::mutable_state::MutableState;
 use crate::witgen::{global_constraints::CombinedRangeConstraintSet, EvalError};
 
-use super::data_structures::identity::Identity;
+use super::data_structures::identity::{BusSend, Identity};
 use super::{
     affine_expression::AlgebraicVariable, processor::OuterQuery, rows::RowPair, EvalResult,
     EvalValue, IncompleteCause, QueryCallback,
@@ -37,11 +37,7 @@ 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_machine_call(
-                bus_interaction.bus_id().unwrap(),
-                &bus_interaction.selected_payload,
-                rows,
-            ),
+            Identity::BusSend(bus_send) => self.process_machine_call(bus_send, rows),
             Identity::Connect(..) => {
                 // TODO this is not the right cause.
                 Ok(EvalValue::incomplete(IncompleteCause::SolvingFailed))
@@ -67,15 +63,15 @@ impl<'a, 'c, T: FieldElement, Q: QueryCallback<T>> IdentityProcessor<'a, 'c, T,
 
     fn process_machine_call(
         &mut self,
-        bus_id: T,
-        left: &'a powdr_ast::analyzed::SelectedExpressions<T>,
+        bus_send: &'a BusSend<T>,
         rows: &RowPair<'_, 'a, T>,
     ) -> EvalResult<'a, T> {
-        if let Some(status) = self.handle_left_selector(&left.selector, rows) {
+        if let Some(status) = self.handle_left_selector(&bus_send.selected_payload.selector, rows) {
             return Ok(status);
         }
 
-        let left = match left
+        let arguments = match bus_send
+            .selected_payload
             .expressions
             .iter()
             .map(|e| rows.evaluate(e))
@@ -85,7 +81,15 @@ 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(bus_id, &left, rows)
+        let bus_id = match rows.evaluate(&bus_send.bus_id) {
+            Ok(bus_id) => match bus_id.constant_value() {
+                Some(bus_id) => bus_id,
+                None => return Ok(EvalValue::incomplete(IncompleteCause::NonConstantBusID)),
+            },
+            Err(incomplete_cause) => return Ok(EvalValue::incomplete(incomplete_cause)),
+        };
+
+        self.mutable_state.call(bus_id, &arguments, rows)
     }
 
     /// Handles the lookup that connects the current machine to the calling machine.

executor/src/witgen/multiplicity_column_generator.rs

@@ -1,5 +1,9 @@
-use std::collections::{BTreeMap, HashMap};
+use std::{
+    collections::{BTreeMap, HashMap},
+    iter::once,
+};
 
+use itertools::Itertools;
 use powdr_ast::{
     analyzed::{AlgebraicExpression, PolyID, PolynomialType, SelectedExpressions},
     parsed::visitor::AllChildren,
@@ -35,9 +39,6 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
     ) -> HashMap<String, Vec<T>> {
         record_start(MULTIPLICITY_WITGEN_NAME);
 
-        // A map from multiplicity column ID to the vector of multiplicities.
-        let mut multiplicity_columns = BTreeMap::new();
-
         let (identities, _) = convert_identities(self.fixed.analyzed);
 
         let all_columns = witness_columns
@@ -74,8 +75,15 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
                 bus_receive.has_arbitrary_multiplicity() && bus_receive.multiplicity.is_some()
             })
             .map(|(bus_id, bus_receive)| {
-                let (size, rhs_tuples) =
-                    self.get_tuples(&terminal_values, &bus_receive.selected_payload);
+                let SelectedExpressions {
+                    selector,
+                    expressions,
+                } = &bus_receive.selected_payload;
+                let (size, rhs_tuples) = self.get_tuples(
+                    &terminal_values,
+                    selector,
+                    &expressions.iter().collect::<Vec<_>>(),
+                );
 
                 let index = rhs_tuples
                     .into_iter()
@@ -104,28 +112,55 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
             })
             .collect::<BTreeMap<_, _>>();
 
+        // A map from multiplicity column ID to the vector of multiplicities.
+        let mut multiplicity_columns = receive_infos
+            .values()
+            .map(|info| (info.multiplicity_column, vec![0; info.size]))
+            .collect::<BTreeMap<_, _>>();
+
         // Increment multiplicities for all bus sends.
-        for (bus_send, bus_receive) in identities.iter().filter_map(|i| match i {
-            Identity::BusSend(bus_send) => receive_infos
-                .get(&bus_send.bus_id().unwrap())
-                .map(|bus_receive| (bus_send, bus_receive)),
+        let bus_sends = identities.iter().filter_map(|i| match i {
+            Identity::BusSend(bus_send) => match bus_send.bus_id() {
+                // As a performance optimization, already filter out sends with a static
+                // bus ID for which we know we don't need to track multiplicities.
+                Some(bus_id) => receive_infos.get(&bus_id).map(|_| bus_send),
+                // For dynamic sends, this optimization is not possible.
+                None => Some(bus_send),
+            },
             _ => None,
-        }) {
-            let (_, lhs_tuples) = self.get_tuples(&terminal_values, &bus_send.selected_payload);
+        });
 
-            let multiplicities = multiplicity_columns
-                .entry(bus_receive.multiplicity_column)
-                .or_insert_with(|| vec![0; bus_receive.size]);
-            assert_eq!(multiplicities.len(), bus_receive.size);
+        for bus_send in bus_sends {
+            let SelectedExpressions {
+                selector,
+                expressions,
+            } = &bus_send.selected_payload;
+
+            // We need to evaluate both the bus_id (to know the run-time receive) and the expressions
+            let bus_id_and_expressions = once(&bus_send.bus_id)
+                .chain(expressions.iter())
+                .collect::<Vec<_>>();
+
+            let (_, bus_id_and_expressions) =
+                self.get_tuples(&terminal_values, selector, &bus_id_and_expressions);
 
             // Looking up the index is slow, so we do it in parallel.
-            let indices = lhs_tuples
+            let columns_and_indices = bus_id_and_expressions
                 .into_par_iter()
-                .map(|(_, tuple)| bus_receive.index[&tuple])
+                .filter_map(|(_, bus_id_and_expressions)| {
+                    receive_infos
+                        .get(&bus_id_and_expressions[0])
+                        .map(|receive_info| {
+                            (
+                                receive_info.multiplicity_column,
+                                receive_info.index[&bus_id_and_expressions[1..]],
+                            )
+                        })
+                })
                 .collect::<Vec<_>>();
 
-            for index in indices {
-                multiplicities[index] += 1;
+            for (multiplicity_column, index) in columns_and_indices {
+                multiplicity_columns.get_mut(&multiplicity_column).unwrap()[index] += 1;
             }
         }
 
@@ -150,12 +185,13 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
     fn get_tuples(
         &self,
         terminal_values: &OwnedTerminalValues<T>,
-        selected_expressions: &SelectedExpressions<T>,
+        selector: &AlgebraicExpression<T>,
+        expressions: &[&AlgebraicExpression<T>],
     ) -> (usize, Vec<(usize, Vec<T>)>) {
-        let machine_size = selected_expressions
-            .expressions
+        let machine_size = expressions
             .iter()
-            .flat_map(|expr| expr.all_children())
+            .flat_map(|e| e.all_children())
+            .chain(selector.all_children())
             .filter_map(|expr| match expr {
                 AlgebraicExpression::Reference(ref r) => match r.poly_id.ptype {
                     PolynomialType::Committed | PolynomialType::Constant => {
@@ -169,7 +205,12 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
             // But in practice, either the machine has a (smaller) witness column, or
             // it's a fixed lookup, so there is only one size.
             .min()
-            .unwrap_or_else(|| panic!("No column references found: {selected_expressions}"));
+            .unwrap_or_else(|| {
+                panic!(
+                    "No column references found: {selector} $ [{}]",
+                    expressions.iter().map(ToString::to_string).join(", ")
+                )
+            });
 
         let tuples = (0..machine_size)
             .into_par_iter()
@@ -178,14 +219,16 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
                     terminal_values.row(row),
                     &self.fixed.intermediate_definitions,
                 );
-                let result = evaluator.evaluate(&selected_expressions.selector);
+                let selector = evaluator.evaluate(selector);
 
-                assert!(result.is_zero() || result.is_one(), "Non-binary selector");
-                result.is_one().then(|| {
+                assert!(
+                    selector.is_zero() || selector.is_one(),
+                    "Non-binary selector"
+                );
+                selector.is_one().then(|| {
                     (
                         row,
-                        selected_expressions
-                            .expressions
+                        expressions
                             .iter()
                             .map(|expression| evaluator.evaluate(expression))
                             .collect::<Vec<_>>(),

executor/src/witgen/vm_processor.rs

@@ -364,9 +364,12 @@ impl<'a, 'c, T: FieldElement, Q: QueryCallback<T>> VmProcessor<'a, 'c, T, Q> {
             .filter_map(|(index, (ident, _))| match ident {
                 Identity::BusSend(send) => send
                     .try_match_static(&self.fixed_data.bus_receives)
-                    .unwrap()
-                    .has_arbitrary_multiplicity()
-                    .then_some((index, &send.selected_payload)),
+                    // We assume that the PC lookup is static.
+                    .and_then(|receive| {
+                        receive
+                            .has_arbitrary_multiplicity()
+                            .then_some((index, &send.selected_payload))
+                    }),
                 _ => None,
             })
             .max_by_key(|(_, left)| left.expressions.len())

pipeline/tests/asm.rs

@@ -133,7 +133,7 @@ fn empty() {
 
 #[test]
 // TODO: https://github.com/powdr-labs/powdr/issues/2292
-#[should_panic = "No column references found: []"]
+#[should_panic = "No column references found: 1 $ []"]
 fn single_operation() {
     let f = "asm/single_operation.asm";
     regular_test_all_fields(f, &[]);
@@ -251,7 +251,6 @@ fn static_bus_multi() {
 }
 
 #[test]
-#[should_panic = "Expected first payload entry to be a static ID"]
 fn dynamic_bus() {
     // Witgen does not currently support this.
     let f = "asm/dynamic_bus.asm";