Implement basic dynamic bus

Author: georgwiese

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,10 @@ 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,
+                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_interaction) => self.process_machine_call(bus_interaction, 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

@@ -35,9 +35,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
@@ -75,11 +72,11 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
             })
             .map(|(bus_id, bus_receive)| {
                 let (size, rhs_tuples) =
-                    self.get_tuples(&terminal_values, &bus_receive.selected_payload);
+                    self.get_tuples(&terminal_values, &bus_receive.selected_payload, None);
 
                 let index = rhs_tuples
                     .into_iter()
-                    .map(|(i, tuple)| {
+                    .map(|(i, tuple, _)| {
                         // There might be multiple identical rows, but it's fine, we can pick any.
                         (tuple, i)
                     })
@@ -104,28 +101,35 @@ 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)),
+        for bus_send in identities.iter().filter_map(|i| match i {
+            Identity::BusSend(bus_send) => 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);
+            let (_, lhs_tuples) = self.get_tuples(
+                &terminal_values,
+                &bus_send.selected_payload,
+                Some(&bus_send.bus_id),
+            );
 
             // Looking up the index is slow, so we do it in parallel.
             let indices = lhs_tuples
                 .into_par_iter()
-                .map(|(_, tuple)| bus_receive.index[&tuple])
+                .filter_map(|(_, tuple, bus_id)| {
+                    receive_infos.get(&bus_id.unwrap()).map(|receive_info| {
+                        (receive_info.multiplicity_column, receive_info.index[&tuple])
+                    })
+                })
                 .collect::<Vec<_>>();
 
-            for index in indices {
-                multiplicities[index] += 1;
+            for (multiplicity_column, index) in indices {
+                multiplicity_columns.get_mut(&multiplicity_column).unwrap()[index] += 1;
             }
         }
 
@@ -151,7 +155,8 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
         &self,
         terminal_values: &OwnedTerminalValues<T>,
         selected_expressions: &SelectedExpressions<T>,
-    ) -> (usize, Vec<(usize, Vec<T>)>) {
+        bus_id: Option<&AlgebraicExpression<T>>,
+    ) -> (usize, Vec<(usize, Vec<T>, Option<T>)>) {
         let machine_size = selected_expressions
             .expressions
             .iter()
@@ -178,17 +183,21 @@ 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(&selected_expressions.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
                             .iter()
                             .map(|expression| evaluator.evaluate(expression))
                             .collect::<Vec<_>>(),
+                        bus_id.map(|bus_id| evaluator.evaluate(bus_id)),
                     )
                 })
             })

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)),
+                    // Filter dynamic lookups
+                    .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

@@ -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";