Feat/shrink cpu byte alu

prover/src/lib.rs

@@ -48,11 +48,12 @@ use crate::tables::trace_builder::Traces;
 use crate::tables::trace_builder::count_table_lengths;
 use crate::tables::types::BusId;
 use crate::test_utils::{
-    E, F, VmAir, create_bitwise_air, create_branch_air, create_commit_air, create_cpu_air,
-    create_decode_air, create_dvrm_air, create_halt_air, create_keccak_air, create_keccak_rc_air,
-    create_keccak_rnd_air, create_load_air, create_lt_air, create_memw_air,
-    create_memw_aligned_air, create_memw_register_air, create_mul_air, create_page_air,
-    create_register_air, create_shift_air,
+    E, F, VmAir, create_bitwise_air, create_branch_air, create_bytewise_air, create_commit_air,
+    create_cpu_air, create_cpu32_air, create_decode_air, create_dvrm_air, create_eq_air,
+    create_halt_air, create_keccak_air, create_keccak_rc_air, create_keccak_rnd_air,
+    create_load_air, create_lt_air, create_memw_air, create_memw_aligned_air,
+    create_memw_register_air, create_mul_air, create_page_air, create_register_air,
+    create_shift_air, create_store_air,
 };
 
 use stark::proof::options::{GoldilocksCubicProofOptions, ProofOptions};
