merge bytecode / memory / precompile lookups into a shared BusInteraction

crates/lean_prover/src/prove_execution.rs

@@ -91,10 +91,14 @@ pub fn prove_execution(
     let mut memory_acc = F::zero_vec(memory.len());
     info_span!("Building memory access count").in_scope(|| {
         for (table, trace) in &traces {
-            for lookup in table.lookups() {
-                for i in &trace.columns[lookup.index] {
-                    for j in 0..lookup.values.len() {
-                        memory_acc[i.to_usize() + j] += F::ONE;
+            let buses = table.bus_interactions();
+            for group in memory_lookup_groups(&buses) {
+                let idx_col = &trace.columns[group.idx_col];
+                let n = group.value_cols.len();
+                for idx in idx_col {
+                    let base = idx.to_usize();
+                    for ofs in 0..n {
+                        memory_acc[base + ofs] += F::ONE;
                     }
                 }
             }
@@ -122,7 +126,7 @@ pub fn prove_execution(
     // logup (GKR)
     let logup_c = prover_state.sample();
     prover_state.duplex();
-    let logup_alphas = prover_state.sample_vec(log2_ceil_usize(max_bus_width_including_bytecode()));
+    let logup_alphas = prover_state.sample_vec(LOG_MAX_BUS_WIDTH);
     let logup_alphas_eq_poly = eval_eq(&logup_alphas);
 
     let logup_statements = prove_generic_logup(
@@ -181,11 +185,11 @@ pub fn prove_execution(
         let bus_numerator_value = logup_statements.bus_numerators_values[table];
         let bus_denominator_value = logup_statements.bus_denominators_values[table];
         let bus_final_value = bus_numerator_value
-            * match table.bus().direction {
+            * match table.bus_interactions()[0].direction {
                 BusDirection::Pull => EF::NEG_ONE,
                 BusDirection::Push => EF::ONE,
             }
-            + bus_beta * (bus_denominator_value - logup_c);
+            + bus_beta * (logup_c - bus_denominator_value);
 
         let eq_suffix = from_end(gkr_point, *log_n_rows).to_vec();
 

crates/lean_prover/src/verify_execution.rs

@@ -78,7 +78,7 @@ pub fn verify_execution(
 
     let logup_c = verifier_state.sample();
     verifier_state.duplex();
-    let logup_alphas = verifier_state.sample_vec(log2_ceil_usize(max_bus_width_including_bytecode()));
+    let logup_alphas = verifier_state.sample_vec(LOG_MAX_BUS_WIDTH);
     let logup_alphas_eq_poly = eval_eq(&logup_alphas);
 
     let logup_statements = verify_generic_logup(
@@ -126,11 +126,11 @@ pub fn verify_execution(
         let bus_numerator_value = logup_statements.bus_numerators_values[table];
         let bus_denominator_value = logup_statements.bus_denominators_values[table];
         let bus_final_value = bus_numerator_value
-            * match table.bus().direction {
+            * match table.bus_interactions()[0].direction {
                 BusDirection::Pull => EF::NEG_ONE,
                 BusDirection::Push => EF::ONE,
             }
-            + bus_beta * (bus_denominator_value - logup_c);
+            + bus_beta * (logup_c - bus_denominator_value);
 
         initial_sum += eta_power * bus_final_value;
 

crates/lean_vm/src/core/constants.rs

@@ -48,7 +48,7 @@ mod tests {
     #[test]
     fn ensure_no_overflow_in_logup() {
         fn memory_lookups_count<T: TableT>(t: &T) -> usize {
-            t.lookups().iter().map(|l| l.values.len()).sum::<usize>()
+            t.bus_interactions().iter().filter(|bus| bus.is_memory_lookup()).count()
         }
         // memory lookup
         let mut max_memory_logup_sum: u64 = 0;

crates/lean_vm/src/tables/execution/mod.rs

@@ -25,34 +25,33 @@ impl<const BUS: bool> TableT for ExecutionTable<BUS> {
         N_TOTAL_EXECUTION_COLUMNS + N_TEMPORARY_EXEC_COLUMNS
     }
 
-    fn lookups(&self) -> Vec<LookupIntoMemory> {
-        vec![
-            LookupIntoMemory {
-                index: COL_MEM_ADDRESS_A,
-                values: vec![COL_MEM_VALUE_A],
-            },
-            LookupIntoMemory {
-                index: COL_MEM_ADDRESS_B,
-                values: vec![COL_MEM_VALUE_B],
-            },
-            LookupIntoMemory {
-                index: COL_MEM_ADDRESS_C,
-                values: vec![COL_MEM_VALUE_C],
-            },
-        ]
-    }
-
-    fn bus(&self) -> Bus {
-        Bus {
+    fn bus_interactions(&self) -> Vec<BusInteraction> {
+        let bytecode_lookup = BusInteraction {
+            direction: BusDirection::Push,
+            multiplicity: BusMultiplicity::One,
+            domainsep: BusData::Constant(LOGUP_BYTECODE_DOMAINSEP),
+            data: (0..N_INSTRUCTION_COLUMNS)
+                .map(|i| BusData::Column(N_RUNTIME_COLUMNS + i))
+                .chain(std::iter::once(BusData::Column(COL_PC)))
+                .collect(),
+        };
+        let precompile_bus = BusInteraction {
             direction: BusDirection::Push,
-            multiplicity: COL_IS_PRECOMPILE,
+            multiplicity: BusMultiplicity::Column(COL_IS_PRECOMPILE),
             domainsep: BusData::Column(COL_PRECOMPILE_DOMAINSEP),
             data: vec![
                 BusData::Column(COL_EXEC_NU_A),
                 BusData::Column(COL_EXEC_NU_B),
                 BusData::Column(COL_EXEC_NU_C),
             ],
-        }
+        };
+        // Convention shared with the other tables: the unique Multiplicity::Column bus
+        // comes first; everything that follows is Multiplicity::One.
+        let mut buses = vec![precompile_bus, bytecode_lookup];
+        buses.extend(memory_lookups_consecutive(COL_MEM_ADDRESS_A, COL_MEM_VALUE_A, 1));
+        buses.extend(memory_lookups_consecutive(COL_MEM_ADDRESS_B, COL_MEM_VALUE_B, 1));
+        buses.extend(memory_lookups_consecutive(COL_MEM_ADDRESS_C, COL_MEM_VALUE_C, 1));
+        buses
     }
 
     fn padding_row(&self, zero_vec_ptr: usize, _null_hash_ptr: usize, ending_pc: usize) -> Vec<F> {

crates/lean_vm/src/tables/extension_op/mod.rs

@@ -87,34 +87,21 @@ impl<const BUS: bool> TableT for ExtensionOpPrecompile<BUS> {
         Table::extension_op()
     }
 
-    fn lookups(&self) -> Vec<LookupIntoMemory> {
-        vec![
-            LookupIntoMemory {
-                index: COL_IDX_A,
-                values: (COL_VA..COL_VA + DIMENSION).collect(),
-            },
-            LookupIntoMemory {
-                index: COL_IDX_B,
-                values: (COL_VB..COL_VB + DIMENSION).collect(),
-            },
-            LookupIntoMemory {
-                index: COL_IDX_RES,
-                values: (COL_VRES..COL_VRES + DIMENSION).collect(),
-            },
-        ]
-    }
-
-    fn bus(&self) -> Bus {
-        Bus {
+    fn bus_interactions(&self) -> Vec<BusInteraction> {
+        let mut buses = vec![BusInteraction {
             direction: BusDirection::Pull,
-            multiplicity: COL_MULTIPLICITY_EXTENSION_OP,
+            multiplicity: BusMultiplicity::Column(COL_MULTIPLICITY_EXTENSION_OP),
             domainsep: BusData::Column(COL_DOMAINSEP_EXTENSION_OP),
             data: vec![
                 BusData::Column(COL_IDX_A),
                 BusData::Column(COL_IDX_B),
                 BusData::Column(COL_IDX_RES),
             ],
-        }
+        }];
+        buses.extend(memory_lookups_consecutive(COL_IDX_A, COL_VA, DIMENSION));
+        buses.extend(memory_lookups_consecutive(COL_IDX_B, COL_VB, DIMENSION));
+        buses.extend(memory_lookups_consecutive(COL_IDX_RES, COL_VRES, DIMENSION));
+        buses
     }
 
     fn n_columns_total(&self) -> usize {

crates/lean_vm/src/tables/poseidon_16/mod.rs

@@ -138,44 +138,42 @@ impl<const BUS: bool> TableT for Poseidon16Precompile<BUS> {
         Table::poseidon16()
     }
 
-    fn lookups(&self) -> Vec<LookupIntoMemory> {
-        vec![
-            LookupIntoMemory {
-                index: POSEIDON_16_COL_EFFECTIVE_INDEX_LEFT_FIRST,
-                values: (POSEIDON_16_COL_INPUT_START..POSEIDON_16_COL_INPUT_START + HALF_DIGEST_LEN).collect(),
-            },
-            LookupIntoMemory {
-                index: POSEIDON_16_COL_EFFECTIVE_INDEX_LEFT_SECOND,
-                values: (POSEIDON_16_COL_INPUT_START + HALF_DIGEST_LEN..POSEIDON_16_COL_INPUT_START + DIGEST_LEN)
-                    .collect(),
-            },
-            LookupIntoMemory {
-                index: POSEIDON_16_COL_INDEX_INPUT_RIGHT,
-                values: (POSEIDON_16_COL_INPUT_START + DIGEST_LEN..POSEIDON_16_COL_INPUT_START + DIGEST_LEN * 2)
-                    .collect(),
-            },
-            LookupIntoMemory {
-                index: POSEIDON_16_COL_INDEX_INPUT_RES,
-                values: (POSEIDON_16_COL_OUTPUT_LEFT..POSEIDON_16_COL_OUTPUT_LEFT + DIGEST_LEN * 2).collect(),
-            },
-        ]
-    }
-
     fn n_columns_total(&self) -> usize {
         num_cols_total_poseidon_16()
     }
 
-    fn bus(&self) -> Bus {
-        Bus {
+    fn bus_interactions(&self) -> Vec<BusInteraction> {
+        let mut buses = vec![BusInteraction {
             direction: BusDirection::Pull,
-            multiplicity: POSEIDON_16_COL_MULTIPLICITY,
+            multiplicity: BusMultiplicity::Column(POSEIDON_16_COL_MULTIPLICITY),
             domainsep: BusData::Column(POSEIDON_16_COL_DOMAINSEP),
             data: vec![
                 BusData::Column(POSEIDON_16_COL_INDEX_INPUT_LEFT),
                 BusData::Column(POSEIDON_16_COL_INDEX_INPUT_RIGHT),
                 BusData::Column(POSEIDON_16_COL_INDEX_INPUT_RES),
             ],
-        }
+        }];
+        buses.extend(memory_lookups_consecutive(
+            POSEIDON_16_COL_EFFECTIVE_INDEX_LEFT_FIRST,
+            POSEIDON_16_COL_INPUT_START,
+            HALF_DIGEST_LEN,
+        ));
+        buses.extend(memory_lookups_consecutive(
+            POSEIDON_16_COL_EFFECTIVE_INDEX_LEFT_SECOND,
+            POSEIDON_16_COL_INPUT_START + HALF_DIGEST_LEN,
+            HALF_DIGEST_LEN,
+        ));
+        buses.extend(memory_lookups_consecutive(
+            POSEIDON_16_COL_INDEX_INPUT_RIGHT,
+            POSEIDON_16_COL_INPUT_START + DIGEST_LEN,
+            DIGEST_LEN,
+        ));
+        buses.extend(memory_lookups_consecutive(
+            POSEIDON_16_COL_INDEX_INPUT_RES,
+            POSEIDON_16_COL_OUTPUT_LEFT,
+            DIGEST_LEN * 2,
+        ));
+        buses
     }
 
     fn padding_row(&self, zero_vec_ptr: usize, null_hash_ptr: usize, _ending_pc: usize) -> Vec<F> {

crates/lean_vm/src/tables/table_enum.rs

@@ -5,7 +5,8 @@ use crate::*;
 
 pub const N_TABLES: usize = 3;
 pub const ALL_TABLES: [Table; N_TABLES] = [Table::execution(), Table::extension_op(), Table::poseidon16()];
-pub const MAX_PRECOMPILE_BUS_WIDTH: usize = 4;
+pub const MAX_BUS_WIDTH: usize = N_INSTRUCTION_COLUMNS + 2; // + 1 for PC, + 1 for domainsep
+pub const LOG_MAX_BUS_WIDTH: usize = log2_ceil_usize(MAX_BUS_WIDTH);
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[repr(usize)]
@@ -60,14 +61,11 @@ impl TableT for Table {
     fn table(&self) -> Table {
         delegate_to_inner!(self, table)
     }
-    fn lookups(&self) -> Vec<LookupIntoMemory> {
-        delegate_to_inner!(self, lookups)
-    }
     fn is_execution_table(&self) -> bool {
         delegate_to_inner!(self, is_execution_table)
     }
-    fn bus(&self) -> Bus {
-        delegate_to_inner!(self, bus)
+    fn bus_interactions(&self) -> Vec<BusInteraction> {
+        delegate_to_inner!(self, bus_interactions)
     }
     fn padding_row(&self, zero_vec_ptr: usize, null_hash_ptr: usize, ending_pc: usize) -> Vec<PF<EF>> {
         delegate_to_inner!(self, padding_row, zero_vec_ptr, null_hash_ptr, ending_pc)
@@ -106,12 +104,6 @@ impl Air for Table {
     }
 }
 
-pub fn max_bus_width_including_bytecode() -> usize {
-    MAX_PRECOMPILE_BUS_WIDTH
-        .max(N_INSTRUCTION_COLUMNS + 2) // 2 = +1 for PC, +1 for domainsep
-        .next_power_of_two()
-}
-
 pub fn max_air_constraints() -> usize {
     ALL_TABLES.iter().map(|table| table.n_constraints()).max().unwrap()
 }
@@ -128,9 +120,13 @@ mod tests {
     }
 
     #[test]
-    fn test_max_precompile_bus_width() {
-        // +1 for the domainsep
-        let expected_max_bus_width = ALL_TABLES.iter().map(|table| table.bus().data.len() + 1).max().unwrap();
-        assert_eq!(MAX_PRECOMPILE_BUS_WIDTH, expected_max_bus_width);
+    fn test_max_bus_width() {
+        let expected_max_bus_width = ALL_TABLES
+            .iter()
+            .flat_map(|table| table.bus_interactions())
+            .map(|bus| bus.data.len() + 1)
+            .max()
+            .unwrap();
+        assert_eq!(MAX_BUS_WIDTH, expected_max_bus_width);
     }
 }