Rewrite multiplicity generator, remove Connection (#2538)

georgwiese · web-flow · commit 3dab8713c7fd · 2025-03-14T10:36:27.000Z
The multiplicity column generator was the last place where we used
`Connection`. This PR rewrites it and removes `Connection`.
diff --git a/executor/src/witgen/machines/mod.rs b/executor/src/witgen/machines/mod.rs
@@ -3,7 +3,7 @@ use std::fmt::Display;
 
 use bit_vec::BitVec;
 use dynamic_machine::DynamicMachine;
-use powdr_ast::analyzed::{self, AlgebraicExpression, ContainsNextRef, DegreeRange, PolyID};
+use powdr_ast::analyzed::{self, ContainsNextRef, DegreeRange, PolyID};
 
 use powdr_number::DegreeType;
 use powdr_number::FieldElement;
@@ -244,66 +244,6 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
     }
 }
 
-#[derive(Clone, Copy, Debug)]
-/// A connection is a witness generation directive to propagate rows across machines
-pub struct Connection<'a, T> {
-    pub left: &'a analyzed::SelectedExpressions<T>,
-    pub right: &'a analyzed::SelectedExpressions<T>,
-    /// For [ConnectionKind::Permutation], rows of `left` are a permutation of rows of `right`. For [ConnectionKind::Lookup], all rows in `left` are in `right`.
-    pub kind: ConnectionKind,
-    /// If the connection comes from a phantom lookup, this is the multiplicity column.
-    /// For each row of `right` it counts how often that row occurs in `left`.
-    /// Note that multiple connections can share the same multiplicity column.
-    pub multiplicity_column: Option<PolyID>,
-}
-
-impl<T: Display> Display for Connection<'_, T> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{} {} {}", self.left, self.kind, self.right)
-    }
-}
-
-impl<'a, T: FieldElement> Connection<'a, T> {
-    /// Creates a connection if the identity is a bus send.
-    pub fn try_new(
-        identity: &'a Identity<T>,
-        bus_receives: &'a BTreeMap<T, BusReceive<T>>,
-    ) -> Option<Self> {
-        match identity {
-            Identity::BusSend(bus_interaction) => {
-                let send = bus_interaction;
-                let receive = send
-                    .try_match_static(bus_receives)
-                    .expect("No matching receive!");
-                let multiplicity_column = if receive.has_arbitrary_multiplicity() {
-                    receive
-                        .multiplicity
-                        .as_ref()
-                        .and_then(|multiplicity| match multiplicity {
-                            AlgebraicExpression::Reference(reference) => Some(reference.poly_id),
-                            // For receives of permutations, we would have complex expressions here.
-                            _ => None,
-                        })
-                } else {
-                    // For permutations, the selector is already generated by "normal" witgen.
-                    None
-                };
-                Some(Connection {
-                    left: &send.selected_payload,
-                    right: &receive.selected_payload,
-                    kind: if receive.has_arbitrary_multiplicity() {
-                        ConnectionKind::Lookup
-                    } else {
-                        ConnectionKind::Permutation
-                    },
-                    multiplicity_column,
-                })
-            }
-            _ => None,
-        }
-    }
-}
-
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
 pub enum ConnectionKind {
     Lookup,
diff --git a/executor/src/witgen/multiplicity_column_generator.rs b/executor/src/witgen/multiplicity_column_generator.rs
@@ -1,22 +1,19 @@
 use std::collections::{BTreeMap, HashMap};
 
 use powdr_ast::{
-    analyzed::{AlgebraicExpression, PolynomialType, SelectedExpressions},
+    analyzed::{AlgebraicExpression, PolyID, PolynomialType, SelectedExpressions},
     parsed::visitor::AllChildren,
 };
 use powdr_executor_utils::expression_evaluator::{ExpressionEvaluator, OwnedTerminalValues};
 use powdr_number::FieldElement;
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
 
 use crate::witgen::{
-    data_structures::identity::convert_identities,
-    machines::{
-        profiling::{record_end, record_start},
-        Connection,
-    },
+    data_structures::identity::{convert_identities, Identity},
+    machines::profiling::{record_end, record_start},
 };
 
-use super::FixedData;
+use super::{util::try_to_simple_poly, FixedData};
 
 static MULTIPLICITY_WITGEN_NAME: &str = "multiplicity witgen";
 
@@ -30,28 +27,18 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
     }
 
     /// Takes a map of witness columns and extends it with the multiplicity columns
