feat: simplify auxiliary trace building logic

processor/src/chiplets/aux_trace/main_trace.rs

@@ -1,7 +1,5 @@
-use core::ops::Range;
-
 use super::{ColMatrix, Felt};
-
+use core::ops::Range;
 use miden_air::trace::{
     chiplets::{
         hasher::STATE_WIDTH, BITWISE_A_COL_IDX, BITWISE_B_COL_IDX, BITWISE_OUTPUT_COL_IDX,
@@ -11,11 +9,17 @@ use miden_air::trace::{
     decoder::{HASHER_STATE_OFFSET, NUM_HASHER_COLUMNS, USER_OP_HELPERS_OFFSET},
     CHIPLETS_OFFSET, CLK_COL_IDX, CTX_COL_IDX, DECODER_TRACE_OFFSET, STACK_TRACE_OFFSET,
 };
-
 use vm_core::{utils::range, ONE, ZERO};
+
+// CONSTANTS
+// ================================================================================================
+
 const DECODER_HASHER_RANGE: Range<usize> =
     range(DECODER_TRACE_OFFSET + HASHER_STATE_OFFSET, NUM_HASHER_COLUMNS);
 
+// HELPER STRUCT AND METHODS
+// ================================================================================================
+
 pub struct MainTrace<'a> {
     columns: &'a ColMatrix<Felt>,
 }
@@ -27,7 +31,8 @@ impl<'a> MainTrace<'a> {
         }
     }
 
