refactor: clean up aux column generator for stack overflow table (#1213)

processor/src/stack/aux_trace.rs

@@ -29,67 +29,46 @@ impl AuxTraceBuilder {
 }
 
 impl<E: FieldElement<BaseField = Felt>> AuxColumnBuilder<E> for AuxTraceBuilder {
+    /// Initializes the overflow stack auxiliary column.
     fn init_responses(&self, _main_trace: &MainTrace, alphas: &[E]) -> E {
-        init_overflow_table(self, alphas)
+        let mut initial_column_value = E::ONE;
+        for row in self.overflow_table_rows.iter().take(self.num_init_rows) {
+            let value = (*row).to_value(alphas);
+            initial_column_value *= value;
+        }
+        initial_column_value
     }
-    fn get_requests_at(&self, main_trace: &MainTrace, alphas: &[E], i: usize) -> E {
-        stack_overflow_table_removals(main_trace, alphas, i)
-    }
-
-    fn get_responses_at(&self, main_trace: &MainTrace, alphas: &[E], i: usize) -> E {
-        stack_overflow_table_inclusions(main_trace, alphas, i)
-    }
-}
 
-/// Initializes the overflow stack auxiliary column.
-fn init_overflow_table<E>(overflow_table: &AuxTraceBuilder, alphas: &[E]) -> E
-where
-    E: FieldElement<BaseField = Felt>,
-{
-    let mut initial_column_value = E::ONE;
-    for row in overflow_table.overflow_table_rows.iter().take(overflow_table.num_init_rows) {
-        let value = (*row).to_value(alphas);
-        initial_column_value *= value;
-    }
-    initial_column_value
-}
-
-/// Adds a row to the stack overflow table.
-fn stack_overflow_table_inclusions<E>(main_trace: &MainTrace, alphas: &[E], i: usize) -> E
-where
-    E: FieldElement<BaseField = Felt>,
-{
-    let is_right_shift = main_trace.is_right_shift(i);
+    /// Removes a row from the stack overflow table.
+    fn get_requests_at(&self, main_trace: &MainTrace, alphas: &[E], i: usize) -> E {
+        let is_left_shift = main_trace.is_left_shift(i);
+        let is_non_empty_overflow = main_trace.is_non_empty_overflow(i);
 
-    if is_right_shift {
-        let k0 = main_trace.clk(i);
-        let s15 = main_trace.stack_element(15, i);
-        let b1 = main_trace.parent_overflow_address(i);
+        if is_left_shift && is_non_empty_overflow {
+            let b1 = main_trace.parent_overflow_address(i);
+            let s15_prime = main_trace.stack_element(15, i + 1);
+            let b1_prime = main_trace.parent_overflow_address(i + 1);
 
-        alphas[0] + alphas[1].mul_base(k0) + alphas[2].mul_base(s15) + alphas[3].mul_base(b1)
-    } else {
-        E::ONE
+            let row = OverflowTableRow::new(b1, s15_prime, b1_prime);
+            row.to_value(alphas)
+        } else {
+            E::ONE
+        }
     }
-}
 
-/// Removes a row from the stack overflow table.
-fn stack_overflow_table_removals<E>(main_trace: &MainTrace, alphas: &[E], i: usize) -> E
-where
-    E: FieldElement<BaseField = Felt>,
-{
-    let is_left_shift = main_trace.is_left_shift(i);
-    let is_non_empty_overflow = main_trace.is_non_empty_overflow(i);
+    /// Adds a row to the stack overflow table.
+    fn get_responses_at(&self, main_trace: &MainTrace, alphas: &[E], i: usize) -> E {
+        let is_right_shift = main_trace.is_right_shift(i);
 
-    if is_left_shift && is_non_empty_overflow {
-        let b1 = main_trace.parent_overflow_address(i);
-        let s15_prime = main_trace.stack_element(15, i + 1);
-        let b1_prime = main_trace.parent_overflow_address(i + 1);
+        if is_right_shift {
+            let k0 = main_trace.clk(i);
+            let s15 = main_trace.stack_element(15, i);
+            let b1 = main_trace.parent_overflow_address(i);
 
-        alphas[0]
-            + alphas[1].mul_base(b1)
-            + alphas[2].mul_base(s15_prime)
-            + alphas[3].mul_base(b1_prime)
-    } else {
-        E::ONE
+            let row = OverflowTableRow::new(k0, s15, b1);
+            row.to_value(alphas)
+        } else {
+            E::ONE
+        }
     }
 }

processor/src/stack/mod.rs

@@ -240,7 +240,7 @@ impl Stack {
 
         // Update the overflow table.
         let to_overflow = self.trace.get_stack_value_at(self.clk, MAX_TOP_IDX);
-        self.overflow.push(to_overflow, self.clk as u64);
+        self.overflow.push(to_overflow, Felt::from(self.clk));
 
         // Stack depth always increases on right shift.
         self.active_depth += 1;

processor/src/stack/overflow.rs

@@ -16,9 +16,6 @@ pub struct OverflowTable {
     /// A list of indices into the `all_rows` vector which describes the rows currently in the
     /// overflow table.
     active_rows: Vec<usize>,
-    /// A list of updates made to the overflow table during program execution. For each update we
-    /// also track the cycle at which the update happened.
-    update_trace: Vec<(u64, OverflowTableUpdate)>,
     /// A map which records the full state of the overflow table at every cycle during which an
     /// update happened. This map is populated only when `trace_enabled` = true.
     trace: BTreeMap<u64, Vec<Felt>>,
@@ -42,7 +39,6 @@ impl OverflowTable {
         Self {
             all_rows: Vec::new(),
             active_rows: Vec::new(),
-            update_trace: Vec::new(),
             trace: BTreeMap::new(),
             trace_enabled: enable_trace,
             num_init_rows: 0,
@@ -62,7 +58,7 @@ impl OverflowTable {
 
         let mut clk = Felt::MODULUS - init_values.len() as u64;
         for &val in init_values.iter().rev() {
-            overflow_table.push(val, clk);
+            overflow_table.push(val, Felt::new(clk));
             clk += 1;
         }
 
@@ -75,12 +71,12 @@ impl OverflowTable {
     /// Pushes the specified value into the overflow table.
     ///
     /// Parameter clk specifies the clock cycle at which the value is added to the table.
-    pub fn push(&mut self, value: Felt, clk: u64) {
+    pub fn push(&mut self, value: Felt, clk: Felt) {
         // ZERO address indicates that the overflow table is empty, and thus, no actual value
         // should be inserted into the table with this address. This is not a problem since for
         // every real program, we first execute an operation marking the start of a code block,
         // and thus, no operation can shift the stack to the right at clk = 0.
-        debug_assert_ne!(clk, 0, "cannot add value to overflow at clk=0");
+        debug_assert_ne!(clk, ZERO, "cannot add value to overflow at clk=0");
 
         // create and record the new row, and also put it at the top of the overflow table
         let row_idx = self.all_rows.len() as u32;
@@ -89,14 +85,11 @@ impl OverflowTable {
         self.active_rows.push(row_idx as usize);
 
         // set the last row address to the address of the newly added row
-        self.last_row_addr = Felt::from(clk);
-
-        // mark this clock cycle as the cycle at which a new row was inserted into the table
-        self.update_trace.push((clk, OverflowTableUpdate::RowInserted(row_idx)));
+        self.last_row_addr = clk;
 
         if self.trace_enabled {
             // insert a copy of the current table state into the trace
-            self.save_current_state(clk);
+            self.save_current_state(clk.as_int());
         }
     }
 
@@ -119,10 +112,6 @@ impl OverflowTable {
         let removed_value = last_row.val;
         self.last_row_addr = last_row.prev;
 
-        // mark this clock cycle as the clock cycle at which a row was removed from the table
-        self.update_trace
-            .push((clk, OverflowTableUpdate::RowRemoved(last_row_idx as u32)));
-
         if self.trace_enabled {
             // insert a copy of the current table state into the trace
             self.save_current_state(clk);
@@ -252,12 +241,8 @@ pub struct OverflowTableRow {
 }
 
 impl OverflowTableRow {
-    pub fn new(clk: u64, val: Felt, prev: Felt) -> Self {
-        Self {
-            val,
-            clk: Felt::from(clk),
-            prev,
-        }
+    pub fn new(clk: Felt, val: Felt, prev: Felt) -> Self {
+        Self { val, clk, prev }
     }
 }
 
@@ -271,17 +256,3 @@ impl OverflowTableRow {
             + alphas[3].mul_base(self.prev)
     }
 }
-
-// OVERFLOW TABLE UPDATES
-// ================================================================================================
-
-/// Describes an update to the stack overflow table. There could be two types of updates:
-/// - A single row can be added to the table. This happens during a right shift.
-/// - A single row can be removed from the table. This happens during a left shift.
-///
-/// For each update we also record the index of the row that was added/removed from the table.
-#[derive(Debug, Copy, Clone, PartialEq, Eq)]
-pub enum OverflowTableUpdate {
-    RowInserted(u32),
-    RowRemoved(u32),
-}

processor/src/stack/tests.rs

@@ -53,9 +53,9 @@ fn initialize_overflow() {
     ];
     let init_addr = Felt::MODULUS - 3;
     let expected_overflow_rows = vec![
-        OverflowTableRow::new(init_addr, ONE, ZERO),
-        OverflowTableRow::new(init_addr + 1, Felt::new(2), Felt::new(init_addr)),
-        OverflowTableRow::new(init_addr + 2, Felt::new(3), Felt::new(init_addr + 1)),
+        OverflowTableRow::new(Felt::new(init_addr), ONE, ZERO),
+        OverflowTableRow::new(Felt::new(init_addr + 1), Felt::new(2), Felt::new(init_addr)),
+        OverflowTableRow::new(Felt::new(init_addr + 2), Felt::new(3), Felt::new(init_addr + 1)),
     ];
     let expected_overflow_active_rows = vec![0, 1, 2];
 

processor/src/trace/tests/stack.rs

@@ -37,10 +37,10 @@ fn p1_trace() {
     let p1 = aux_columns.get_column(P1_COL_IDX);
 
     let row_values = [
-        OverflowTableRow::new(2, ONE, ZERO).to_value(&alphas),
-        OverflowTableRow::new(3, TWO, TWO).to_value(&alphas),
-        OverflowTableRow::new(6, TWO, TWO).to_value(&alphas),
-        OverflowTableRow::new(10, ZERO, ZERO).to_value(&alphas),
+        OverflowTableRow::new(Felt::new(2), ONE, ZERO).to_value(&alphas),
+        OverflowTableRow::new(Felt::new(3), TWO, TWO).to_value(&alphas),
+        OverflowTableRow::new(Felt::new(6), TWO, TWO).to_value(&alphas),
+        OverflowTableRow::new(Felt::new(10), ZERO, ZERO).to_value(&alphas),
     ];
 
     // make sure the first entry is ONE