Add folded payload columns to PhantomBusInteraction (#2426)

Depends on #2427 (because otherwise a referenced column might be removed
by the optimizer).

This PR is completely analogous to #2412: We add a list of expressions
that evaluate to the folded payload to `PhantomBusInteraction`. This
lets us easily find the referenced columns in the manual bus witgen, see
#2428.

ast/src/analyzed/display.rs

@@ -431,11 +431,16 @@ 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(", "),
             self.latch,
+            self.folded_expressions
+                .0
+                .iter()
+                .map(ToString::to_string)
+                .format(", "),
             self.accumulator_columns
                 .iter()
                 .map(ToString::to_string)

ast/src/analyzed/mod.rs

@@ -1008,6 +1008,7 @@ impl<T> Children<AlgebraicExpression<T>> for ExpressionList<T> {
 #[derive(
     Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema, Hash,
 )]
+/// For documentation, see the equivalent `Constr` variant in std/prelude.asm.
 pub struct PhantomBusInteractionIdentity<T> {
     // The ID is globally unique among identities.
     pub id: u64,
@@ -1016,6 +1017,10 @@ pub struct PhantomBusInteractionIdentity<T> {
     pub bus_id: AlgebraicExpression<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.
     pub accumulator_columns: Vec<AlgebraicReference>,
 }
 

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, acc;
+    col witness right_selector, left_latch, a, b, multiplicities, folded, acc;
     {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 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::acc]);
-              Constr::PhantomBusInteraction(-main::multiplicities, 42, [main::b], main::right_latch, [main::acc]);",
+            // 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]);",
         );
         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 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::acc]);
-              Constr::PhantomBusInteraction(-(main::right_latch * main::right_selector), 42, [main::b], main::right_latch * main::right_selector, [main::acc]);",
+            // 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]);",
         );
         assert_eq!(
             send.selected_payload.to_string(),

pil-analyzer/src/condenser.rs

@@ -802,7 +802,11 @@ fn to_constraint<T: FieldElement>(
                 _ => panic!("Expected array, got {:?}", fields[2]),
             }),
             latch: to_expr(&fields[3]),
-            accumulator_columns: match fields[4].as_ref() {
+            folded_expressions: ExpressionList(match fields[4].as_ref() {
+                Value::Array(fields) => fields.iter().map(|f| to_expr(f)).collect(),
+                _ => panic!("Expected array, got {:?}", fields[4]),
+            }),
+            accumulator_columns: match fields[5].as_ref() {
                 Value::Array(fields) => fields
                     .iter()
                     .map(|f| match to_expr(f) {
@@ -813,7 +817,7 @@ fn to_constraint<T: FieldElement>(
                         _ => panic!("Expected reference, got {f:?}"),
                     })
                     .collect(),
-                _ => panic!("Expected array, got {:?}", fields[2]),
+                _ => panic!("Expected array, got {:?}", fields[5]),
             },
         }
         .into(),

std/prelude.asm

@@ -54,8 +54,12 @@ enum Constr {
     ///   would be in an equivalent lookup or permutation:
     ///   - It should always evaluate to a binary value.
     ///   - If it evaluates to zero, the multiplicity must be zero.
+    /// - A list of expressions that evaluate to the value of the folded payload,
+    ///   i.e., `beta - fingerprint_with_id(id, payload, alpha)`
+    ///   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[])
+    PhantomBusInteraction(expr, expr, expr[], expr, expr[], expr[])
 }
 
 /// This is the result of the "$" operator. It can be used as the left and

std/protocols/bus.asm

@@ -91,7 +91,7 @@ let bus_interaction: expr, expr[], expr, expr -> () = constr |id, payload, multi
     constrain_eq_ext(update_expr, from_base(0));
 
     // Add phantom bus interaction
-    Constr::PhantomBusInteraction(multiplicity, id, payload, latch, acc);
+    Constr::PhantomBusInteraction(multiplicity, id, payload, latch, unpack_ext_array(folded), acc);
 };
 
 /// Compute acc' = acc * (1 - is_first') + multiplicity' / fingerprint(id, payload...),