@@ -84,6 +85,11 @@ pub struct TableCounts {
     pub shift: usize,
     pub branch: usize,
     pub memw_register: usize,
+    // Auxiliary ALU / memory / CPU32 dispatch chips
+    pub eq: usize,
+    pub bytewise: usize,
+    pub store: usize,
+    pub cpu32: usize,
 }
 
 impl TableCounts {
@@ -99,6 +105,10 @@ impl TableCounts {
             + self.shift
             + self.branch
             + self.memw_register
+            + self.eq
+            + self.bytewise
+            + self.store
+            + self.cpu32
     }
 
     /// Validate that all required tables have at least one chunk.
@@ -117,6 +127,10 @@ impl TableCounts {
             ("shift", self.shift),
             ("branch", self.branch),
             ("memw_register", self.memw_register),
+            ("eq", self.eq),
+            ("bytewise", self.bytewise),
+            ("store", self.store),
+            ("cpu32", self.cpu32),
         ];
         for (name, count) in checks {
             if count == 0 {
@@ -212,6 +226,11 @@ pub(crate) struct VmAirs {
     pub register: VmAir,
     pub pages: Vec<VmAir>,
     pub memw_registers: Vec<VmAir>,
+    // Auxiliary ALU / memory / CPU32 dispatch chips
+    pub eqs: Vec<VmAir>,
+    pub bytewises: Vec<VmAir>,
+    pub stores: Vec<VmAir>,
+    pub cpu32s: Vec<VmAir>,
 }
 
 impl VmAirs {
@@ -269,6 +288,18 @@ impl VmAirs {
         {
             pairs.push((air, trace, &()));
         }
+        for (air, trace) in self.eqs.iter().zip(traces.eqs.iter_mut()) {
+            pairs.push((air, trace, &()));
+        }
+        for (air, trace) in self.bytewises.iter().zip(traces.bytewises.iter_mut()) {
+            pairs.push((air, trace, &()));
+        }
+        for (air, trace) in self.stores.iter().zip(traces.stores.iter_mut()) {
+            pairs.push((air, trace, &()));
+        }
+        for (air, trace) in self.cpu32s.iter().zip(traces.cpu32s.iter_mut()) {
+            pairs.push((air, trace, &()));
+        }
 
         pairs
     }
@@ -319,6 +350,18 @@ impl VmAirs {
         for air in &self.memw_registers {
             refs.push(air);
         }
+        for air in &self.eqs {
+            refs.push(air);
+        }
+        for air in &self.bytewises {
+            refs.push(air);
+        }
+        for air in &self.stores {
+            refs.push(air);
+        }
+        for air in &self.cpu32s {
+            refs.push(air);
+        }
 
         refs
     }
@@ -424,6 +467,18 @@ impl VmAirs {
         let memw_registers: Vec<_> = (0..table_counts.memw_register)
             .map(|i| create_memw_register_air(proof_options).with_name(&format!("MEMW_R[{}]", i)))
             .collect();
+        let eqs: Vec<_> = (0..table_counts.eq)
+            .map(|i| create_eq_air(proof_options).with_name(&format!("EQ[{}]", i)))
+            .collect();
+        let bytewises: Vec<_> = (0..table_counts.bytewise)
+            .map(|i| create_bytewise_air(proof_options).with_name(&format!("BYTEWISE[{}]", i)))
+            .collect();
+        let stores: Vec<_> = (0..table_counts.store)
+            .map(|i| create_store_air(proof_options).with_name(&format!("STORE[{}]", i)))
+            .collect();
+        let cpu32s: Vec<_> = (0..table_counts.cpu32)
+            .map(|i| create_cpu32_air(proof_options).with_name(&format!("CPU32[{}]", i)))
+            .collect();
 
         #[cfg(feature = "debug-checks")]
         debug_report::print_bus_legend();
@@ -448,6 +503,10 @@ impl VmAirs {
             register,
             pages,
             memw_registers,
+            eqs,
+            bytewises,
+            stores,
+            cpu32s,
         }
     }
 }

prover/src/statement.rs

@@ -16,7 +16,7 @@ use crate::test_utils::E;
 use crate::{RuntimePageRange, TableCounts};
 
 /// Domain-separation tag. Bump the suffix (`_V2`, ...) on any encoding change.
-const DOMAIN_TAG: &[u8] = b"LAMBDAVM_STARK_STATEMENT_V1";
+const DOMAIN_TAG: &[u8] = b"LAMBDAVM_STARK_STATEMENT_V2";
 
 fn elf_digest(elf: &[u8]) -> [u8; 32] {
     let mut h = Keccak256::new();
@@ -55,6 +55,10 @@ pub(crate) fn absorb_statement(
         shift,
         branch,
         memw_register,
+        eq,
+        bytewise,
+        store,
+        cpu32,
     } = table_counts;
     for count in [
         cpu,
@@ -67,6 +71,10 @@ pub(crate) fn absorb_statement(
         shift,
         branch,
         memw_register,
+        eq,
+        bytewise,
+        store,
+        cpu32,
     ] {
         t.append_bytes(&(count as u64).to_le_bytes());
     }

prover/src/tables/bitwise.rs

@@ -38,7 +38,7 @@ use stark::trace::{TraceTable, columns2rows};
 #[cfg(feature = "parallel")]
 use rayon::prelude::*;
 
-use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField};
+use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, alu_op};
 
 // =========================================================================
 // Column indices for BITWISE table
@@ -92,8 +92,14 @@ pub mod cols {
     pub const MU_IS_B20: usize = 19;
     /// Multiplicity for HWSL lookups
     pub const MU_HWSL: usize = 20;
+    /// Multiplicity for `BYTE_ALU[opsel=AND]` lookups
+    pub const MU_BYTE_ALU_AND: usize = 21;
+    /// Multiplicity for `BYTE_ALU[opsel=OR]` lookups
+    pub const MU_BYTE_ALU_OR: usize = 22;
+    /// Multiplicity for `BYTE_ALU[opsel=XOR]` lookups
+    pub const MU_BYTE_ALU_XOR: usize = 23;
     /// Total number of columns
-    pub const NUM_COLUMNS: usize = 21;
+    pub const NUM_COLUMNS: usize = 24;
 }
 
 /// Number of rows in the BITWISE table: 256 * 256 * 16 = 2^20
@@ -442,6 +448,9 @@ pub fn update_multiplicities(
             BitwiseOperationType::IsHalf => cols::MU_IS_HALF,
             BitwiseOperationType::IsB20 => cols::MU_IS_B20,
             BitwiseOperationType::Hwsl => cols::MU_HWSL,
+            BitwiseOperationType::ByteAluAnd => cols::MU_BYTE_ALU_AND,
+            BitwiseOperationType::ByteAluOr => cols::MU_BYTE_ALU_OR,
+            BitwiseOperationType::ByteAluXor => cols::MU_BYTE_ALU_XOR,
         };
 
         // Increment multiplicity
@@ -477,8 +486,10 @@ pub(crate) fn trim_zero_rows(
     let kept_rows: Vec<usize> = (0..num_rows)
         .filter(|&row| {
             let row_data = trace.main_table.get_row(row);
-            // Check all multiplicity columns (indices 11-20)
-            (cols::MU_AND..=cols::MU_HWSL).any(|col| row_data[col] != FE::zero())
+            // Check all multiplicity columns (MU_AND..=MU_BYTE_ALU_XOR), including
+            // the BYTE_ALU columns (rows used only by a BYTE_ALU lookup
+            // must not be trimmed).
+            (cols::MU_AND..=cols::MU_BYTE_ALU_XOR).any(|col| row_data[col] != FE::zero())
         })
         .collect();
 
@@ -519,6 +530,10 @@ pub enum BitwiseOperationType {
     IsHalf,
     IsB20,
     Hwsl,
+    // Unified `BYTE_ALU` lookups, keyed by opsel.
+    ByteAluAnd,
+    ByteAluOr,
+    ByteAluXor,
 }
 
 /// A lookup request to the BITWISE precomputed table.
@@ -807,5 +822,67 @@ pub fn bus_interactions() -> Vec<BusInteraction> {
                 },
             ],
         ),
+        // BYTE_ALU[opsel, X, Y] -> out.
+        // Unifies AND/OR/XOR into one bus keyed by the `alu_op` descriptor.
+        // Implemented as one receiver per opsel, reusing the precomputed
+        // AND/OR/XOR result columns (the "single 2^20 column" in bitwise.typ is
+        // an optimization note, not a requirement).
+        BusInteraction::receiver(
+            BusId::ByteAlu,
+            Multiplicity::Column(cols::MU_BYTE_ALU_AND),
+            vec![
+                BusValue::constant(alu_op::AND as u64),
+                BusValue::Packed {
+                    start_column: cols::X,
+                    packing: Packing::Direct,
+                },
+                BusValue::Packed {
+                    start_column: cols::Y,
+                    packing: Packing::Direct,
+                },
+                BusValue::Packed {
+                    start_column: cols::AND,
+                    packing: Packing::Direct,
+                },
+            ],
+        ),
+        BusInteraction::receiver(
+            BusId::ByteAlu,
+            Multiplicity::Column(cols::MU_BYTE_ALU_OR),
+            vec![
+                BusValue::constant(alu_op::OR as u64),
+                BusValue::Packed {
+                    start_column: cols::X,
+                    packing: Packing::Direct,
+                },
+                BusValue::Packed {
+                    start_column: cols::Y,
+                    packing: Packing::Direct,
+                },
+                BusValue::Packed {
+                    start_column: cols::OR,
+                    packing: Packing::Direct,
+                },
+            ],
+        ),
+        BusInteraction::receiver(
+            BusId::ByteAlu,
+            Multiplicity::Column(cols::MU_BYTE_ALU_XOR),
+            vec![
+                BusValue::constant(alu_op::XOR as u64),
+                BusValue::Packed {
+                    start_column: cols::X,
+                    packing: Packing::Direct,
+                },
+                BusValue::Packed {
+                    start_column: cols::Y,
+                    packing: Packing::Direct,
+                },
+                BusValue::Packed {
+                    start_column: cols::XOR,
+                    packing: Packing::Direct,
+                },
+            ],
+        ),
     ]
 }