-    /// reference in the phantom lookups.
+    /// referenced by bus sends with a non-binary multiplicity.
     pub fn generate(
         &self,
         witness_columns: HashMap<String, Vec<T>>,
         publics: BTreeMap<String, Option<T>>,
     ) -> HashMap<String, Vec<T>> {
         record_start(MULTIPLICITY_WITGEN_NAME);
 
-        log::trace!("Starting multiplicity witness generation.");
-        let start = std::time::Instant::now();
-
-        // Several range constraints might point to the same target
+        // 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 phantom_lookups = identities
-            .iter()
-            .filter_map(|identity| {
-                Connection::try_new(identity, &self.fixed.bus_receives)
-                    .and_then(|connection| connection.multiplicity_column.map(|_| connection))
-            })
-            .collect::<Vec<_>>();
 
         let all_columns = witness_columns
             .into_iter()
@@ -78,64 +65,65 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
             challenge_values: self.fixed.challenges.clone(),
         };
 
-        log::trace!(
-            "  Done building trace values, took: {}s",
-            start.elapsed().as_secs_f64()
-        );
-
-        for lookup in phantom_lookups {
-            log::trace!("  Updating multiplicity for phantom lookup: {lookup}");
-            let start = std::time::Instant::now();
-
-            let (rhs_machine_size, rhs_tuples) = self.get_tuples(&terminal_values, lookup.right);
-
-            log::trace!(
-                "    Done collecting RHS tuples, took {}s",
-                start.elapsed().as_secs_f64()
-            );
-            let start = std::time::Instant::now();
-
-            let index = rhs_tuples
-                .iter()
-                .map(|(i, tuple)| {
-                    // There might be multiple identical rows, but it's fine, we can pick any.
-                    (tuple, *i)
-                })
-                .collect::<HashMap<_, _>>();
-
-            log::trace!(
-                "    Done building index, took {}s",
-                start.elapsed().as_secs_f64()
-            );
-            let start = std::time::Instant::now();
+        // Index all bus receives with arbitrary multiplicity.
+        let receive_infos = self
+            .fixed
+            .bus_receives
+            .iter()
+            .filter(|(_, bus_receive)| {
+                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 (_, lhs_tuples) = self.get_tuples(&terminal_values, lookup.left);
+                let index = rhs_tuples
+                    .into_iter()
+                    .map(|(i, tuple)| {
+                        // There might be multiple identical rows, but it's fine, we can pick any.
+                        (tuple, i)
+                    })
+                    .collect::<HashMap<_, _>>();
+
+                let multiplicity = bus_receive.multiplicity.as_ref().unwrap();
+                (
+                    *bus_id,
+                    ReceiveInfo {
+                        multiplicity_column: try_to_simple_poly(multiplicity)
+                            .unwrap_or_else(|| {
+                                panic!("Expected simple reference, got: {multiplicity}")
+                            })
+                            .poly_id,
+                        size,
+                        index,
+                    },
+                )
+            })
+            .collect::<BTreeMap<_, _>>();
 
-            log::trace!(
-                "    Done collecting LHS tuples, took: {}s",
-                start.elapsed().as_secs_f64()
-            );
-            let start = std::time::Instant::now();
+        // 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)),
+            _ => None,
+        }) {
+            let (_, lhs_tuples) = self.get_tuples(&terminal_values, &bus_send.selected_payload);
 
-            let multiplicity_column_id = lookup.multiplicity_column.unwrap();
             let multiplicities = multiplicity_columns
-                .entry(multiplicity_column_id)
-                .or_insert_with(|| vec![0; rhs_machine_size]);
-            assert_eq!(multiplicities.len(), rhs_machine_size);
+                .entry(bus_receive.multiplicity_column)
+                .or_insert_with(|| vec![0; bus_receive.size]);
+            assert_eq!(multiplicities.len(), bus_receive.size);
 
             // Looking up the index is slow, so we do it in parallel.
             let indices = lhs_tuples
                 .into_par_iter()
-                .map(|(_, tuple)| index[&tuple])
+                .map(|(_, tuple)| bus_receive.index[&tuple])
                 .collect::<Vec<_>>();
 
             for index in indices {
                 multiplicities[index] += 1;
             }
-            log::trace!(
-                "    Done updating multiplicities, took: {}s",
-                start.elapsed().as_secs_f64()
-            );
         }
 
         let columns = terminal_values
@@ -206,3 +194,10 @@ impl<'a, T: FieldElement> MultiplicityColumnGenerator<'a, T> {
         (machine_size, tuples)
     }
 }
+
+struct ReceiveInfo<T> {
+    multiplicity_column: PolyID,
+    size: usize,
+    /// Maps a tuple of values to its index in the trace.
+    index: HashMap<Vec<T>, usize>,
+}
