Bus multi interaction (arbitrary number of columns) (#2457)

CI fails from my attempt to break this up to several PR, I think it's
probably because materializing the helper column isn't compatible with
our permutation/lookup -> bus witgen infrastructure. #2458, which is
#2457 plus not materializing the helper column in some cases, is passing
all CI. Is it possible to merge this with CI failures which will be
immediately fixed by #2458?

Ready for a final review. Broken up from
#2449 so that this works end to
end for arbitrary number of columns (both even and odd).

Instead of using and Option for helper columns, which currently bugs
out, this version uses an array and forces the case of odd number of bus
interactions to materialize a helper_last = multiplicity_last /
payloads_last. This has the advantage of passing tests end to end, so we
can debug in another PR.

One curious case is that CI test for block_to_block_with_bus_composite()
failed at the prover stage on the helper column witness constraint
before if I use the bus multi version for it. Maybe some updates on
witness generation for helper is needed (not on the bus accumulator but
on the conversion from non-native permutation/lookup to bus)?
@georgwiese

Author: qwang98

ast/src/analyzed/display.rs

@@ -431,7 +431,7 @@ impl<T: Display> Display for PhantomBusInteractionIdentity<T> {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         write!(
             f,
-            "Constr::PhantomBusInteraction({}, {}, [{}], {}, [{}], [{}]);",
+            "Constr::PhantomBusInteraction({}, {}, [{}], {}, [{}], [{}], {});",
             self.multiplicity,
             self.bus_id,
             self.payload.0.iter().map(ToString::to_string).format(", "),
@@ -445,6 +445,13 @@ impl<T: Display> Display for PhantomBusInteractionIdentity<T> {
                 .iter()
                 .map(ToString::to_string)
                 .format(", "),
+            match &self.helper_columns {
+                Some(helper_columns) => format!(
+                    "Option::Some([{}])",
+                    helper_columns.iter().map(ToString::to_string).format(", ")
+                ),
+                None => "Option::None".to_string(),
+            },
         )
     }
 }

ast/src/analyzed/mod.rs

@@ -1037,10 +1037,15 @@ pub struct PhantomBusInteractionIdentity<T> {
     pub payload: ExpressionList<T>,
     pub latch: AlgebraicExpression<T>,
     pub folded_expressions: ExpressionList<T>,
-    // Note that in PIL, this is a list of expressions, but we'd
-    // always expect direct column references, so this is unpacked
-    // when converting from PIL to this struct.
+    // Note that in PIL, `accumulator_columns` and
+    // `helper_columns` are lists of expressions, but we'd
+    // always expect direct column references, because
+    // they always materialize as witness columns,
+    // so they are unpacked when converting from PIL
+    // to this struct, whereas `folded_expressions`
+    // can be linear and thus optimized away by pilopt.
     pub accumulator_columns: Vec<AlgebraicReference>,
+    pub helper_columns: Option<Vec<AlgebraicReference>>,
 }
 
 impl<T> Children<AlgebraicExpression<T>> for PhantomBusInteractionIdentity<T> {

executor/src/witgen/bus_accumulator/mod.rs

@@ -21,6 +21,8 @@ mod extension_field;
 mod fp2;
 mod fp4;
 
+pub type InteractionColumns<T> = (Vec<Vec<T>>, Vec<Vec<T>>, Vec<Vec<T>>);
+
 /// Generates the second-stage columns for the bus accumulator.
 pub fn generate_bus_accumulator_columns<'a, T>(
     pil: &'a Analyzed<T>,
@@ -127,8 +129,9 @@ impl<'a, T: FieldElement, Ext: ExtensionField<T> + Sync> BusAccumulatorGenerator
             .bus_interactions
             .par_iter()
             .flat_map(|bus_interaction| {
-                let (folded, acc) = self.interaction_columns(bus_interaction);
+                let (folded, helper, acc) = self.interaction_columns(bus_interaction);
                 collect_folded_columns(bus_interaction, folded)
+                    .chain(collect_helper_columns(bus_interaction, helper))
                     .chain(collect_acc_columns(bus_interaction, acc))
                     .collect::<Vec<_>>()
             })
@@ -179,14 +182,20 @@ impl<'a, T: FieldElement, Ext: ExtensionField<T> + Sync> BusAccumulatorGenerator
         result
     }
 
+    /// Given a bus interaction and existing witness values,
+    /// calculates and returns a triple tuple of:
+    /// - the folded columns (one per bus interaction)
+    /// - one helper column per pair of bus interactions
+    /// - the accumulator column (shared by all interactions)
     fn interaction_columns(
         &self,
         bus_interaction: &PhantomBusInteractionIdentity<T>,
-    ) -> (Vec<Vec<T>>, Vec<Vec<T>>) {
+    ) -> InteractionColumns<T> {
         let intermediate_definitions = self.pil.intermediate_definitions();
 
         let size = self.values.height();
         let mut folded_list = vec![Ext::zero(); size];
+        let mut helper_list = vec![Ext::zero(); size];
         let mut acc_list = vec![Ext::zero(); size];
 
         for i in 0..size {
@@ -206,28 +215,40 @@ impl<'a, T: FieldElement, Ext: ExtensionField<T> + Sync> BusAccumulatorGenerator
                 .collect::<Vec<_>>();
             let folded = self.beta - self.fingerprint(&tuple);
 
+            let to_add = folded.inverse() * multiplicity;
+
+            let helper = match bus_interaction.helper_columns {
+                Some(_) => to_add,
+                None => Ext::zero(),
+            };
+
             let new_acc = match multiplicity.is_zero() {
                 true => current_acc,
-                false => current_acc + folded.inverse() * multiplicity,
+                false => current_acc + to_add,
             };
 
             folded_list[i] = folded;
+            helper_list[i] = helper;
             acc_list[i] = new_acc;
         }
 
         // Transpose from row-major to column-major & flatten.
         let mut folded = vec![Vec::with_capacity(size); Ext::size()];
+        let mut helper = vec![Vec::with_capacity(size); Ext::size()];
         let mut acc = vec![Vec::with_capacity(size); Ext::size()];
         for row_index in 0..size {
             for (col_index, x) in folded_list[row_index].to_vec().into_iter().enumerate() {
                 folded[col_index].push(x);
             }
+            for (col_index, x) in helper_list[row_index].to_vec().into_iter().enumerate() {
+                helper[col_index].push(x);
+            }
             for (col_index, x) in acc_list[row_index].to_vec().into_iter().enumerate() {
                 acc[col_index].push(x);
             }
         }
 
-        (folded, acc)
+        (folded, helper, acc)
     }
 
     /// Fingerprints a tuples of field elements, using the pre-computed powers of alpha.
@@ -279,3 +300,20 @@ fn collect_acc_columns<T>(
         .zip_eq(acc)
         .map(|(column_reference, column)| (column_reference.poly_id, column))
 }
+
+fn collect_helper_columns<T>(
+    bus_interaction: &PhantomBusInteractionIdentity<T>,
+    helper: Vec<Vec<T>>,
+) -> impl Iterator<Item = (PolyID, Vec<T>)> {
+    match &bus_interaction.helper_columns {
+        Some(helper_columns) => {
+            let pairs: Vec<_> = helper_columns
+                .iter()
+                .zip_eq(helper)
+                .map(|(column_reference, column)| (column_reference.poly_id, column))
+                .collect();
+            pairs.into_iter()
+        }
+        None => Vec::new().into_iter(),
+    }
+}

executor/src/witgen/data_structures/identity.rs

@@ -356,7 +356,7 @@ mod test {
             r"
 namespace main(4);
     col fixed right_latch = [0, 1]*;
-    col witness right_selector, left_latch, a, b, multiplicities, folded, acc;
+    col witness right_selector, left_latch, a, b, multiplicities, folded, acc, helper;
     {constraint}
     
     // Selectors should be binary
@@ -420,9 +420,9 @@ namespace main(4);
         // std::protocols::lookup_via_bus::lookup_send and
         // std::protocols::lookup_via_bus::lookup_receive.
         let (send, receive) = get_generated_bus_interaction_pair(
-            // The folded expressions and accumulator is ignored in both the bus send and receive, so we just use the same.
-            r"Constr::PhantomBusInteraction(main::left_latch, 42, [main::a], main::left_latch, [main::folded], [main::acc]);
-              Constr::PhantomBusInteraction(-main::multiplicities, 42, [main::b], main::right_latch, [main::folded], [main::acc]);",
+            // The folded expressions, accumulator, and helper columns are ignored in both the bus send and receive, so we just use the same.
+            r"Constr::PhantomBusInteraction(main::left_latch, 42, [main::a], main::left_latch, [main::folded], [main::acc], Option::None);
+              Constr::PhantomBusInteraction(-main::multiplicities, 42, [main::b], main::right_latch, [main::folded], [main::acc], Option::None);",
         );
         assert_eq!(
             send.selected_payload.to_string(),
@@ -478,9 +478,9 @@ namespace main(4);
         // std::protocols::permutation_via_bus::permutation_send and
         // std::protocols::permutation_via_bus::permutation_receive.
         let (send, receive) = get_generated_bus_interaction_pair(
-            // The folded expressions and accumulator is ignored in both the bus send and receive, so we just use the same.
-            r"Constr::PhantomBusInteraction(main::left_latch, 42, [main::a], main::left_latch, [main::folded], [main::acc]);
-              Constr::PhantomBusInteraction(-(main::right_latch * main::right_selector), 42, [main::b], main::right_latch * main::right_selector, [main::folded], [main::acc]);",
+            // The folded expressions, accumulator, and helper columns are ignored in both the bus send and receive, so we just use the same.
+            r"Constr::PhantomBusInteraction(main::left_latch, 42, [main::a], main::left_latch, [main::folded], [main::acc], Option::None);
+              Constr::PhantomBusInteraction(-(main::right_latch * main::right_selector), 42, [main::b], main::right_latch * main::right_selector, [main::folded], [main::acc], Option::None);",
         );
         assert_eq!(
             send.selected_payload.to_string(),

pil-analyzer/src/condenser.rs

@@ -819,6 +819,34 @@ fn to_constraint<T: FieldElement>(
                     .collect(),
                 _ => panic!("Expected array, got {:?}", fields[5]),
             },
+            helper_columns: match fields[6].as_ref() {
+                Value::Enum(enum_value) => {
+                    assert_eq!(enum_value.enum_decl.name, "std::prelude::Option");
+                    match enum_value.variant {
+                        "None" => None,
+                        "Some" => {
+                            let fields = enum_value.data.as_ref().unwrap();
+                            assert_eq!(fields.len(), 1);
+                            match fields[0].as_ref() {
+                                Value::Array(fields) => fields
+                                    .iter()
+                                    .map(|f| match to_expr(f) {
+                                        AlgebraicExpression::Reference(reference) => {
+                                            assert!(!reference.next);
+                                            reference
+                                        }
+                                        _ => panic!("Expected reference, got {f:?}"),
+                                    })
+                                    .collect::<Vec<_>>()
+                                    .into(),
+                                _ => panic!("Expected array, got {:?}", fields[0]),
+                            }
+                        }
+                        _ => panic!("Expected Some or None, got {0}", enum_value.variant),
+                    }
+                }
+                _ => panic!("Expected Enum, got {:?}", fields[6]),
+            },
         }
         .into(),
         _ => panic!("Expected constraint but got {constraint}"),

std/prelude.asm

@@ -59,7 +59,12 @@ enum Constr {
     ///   Note that this could refer to witness columns, intermediate columns, or
     ///   in-lined expressions.
     /// - The list of accumulator columns.
-    PhantomBusInteraction(expr, expr, expr[], expr, expr[], expr[])
+    /// - The list of helper columns that are intermediate values 
+    ///   (but materialized witnesses) to help calculate
+    ///   the accumulator columns, so that constraints are always bounded to
+    ///   degree 3. Each set of helper columns is always shared by two bus
+    ///   interactions.
+    PhantomBusInteraction(expr, expr, expr[], expr, expr[], expr[], Option<expr[]>)
 }
 
 /// This is the result of the "$" operator. It can be used as the left and