fix first column assertions

air/src/constraints/logup.rs

@@ -1,25 +1,42 @@
 use alloc::vec::Vec;
-use vm_core::{ExtensionOf, Felt, FieldElement};
+use vm_core::{stack::STACK_TOP_SIZE, ExtensionOf, Felt, FieldElement};
 use winter_air::{
     Assertion, EvaluationFrame, LagrangeKernelRandElements, TransitionConstraintDegree,
 };
 
 use crate::{
     gkr_proof::inner_product,
-    trace::logup::{LAGRANGE_KERNEL_COL_IDX, S_COL_IDX},
+    trace::{
+        logup::{LAGRANGE_KERNEL_COL_IDX, S_COL_IDX},
+        range::V_COL_IDX,
+        stack::{B0_COL_IDX, B1_COL_IDX},
+    },
     utils::are_equal,
-    TRACE_WIDTH,
+    CLK_COL_IDX, FMP_COL_IDX, STACK_TRACE_OFFSET, TRACE_WIDTH,
 };
 
 /// The number of auxiliary assertions.
-pub const NUM_ASSERTIONS: usize = TRACE_WIDTH;
+pub const NUM_ASSERTIONS: usize = TRACE_WIDTH - 21;
 
 /// The number of auxiliary assertions.
 pub const NUM_AUX_ASSERTIONS: usize = 2;
 
 pub fn get_assertions_first_step(result: &mut Vec<Assertion<Felt>>, main_trace_first_row: &[Felt]) {
+    const STACK_TRACE_END: usize = STACK_TRACE_OFFSET + STACK_TOP_SIZE - 1;
+    const B0_COL_IDX_MAIN_TRACE: usize = STACK_TRACE_OFFSET + B0_COL_IDX;
+    const B1_COL_IDX_MAIN_TRACE: usize = STACK_TRACE_OFFSET + B1_COL_IDX;
+
     for (col_idx, col) in main_trace_first_row.iter().enumerate().take(TRACE_WIDTH) {
-        result.push(Assertion::single(col_idx, 0, *col))
+        match col_idx {
+            // Hack: These columns already have an assertion
+            CLK_COL_IDX
+            | FMP_COL_IDX
+            | STACK_TRACE_OFFSET..=STACK_TRACE_END
+            | B0_COL_IDX_MAIN_TRACE
+            | B1_COL_IDX_MAIN_TRACE
+            | V_COL_IDX => (),
+            _ => result.push(Assertion::single(col_idx, 0, *col)),
+        }
     }
 }
 

air/src/constraints/range.rs

@@ -27,9 +27,10 @@ pub const CONSTRAINT_DEGREES: [usize; NUM_CONSTRAINTS] = [
 // --- MAIN TRACE ---------------------------------------------------------------------------------
 
 /// Returns the range checker's boundary assertions for the main trace at the first step.
-pub fn get_assertions_first_step(result: &mut Vec<Assertion<Felt>>) {
+pub fn get_assertions_first_step(result: &mut Vec<Assertion<Felt>>, main_trace_first_row: &[Felt]) {
     let step = 0;
     result.push(Assertion::single(V_COL_IDX, step, ZERO));
+    assert_eq!(main_trace_first_row[V_COL_IDX], ZERO);
 }
 
 /// Returns the range checker's boundary assertions for the main trace at the last step.

air/src/constraints/stack/mod.rs

@@ -185,15 +185,21 @@ pub fn enforce_general_constraints<E: FieldElement>(
 // --- MAIN TRACE ---------------------------------------------------------------------------------
 
 /// Returns the stack's boundary assertions for the main trace at the first step.
-pub fn get_assertions_first_step(result: &mut Vec<Assertion<Felt>>, stack_inputs: &[Felt]) {
+pub fn get_assertions_first_step(
+    result: &mut Vec<Assertion<Felt>>,
+    stack_inputs: &[Felt],
+    main_trace_first_row: &[Felt],
+) {
     // stack columns at the first step should be set to stack inputs, excluding overflow inputs.
     for (i, &value) in stack_inputs.iter().take(STACK_TOP_SIZE).enumerate() {
         result.push(Assertion::single(STACK_TRACE_OFFSET + i, 0, value));
+        assert_eq!(main_trace_first_row[STACK_TRACE_OFFSET + i], value);
     }
 
     // if there are remaining slots on top of the stack without specified values, set them to ZERO.
     for i in stack_inputs.len()..STACK_TOP_SIZE {
         result.push(Assertion::single(STACK_TRACE_OFFSET + i, 0, ZERO));
+        assert_eq!(main_trace_first_row[STACK_TRACE_OFFSET + i], ZERO);
     }
 
     // get the initial values for the bookkeeping columns.
@@ -206,10 +212,12 @@ pub fn get_assertions_first_step(result: &mut Vec<Assertion<Felt>>, stack_inputs
 
     // b0 should be initialized to the depth of the stack.
     result.push(Assertion::single(B0_COL_IDX, 0, Felt::new(depth as u64)));
+    assert_eq!(main_trace_first_row[B0_COL_IDX], Felt::new(depth as u64));
 
     // b1 should be initialized to the address of the last row in the overflow table, which is 0
     // when the overflow table is empty and -1 mod p when the stack is initialized with overflow.
     result.push(Assertion::single(B1_COL_IDX, 0, overflow_addr));
+    assert_eq!(main_trace_first_row[B1_COL_IDX], overflow_addr);
 }
 
 /// Returns the stack's boundary assertions for the main trace at the last step.

air/src/lib.rs

@@ -151,16 +151,23 @@ impl Air for ProcessorAir {
 
         // --- set assertions for the first step --------------------------------------------------
         // first value of clk is 0
+        // TODOP: Fix this main_trace_first_row hack in this PR?
         result.push(Assertion::single(CLK_COL_IDX, 0, ZERO));
+        assert_eq!(self.main_trace_first_row[CLK_COL_IDX], ZERO);
 
         // first value of fmp is 2^30
         result.push(Assertion::single(FMP_COL_IDX, 0, Felt::new(2u64.pow(30))));
+        assert_eq!(self.main_trace_first_row[FMP_COL_IDX], Felt::new(2u64.pow(30)));
 
         // add initial assertions for the stack.
-        stack::get_assertions_first_step(&mut result, self.stack_inputs.values());
+        stack::get_assertions_first_step(
+            &mut result,
+            self.stack_inputs.values(),
+            &self.main_trace_first_row,
+        );
 
         // add initial assertions for the range checker.
-        range::get_assertions_first_step(&mut result);
+        range::get_assertions_first_step(&mut result, &self.main_trace_first_row);
 
         // add initial assertions for logup
         logup::get_assertions_first_step(&mut result, &self.main_trace_first_row);