Trace-gen cleanup: single MU-column source, simpler push_reg_access, tighter visibility, stale docs

prover/src/tables/bitwise.rs

@@ -383,8 +383,9 @@ pub fn generate_bitwise_trace() -> TraceTable<GoldilocksField, GoldilocksExtensi
                 table.set_half(row_idx, cols::SLL, sll as u16);
                 table.set_half(row_idx, cols::SLLC, sllc as u16);
 
-                // Multiplicity columns start at zero
-                // They will be updated by update_multiplicities()
+                // Multiplicity columns start at zero. They are filled by
+                // `BitwiseHistogram::fill_multiplicities`; `update_multiplicities`
+                // only tops up the continuation L2G lookups afterward.
             }
         }
     }
@@ -423,13 +424,13 @@ pub fn update_multiplicities(
 /// Number of distinct BITWISE lookup types (one multiplicity column each).
 /// Derived from [`BitwiseOperationType::ALL`], which the compile-time guard
 /// below keeps in lockstep with [`lookup_type_index`].
-pub const NUM_LOOKUP_TYPES: usize = BitwiseOperationType::ALL.len();
+pub(crate) const NUM_LOOKUP_TYPES: usize = BitwiseOperationType::ALL.len();
 
 /// Dense index in `[0, NUM_LOOKUP_TYPES)` for a lookup type. Ordering is an
 /// internal detail of the histogram; [`BitwiseOperationType::ALL`] is its
 /// inverse, enforced at compile time.
 #[inline]
-pub const fn lookup_type_index(t: BitwiseOperationType) -> usize {
+pub(crate) const fn lookup_type_index(t: BitwiseOperationType) -> usize {
     match t {
         BitwiseOperationType::Msb8 => 0,
         BitwiseOperationType::Msb16 => 1,
@@ -444,45 +445,64 @@ pub const fn lookup_type_index(t: BitwiseOperationType) -> usize {
     }
 }
 
+/// The MU_* multiplicity column for each lookup type, in [`lookup_type_index`]
+/// order. This is the single source of truth for the type→column mapping: both
+/// the per-op path ([`mu_column`]) and the histogram fill ([`type_mu_column`])
+/// index into this one array. The compile-time block below checks the entries
+/// are pairwise distinct, so a duplicate column is a build error rather than a
+/// silent overwrite in [`BitwiseHistogram::fill_multiplicities`].
+const MU_COLUMNS: [usize; NUM_LOOKUP_TYPES] = [
+    cols::MU_MSB8,         // Msb8
+    cols::MU_MSB16,        // Msb16
+    cols::MU_ZERO,         // Zero
+    cols::MU_ARE_BYTES,    // AreBytes
+    cols::MU_IS_HALF,      // IsHalf
+    cols::MU_IS_B20,       // IsB20
+    cols::MU_HWSL,         // Hwsl
+    cols::MU_BYTE_ALU_AND, // ByteAluAnd
+    cols::MU_BYTE_ALU_OR,  // ByteAluOr
+    cols::MU_BYTE_ALU_XOR, // ByteAluXor
+];
+
 /// Multiplicity column for a lookup type. Used by the per-op path
 /// ([`update_multiplicities`]), which is still live production code: continuation
 /// epochs add their L2G lookups through it on top of the histogram-filled trace.
 #[inline]
-pub const fn mu_column(t: BitwiseOperationType) -> usize {
-    match t {
-        BitwiseOperationType::Msb8 => cols::MU_MSB8,
-        BitwiseOperationType::Msb16 => cols::MU_MSB16,
-        BitwiseOperationType::Zero => cols::MU_ZERO,
-        BitwiseOperationType::AreBytes => cols::MU_ARE_BYTES,
-        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,
-    }
+pub(crate) const fn mu_column(t: BitwiseOperationType) -> usize {
+    MU_COLUMNS[lookup_type_index(t)]
 }
 
 /// Multiplicity column for the histogram lane at dense index `type_idx`
 /// (inverse of [`lookup_type_index`]). Used by [`BitwiseHistogram::fill_multiplicities`].
 ///
-/// Defined as `mu_column ∘ ALL`, so it cannot disagree with the authoritative
-/// per-type match — no assumption about MU column contiguity or ordering.
+/// Reads directly from [`MU_COLUMNS`], the single type→column source of truth.
 #[inline]
 const fn type_mu_column(type_idx: usize) -> usize {
-    mu_column(BitwiseOperationType::ALL[type_idx])
+    MU_COLUMNS[type_idx]
 }
 
-// Compile-time guard: `ALL` must list every lookup type exactly once, in
-// `lookup_type_index` order (i.e. it is the exact inverse of that mapping).
-// Adding a variant forces the `lookup_type_index` match to be extended
-// (exhaustiveness), and this assert then forces `ALL` — and with it
-// `NUM_LOOKUP_TYPES` — to follow. A mismatch is a compile error, not a test
-// failure: a wrong MU column would silently unbalance the BITWISE bus.
+// Compile-time guards on the type↔column bookkeeping.
+//
+// 1. `ALL` must list every lookup type exactly once, in `lookup_type_index`
+//    order (i.e. it is the exact inverse of that mapping). Adding a variant
+//    forces the `lookup_type_index` match to be extended (exhaustiveness), and
+//    this assert then forces `ALL` — and with it `NUM_LOOKUP_TYPES` — to follow.
+// 2. The type→column map is now derived from the single `MU_COLUMNS` array, and
+//    its entries are checked pairwise distinct (injective). A wrong or duplicated
+//    MU column would silently unbalance the BITWISE bus, so both are compile
+//    errors, not test failures.
 const _: () = {
     let mut i = 0;
     while i < NUM_LOOKUP_TYPES {
         assert!(lookup_type_index(BitwiseOperationType::ALL[i]) == i);
+        let mut j = i + 1;
+        while j < NUM_LOOKUP_TYPES {
+            assert!(
+                MU_COLUMNS[i] != MU_COLUMNS[j],
+                "MU_COLUMNS entries must map distinct lookup types to distinct columns"
+            );
+            j += 1;
+        }
         i += 1;
     }
 };
@@ -499,7 +519,7 @@ const _: () = {
 /// [`update_multiplicities`] produces (both just sum the same lookups per cell).
 ///
 /// Memory: `NUM_ROWS * NUM_LOOKUP_TYPES * 8` bytes = 2^20 * 10 * 8 = 80 MiB.
-pub struct BitwiseHistogram {
+pub(crate) struct BitwiseHistogram {
     counters: Box<[u64]>,
 }
 
@@ -508,22 +528,22 @@ impl BitwiseHistogram {
     // No `Default` impl on purpose: `new()` allocates 80 MiB, so a stray
     // `..Default::default()` / `#[derive(Default)]` must not silently do that.
     #[allow(clippy::new_without_default)]
-    pub fn new() -> Self {
+    pub(crate) fn new() -> Self {
         Self {
             counters: vec![0u64; NUM_ROWS * NUM_LOOKUP_TYPES].into_boxed_slice(),
         }
     }
 
     /// Increment the counter for one lookup.
     #[inline]
-    pub fn bump(&mut self, op: BitwiseOperation) {
+    pub(crate) fn bump(&mut self, op: BitwiseOperation) {
         self.bump_n(op, 1);
     }
 
     /// Add `n` occurrences of one lookup in a single step (e.g. CPU padding rows,
     /// which all send identical all-zero lookups).
     #[inline]
-    pub fn bump_n(&mut self, op: BitwiseOperation, n: u64) {
+    pub(crate) fn bump_n(&mut self, op: BitwiseOperation, n: u64) {
         let idx = lookup_type_index(op.lookup_type) * NUM_ROWS + row_index(op.x, op.y, op.z);
         // (x, y) are u8, and row_index debug-asserts z < 16, so in debug builds a
         // corrupt op fails loudly here. In release an out-of-domain z would NOT
@@ -536,14 +556,14 @@ impl BitwiseHistogram {
 
     /// Fold a slice of lookups into the histogram.
     #[inline]
-    pub fn add_ops(&mut self, ops: &[BitwiseOperation]) {
+    pub(crate) fn add_ops(&mut self, ops: &[BitwiseOperation]) {
         for &op in ops {
             self.bump(op);
         }
     }
 
     /// Merge another histogram into this one (commutative, order-independent).
-    pub fn merge(&mut self, other: &BitwiseHistogram) {
+    pub(crate) fn merge(&mut self, other: &BitwiseHistogram) {
         for (a, b) in self.counters.iter_mut().zip(other.counters.iter()) {
             *a += *b;
         }
@@ -558,7 +578,7 @@ impl BitwiseHistogram {
     /// Callers that layer additional lookups on top (continuation epochs add
     /// their L2G lookups via `update_multiplicities`, which increments) must do
     /// so strictly AFTER this fill, never before.
-    pub fn fill_multiplicities(
+    pub(crate) fn fill_multiplicities(
         &self,
         trace: &mut TraceTable<GoldilocksField, GoldilocksExtension>,
     ) {
@@ -594,7 +614,7 @@ impl BitwiseOperationType {
     /// Every lookup type exactly once, in [`lookup_type_index`] order (the
     /// compile-time guard next to [`type_mu_column`] enforces this). The array
     /// length is the single origin of [`NUM_LOOKUP_TYPES`].
-    pub const ALL: [Self; 10] = [
+    pub(crate) const ALL: [Self; 10] = [
         Self::Msb8,
         Self::Msb16,
         Self::Zero,

prover/src/tables/memw_register.rs

@@ -193,7 +193,11 @@ impl RegRow {
 ///
 /// Thin wrapper over [`generate_memw_register_trace_from_rows`] (via
 /// [`RegRow::from_memw`]) so there is exactly one MEMW_R column-write sequence.
-pub fn generate_memw_register_trace(
+///
+/// Test-only: production code fills MEMW_R directly from [`RegRow`]s, so the walk
+/// never routes through this `MemwOperation`-based entry point.
+#[cfg(test)]
+pub(crate) fn generate_memw_register_trace(
     operations: &[MemwOperation],
 ) -> TraceTable<GoldilocksField, GoldilocksExtension> {
     let rows: Vec<RegRow> = operations.iter().map(RegRow::from_memw).collect();
@@ -202,7 +206,7 @@ pub fn generate_memw_register_trace(
 
 /// The MEMW_R column fill from compact [`RegRow`]s. This is the single source of
 /// truth for the MEMW_R trace layout; both the walk's direct fast path and the
-/// `MemwOperation`-based [`generate_memw_register_trace`] land here.
+/// `MemwOperation`-based `generate_memw_register_trace` test wrapper land here.
 pub(crate) fn generate_memw_register_trace_from_rows(
     rows: &[RegRow],
 ) -> TraceTable<GoldilocksField, GoldilocksExtension> {

prover/src/tables/trace_builder.rs

@@ -433,7 +433,8 @@ impl MemwBuckets {
 }
 
 /// Sink for `MemwOperation`s so `collect_register_ops_from_cpu` can feed either a plain
-/// `Vec` (tests/scratch paths) or the classifying [`MemwBuckets`] (the walk).
+/// `Vec` (the `count_table_lengths` trace-sizing pass) or the classifying
+/// [`MemwBuckets`] (the walk).
 trait MemwSink {
     fn push_op(&mut self, op: MemwOperation);
 
@@ -443,28 +444,25 @@ trait MemwSink {
     /// (via the same predicate as `is_register_op`): if the timestamp delta admits the
     /// op into MEMW_R it fills a compact [`RegRow`] DIRECTLY — no `MemwOperation` is
     /// built. Only on the (rare) fallback (delta out of IS_HALF range, or upper-limb
-    /// mismatch) does it call `build_fallback` to materialize the `MemwOperation` and
+    /// mismatch) does it build the `MemwOperation` (via [`build_reg_fallback`]) and
     /// route it to the aligned/general bucket exactly as before.
     ///
-    /// `reg_addr` is `2 * reg_index`; `[val0,val1]`/`[old0,old1]` are the 32-bit halves;
-    /// `old_ts` is the (shared) old_timestamp of both words.
+    /// `reg_addr` is `2 * reg_index`; `val`/`old` are the two 32-bit halves of the new
+    /// and previous register words; `old_ts` is the (shared) old_timestamp of both words.
     #[inline]
-    #[allow(clippy::too_many_arguments)]
     fn push_reg_access(
         &mut self,
         reg_addr: u64,
-        val0: u32,
-        val1: u32,
-        old0: u32,
-        old1: u32,
+        val: [u32; 2],
+        old: [u32; 2],
         timestamp: u64,
         old_ts: u64,
         is_read: bool,
-        build_fallback: impl FnOnce() -> MemwOperation,
     ) {
         // Default impl (plain Vec): register accesses are still ordinary MemwOperations.
-        let _ = (reg_addr, val0, val1, old0, old1, timestamp, old_ts, is_read);
-        self.push_op(build_fallback());
+        self.push_op(build_reg_fallback(
+            reg_addr, val, old, timestamp, old_ts, is_read,
+        ));
     }
 }
 impl MemwSink for Vec<MemwOperation> {
@@ -480,34 +478,49 @@ impl MemwSink for MemwBuckets {
     }
 
     #[inline]
-    #[allow(clippy::too_many_arguments)]
     fn push_reg_access(
         &mut self,
         reg_addr: u64,
-        val0: u32,
-        val1: u32,
-        old0: u32,
-        old1: u32,
+        val: [u32; 2],
+        old: [u32; 2],
         timestamp: u64,
         old_ts: u64,
         is_read: bool,
-        build_fallback: impl FnOnce() -> MemwOperation,
     ) {
         // Mirror `is_register_op` for a width-2 register access whose two words share
         // `old_ts` (always true here by construction). If it passes, fill a RegRow
         // directly; otherwise fall back to the general/aligned MemwOperation path.
         if reg_ts_delta_in_range(timestamp, old_ts) {
             self.register_rows.push(RegRow::new(
-                reg_addr, timestamp, val0, val1, old0, old1, old_ts, is_read,
+                reg_addr, timestamp, val[0], val[1], old[0], old[1], old_ts, is_read,
             ));
         } else {
-            let op = build_fallback();
+            let op = build_reg_fallback(reg_addr, val, old, timestamp, old_ts, is_read);
             debug_assert!(!is_register_op(&op), "reg fallback must not be MEMW_R");
             self.push(op);
         }
     }
 }
 
+/// Materialize the aligned/general `MemwOperation` for a register access that does
+/// NOT fit the MEMW_R fast path. Register values pack as `[lo, hi, 0, …]` (see
+/// [`pack_register_value`]) and both words share `old_ts`, so this rebuilds exactly
+/// the op the fast-path callers would otherwise have routed to the buckets.
+fn build_reg_fallback(
+    reg_addr: u64,
+    val: [u32; 2],
+    old: [u32; 2],
+    timestamp: u64,
+    old_ts: u64,
+    is_read: bool,
+) -> MemwOperation {
+    let value = [val[0], val[1], 0, 0, 0, 0, 0, 0];
+    let old_value = [old[0], old[1], 0, 0, 0, 0, 0, 0];
+    let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0];
+    MemwOperation::new(true, reg_addr, value, timestamp, 2, is_read)
+        .with_old(old_value, old_timestamps)
+}
+
 /// Collects all derived operations from CPU operations in a single pass.
 ///
 /// This includes:
@@ -957,22 +970,16 @@ fn collect_register_ops_from_cpu<S: MemwSink>(
             register_state.read(d.rs1)
         };
         let ts = op.timestamp;
-        // Direct fast path: fill a RegRow when routing to MEMW_R; the closure rebuilds
-        // the identical MemwOperation only on the (rare) general/aligned fallback.
+        // Direct fast path: fill a RegRow when routing to MEMW_R; push_reg_access
+        // rebuilds the identical MemwOperation only on the (rare) general/aligned
+        // fallback. Reads leave the value unchanged, so old == new here.
         memw_ops.push_reg_access(
             reg_addr,
-            reg_value[0],
-            reg_value[1],
-            reg_value[0],
-            reg_value[1],
+            [reg_value[0], reg_value[1]],
+            [reg_value[0], reg_value[1]],
             ts,
             old_ts,
             true,
-            || {
-                let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0];
-                MemwOperation::new(true, reg_addr, reg_value, ts, 2, true)
-                    .with_old(reg_value, old_timestamps)
-            },
         );
         if d.rs1 == 255 {
             register_state.write_pc(op.rv1, op.timestamp);
@@ -989,18 +996,11 @@ fn collect_register_ops_from_cpu<S: MemwSink>(
         let ts = op.timestamp + 1;
         memw_ops.push_reg_access(
             reg_addr,
-            reg_value[0],
-            reg_value[1],
-            reg_value[0],
-            reg_value[1],
+            [reg_value[0], reg_value[1]],
+            [reg_value[0], reg_value[1]],
             ts,
             old_ts,
             true,
-            || {
-                let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0];
-                MemwOperation::new(true, reg_addr, reg_value, ts, 2, true)
-                    .with_old(reg_value, old_timestamps)
-            },
         );
         register_state.write(d.rs2, op.rv2, op.timestamp + 1);
     }
@@ -1014,18 +1014,11 @@ fn collect_register_ops_from_cpu<S: MemwSink>(
         let ts = op.timestamp + 2;
         memw_ops.push_reg_access(
             reg_addr,
-            reg_value[0],
-            reg_value[1],
-            old_value[0],
-            old_value[1],
+            [reg_value[0], reg_value[1]],
+            [old_value[0], old_value[1]],
             ts,
             old_ts,
             false,
-            || {
-                let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0];
-                MemwOperation::new(true, reg_addr, reg_value, ts, 2, false)
-                    .with_old(old_value, old_timestamps)
-            },
         );
         register_state.write(d.rd, op.rvd, op.timestamp + 2);
     }