prover/src/tables/branch.rs

@@ -395,6 +395,8 @@ pub enum BranchConstraintKind {
     /// `(1 - JALR) * carry_1_pc * (1 - carry_1_pc) = 0`
     /// where carry_1_pc = (pc[1] + offset[1] + carry_0_pc - next_pc_unmasked[1]) / 2^32
     PcCarry1IsBit,
+    /// `IS_BIT<JALR>`: `JALR * (1 - JALR) = 0` (spec defense-in-depth assumption)
+    JalrIsBit,
     /// `JALR * carry_0_reg * (1 - carry_0_reg) = 0`
     /// where carry_0_reg = (register[0] + offset[0] - next_pc_unmasked[0]) / 2^32
     RegCarry0IsBit,
@@ -494,6 +496,7 @@ impl BranchConstraint {
         let one = FieldElement::<F>::one();
 
         match self.kind {
+            BranchConstraintKind::JalrIsBit => &jalr * (&one - &jalr),
             BranchConstraintKind::PcCarry0IsBit => {
                 let cond = &one - &jalr;
                 let c = Self::compute_carry_0_for(cols::PC_0, step);
@@ -520,8 +523,12 @@ impl BranchConstraint {
 
 impl TransitionConstraint<GoldilocksField, GoldilocksExtension> for BranchConstraint {
     fn degree(&self) -> usize {
-        // cond (degree 1) * carry (degree 1) * (1 - carry) (degree 1) = degree 3
-        3
+        match self.kind {
+            // JALR * (1 - JALR) = degree 2
+            BranchConstraintKind::JalrIsBit => 2,
+            // cond (degree 1) * carry (degree 1) * (1 - carry) (degree 1) = degree 3
+            _ => 3,
+        }
     }
 
     fn constraint_idx(&self) -> usize {
@@ -539,11 +546,13 @@ impl TransitionConstraint<GoldilocksField, GoldilocksExtension> for BranchConstr
 
 /// Creates all constraints for the BRANCH table.
 ///
-/// Returns 4 constraints (two conditional ADD templates × 2 carries each):
+/// Returns 5 constraints (two conditional ADD templates × 2 carries each, plus
+/// the `IS_BIT<JALR>` defense-in-depth assumption):
 /// - PcCarry0IsBit:  `(1 - JALR) * carry_0 * (1 - carry_0) = 0`  (pc path)
 /// - PcCarry1IsBit:  `(1 - JALR) * carry_1 * (1 - carry_1) = 0`  (pc path)
 /// - RegCarry0IsBit: `JALR * carry_0 * (1 - carry_0) = 0`        (register path)
 /// - RegCarry1IsBit: `JALR * carry_1 * (1 - carry_1) = 0`        (register path)
+/// - JalrIsBit:      `JALR * (1 - JALR) = 0`
 pub fn branch_constraints(constraint_idx_start: usize) -> (Vec<BranchConstraint>, usize) {
     let mut idx = constraint_idx_start;
     let mut next = || {
@@ -556,6 +565,7 @@ pub fn branch_constraints(constraint_idx_start: usize) -> (Vec<BranchConstraint>
         BranchConstraint::new(BranchConstraintKind::PcCarry1IsBit, next()),
         BranchConstraint::new(BranchConstraintKind::RegCarry0IsBit, next()),
         BranchConstraint::new(BranchConstraintKind::RegCarry1IsBit, next()),
+        BranchConstraint::new(BranchConstraintKind::JalrIsBit, next()),
     ];
     (constraints, idx)
 }