-    // System columns
+    // SYSTEM COLUMNS
+    // --------------------------------------------------------------------------------------------
 
     pub fn clk(&self, i: usize) -> Felt {
         self.columns.get_column(CLK_COL_IDX)[i]
@@ -37,12 +42,18 @@ impl<'a> MainTrace<'a> {
         self.columns.get_column(CTX_COL_IDX)[i]
     }
 
-    // Decoder columns
+    // DECODER COLUMNS
+    // --------------------------------------------------------------------------------------------
 
     pub fn addr(&self, i: usize) -> Felt {
         self.columns.get_column(DECODER_TRACE_OFFSET)[i]
     }
 
+    // Helper method to detect change of address.
+    pub fn is_addr_change(&self, i: usize) -> bool {
+        self.addr(i) != self.addr(i + 1)
+    }
+
     pub fn helper_0(&self, i: usize) -> Felt {
         self.columns.get_column(DECODER_TRACE_OFFSET + USER_OP_HELPERS_OFFSET)[i]
     }
@@ -82,7 +93,8 @@ impl<'a> MainTrace<'a> {
             + b6.mul_small(64)
     }
 
-    // Stack columns
+    // STACK COLUMNS
+    // --------------------------------------------------------------------------------------------
 
     pub fn s0(&self, i: usize) -> Felt {
         self.columns.get_column(STACK_TRACE_OFFSET)[i]
@@ -140,7 +152,8 @@ impl<'a> MainTrace<'a> {
         self.columns.get_column(STACK_TRACE_OFFSET + 13)[i]
     }
 
-    // Chiplets columns
+    // CHIPLETS COLUMNS
+    // --------------------------------------------------------------------------------------------
 
     pub fn chiplet_selector_0(&self, i: usize) -> Felt {
         self.columns.get_column(CHIPLETS_OFFSET)[i]
@@ -230,7 +243,9 @@ impl<'a> MainTrace<'a> {
     pub fn chiplet_kernel_root_3(&self, i: usize) -> Felt {
         self.columns.get_column(CHIPLETS_OFFSET + 9)[i]
     }
-    // Merkle path hashing selectors
+
+    //  MERKLE PATH HASHING SELECTORS
+    // --------------------------------------------------------------------------------------------
 
     pub fn f_mv(&self, i: usize) -> bool {
         (i % 8 == 0)
@@ -271,9 +286,4 @@ impl<'a> MainTrace<'a> {
             && self.chiplet_selector_2(i) == ONE
             && self.chiplet_selector_3(i) == ZERO
     }
-
-    // Helper method to detect change of address in the kernel ROM chiplet.
-    pub fn is_addr_change(&self, i: usize) -> bool {
-        self.addr(i) != self.addr(i + 1)
-    }
 }

processor/src/chiplets/aux_trace/mod.rs

@@ -1,7 +1,4 @@
-use super::{
-    trace::{AuxColumnBuilder, LookupTableRow},
-    BTreeMap, ColMatrix, Felt, FieldElement, StarkField, Vec, Word,
-};
+use super::{ColMatrix, Felt, FieldElement, StarkField, Vec};
 
 use miden_air::trace::chiplets::{
     hasher::{
@@ -11,15 +8,10 @@ use miden_air::trace::chiplets::{
     kernel_rom::KERNEL_PROC_LABEL,
     memory::{MEMORY_READ_LABEL, MEMORY_WRITE_LABEL},
 };
-#[cfg(test)]
-pub(crate) use virtual_table::{ChipletsVTableRow, };
 
 use vm_core::{utils::uninit_vector, Operation, ONE, ZERO};
 use winter_prover::math::batch_inversion;
 
-#[cfg(test)]
-mod virtual_table;
-
 mod main_trace;
 use main_trace::MainTrace;
 
@@ -956,10 +948,7 @@ fn get_op_label(s0: Felt, s1: Felt, s2: Felt, s3: Felt) -> Felt {
 fn addr_to_hash_cycle(addr: Felt) -> usize {
     let row = (addr.as_int() - 1) as usize;
     let cycle_row = row % HASH_CYCLE_LEN;
-    debug_assert!(
-        cycle_row == 0 || cycle_row == HASH_CYCLE_LEN - 1,
-        "invalid address for hasher lookup"
-    );
+    debug_assert!(cycle_row == 0 || cycle_row == HASH_CYCLE_LEN - 1, "invalid address for hasher");
 
     cycle_row
 }

processor/src/chiplets/bitwise/mod.rs

@@ -1,10 +1,7 @@
-use super::{
-    trace::LookupTableRow, utils::get_trace_len, ColMatrix, ExecutionError, Felt,
-    FieldElement, StarkField, TraceFragment, Vec, BITWISE_AND_LABEL, BITWISE_XOR_LABEL, ZERO,
-};
+use super::{utils::get_trace_len, ExecutionError, Felt, StarkField, TraceFragment, Vec, ZERO};
 use miden_air::trace::chiplets::bitwise::{
-    A_COL_IDX, A_COL_RANGE, BITWISE_AND, BITWISE_XOR, B_COL_IDX, B_COL_RANGE, OP_CYCLE_LEN,
-    OUTPUT_COL_IDX, PREV_OUTPUT_COL_IDX, TRACE_WIDTH,
+    A_COL_IDX, A_COL_RANGE, BITWISE_AND, BITWISE_XOR, B_COL_IDX, B_COL_RANGE, OUTPUT_COL_IDX,
+    PREV_OUTPUT_COL_IDX, TRACE_WIDTH,
 };
 
 #[cfg(test)]
@@ -151,10 +148,7 @@ impl Bitwise {
     // --------------------------------------------------------------------------------------------
 
     /// Fills the provided trace fragment with trace data from this bitwise helper instance.
-    pub fn fill_trace(
-        self,
-        trace: &mut TraceFragment,
-    ) {
+    pub fn fill_trace(self, trace: &mut TraceFragment) {
         // make sure fragment dimensions are consistent with the dimensions of this trace
         debug_assert_eq!(self.trace_len(), trace.len(), "inconsistent trace lengths");
         debug_assert_eq!(TRACE_WIDTH, trace.width(), "inconsistent trace widths");
@@ -211,36 +205,3 @@ pub fn assert_u32(value: Felt) -> Result<Felt, ExecutionError> {
         Ok(value)
     }
 }
-
-// BITWISE LOOKUPS
-// ================================================================================================
-#[derive(Debug, Copy, Clone, PartialEq, Eq)]
-pub struct BitwiseLookup {
-    // unique label identifying the bitwise operation
-    label: Felt,
-    a: Felt,
-    b: Felt,
-    z: Felt,
-}
-
-impl BitwiseLookup {
-    pub fn new(label: Felt, a: Felt, b: Felt, z: Felt) -> Self {
-        Self { label, a, b, z }
-    }
-}
-
-impl LookupTableRow for BitwiseLookup {
-    /// Reduces this row to a single field element in the field specified by E. This requires
-    /// at least 5 alpha values.
-    fn to_value<E: FieldElement<BaseField = Felt>>(
-        &self,
-        _main_trace: &ColMatrix<Felt>,
-        alphas: &[E],
-    ) -> E {
-        alphas[0]
-            + alphas[1].mul_base(self.label)
-            + alphas[2].mul_base(self.a)
-            + alphas[3].mul_base(self.b)
-            + alphas[4].mul_base(self.z)
-    }
-}

processor/src/chiplets/bitwise/tests.rs

@@ -1,7 +1,7 @@
-use super::{
-    Bitwise, Felt, StarkField, TraceFragment, Vec, A_COL_IDX,
-    A_COL_RANGE, BITWISE_AND, BITWISE_AND_LABEL, BITWISE_XOR, BITWISE_XOR_LABEL, B_COL_IDX,
-    B_COL_RANGE, OP_CYCLE_LEN, OUTPUT_COL_IDX, PREV_OUTPUT_COL_IDX, TRACE_WIDTH,
+use super::{Bitwise, Felt, StarkField, TraceFragment, Vec};
+use miden_air::trace::chiplets::bitwise::{
+    A_COL_IDX, A_COL_RANGE, BITWISE_AND, BITWISE_XOR, B_COL_IDX, B_COL_RANGE, OP_CYCLE_LEN,
+    OUTPUT_COL_IDX, PREV_OUTPUT_COL_IDX, TRACE_WIDTH,
 };
 use test_utils::rand::rand_value;
 use vm_core::ZERO;
@@ -184,10 +184,10 @@ fn bitwise_multiple() {
 // ================================================================================================
 
 /// Builds a trace of the specified length and fills it with data from the provided Bitwise instance.
-fn build_trace(bitwise: Bitwise, num_rows: usize) ->  Vec<Vec<Felt>> {
+fn build_trace(bitwise: Bitwise, num_rows: usize) -> Vec<Vec<Felt>> {
     let mut trace = (0..TRACE_WIDTH).map(|_| vec![ZERO; num_rows]).collect::<Vec<_>>();
     let mut fragment = TraceFragment::trace_to_fragment(&mut trace);
-    bitwise.fill_trace(&mut fragment );
+    bitwise.fill_trace(&mut fragment);
 
     trace
 }
@@ -228,4 +228,3 @@ fn rand_u32() -> Felt {
     let value = rand_value::<u64>() as u32 as u64;
     Felt::new(value)
 }
-