refactor: add step abstraction to trace (#662)

* Start with multi row steps on trace

* Add table and step view functionality

* Save work in progress

* Remove unused feature from cairo-vm

* Add comment in Cargo.toml

* Fix some AIRs with the refactor

* Change Frame to Table struct in stark proof for OODS

* Further refactor in Table APIs

* Apply clippy suggestion

* Fix trace oods evaluations name in proof struct

* Save work in progress

* Fix unhandled merge conflict

* Fix some Cairo constraints

* Fix permutation argument Cairo constraints

* Fix all Cairo constraints

* Fix all compilation

* Fix all tests

* Apply clippy suggestion

* Remove argument in empty() function

* revert a mistaken change

* Remove unused code

* Use step size constant in validate_trace function

* Remove unused method and remove hardcoded 1 in prover

* Remove hardcoded value in trace construction

* Add documentation to code

* Add more documentation

* Change hardcoded step size

* Remove useless test

provers/cairo/src/execution_trace.rs

@@ -292,7 +292,7 @@ pub fn build_cairo_execution_trace(
         add_rc_builtin_columns(&mut trace_cols, range_check_builtin_range.clone(), memory);
     }
 
-    TraceTable::from_columns(trace_cols)
+    TraceTable::from_columns(trace_cols, 1)
 }
 
 // Build range-check builtin columns: rc_0, rc_1, ... , rc_7, rc_value
@@ -620,7 +620,7 @@ mod test {
             FieldElement::from(7),
             FieldElement::from(7),
         ];
-        let table = TraceTable::<Stark252PrimeField>::from_columns(columns);
+        let table = TraceTable::<Stark252PrimeField>::from_columns(columns, 1);
 
         let (col, rc_min, rc_max) = get_rc_holes(&table, &[0, 1, 2]);
         assert_eq!(col, expected_col);
@@ -635,7 +635,10 @@ mod test {
         let data = row.repeat(8);
         let table = Table::new(data, 36);
 
-        let mut main_trace = TraceTable::<Stark252PrimeField> { table };
+        let mut main_trace = TraceTable::<Stark252PrimeField> {
+            table,
+            step_size: 1,
+        };
 
         let rc_holes = vec![
             Felt252::from(1),
@@ -736,7 +739,7 @@ mod test {
         trace_cols[FRAME_DST_ADDR][1] = Felt252::from(9);
         trace_cols[FRAME_OP0_ADDR][1] = Felt252::from(10);
         trace_cols[FRAME_OP1_ADDR][1] = Felt252::from(11);
-        let mut trace = TraceTable::from_columns(trace_cols);
+        let mut trace = TraceTable::from_columns(trace_cols, 1);
 
         let memory_holes = vec![Felt252::from(4), Felt252::from(7), Felt252::from(8)];
         fill_memory_holes(&mut trace, &memory_holes);

provers/cairo/src/tests/integration_tests.rs

@@ -175,7 +175,7 @@ fn test_verifier_rejects_proof_with_changed_range_check_value() {
     last_column[0] = malicious_rc_value;
     malicious_trace_columns[n_cols - 1] = last_column;
 
-    let malicious_trace = TraceTable::from_columns(malicious_trace_columns);
+    let malicious_trace = TraceTable::from_columns(malicious_trace_columns, 1);
     let proof = generate_cairo_proof(&malicious_trace, &pub_inputs, &proof_options).unwrap();
     assert!(!verify_cairo_proof(&proof, &pub_inputs, &proof_options));
 }
@@ -243,7 +243,7 @@ fn test_verifier_rejects_proof_with_changed_output() {
     output_value_column[output_row_idx] = malicious_output_value;
     malicious_trace_columns[output_col_idx + 4] = output_value_column;
 
-    let malicious_trace = TraceTable::from_columns(malicious_trace_columns);
+    let malicious_trace = TraceTable::from_columns(malicious_trace_columns, 1);
     let proof = generate_cairo_proof(&malicious_trace, &pub_inputs, &proof_options).unwrap();
     assert!(!verify_cairo_proof(&proof, &pub_inputs, &proof_options));
 }

provers/stark/src/debug.rs

@@ -28,7 +28,7 @@ pub fn validate_trace<F: IsFFTField, A: AIR<Field = F>>(
         })
         .collect();
 
-    let trace = TraceTable::from_columns(trace_columns);
+    let trace = TraceTable::from_columns(trace_columns, A::STEP_SIZE);
 
     // --------- VALIDATE BOUNDARY CONSTRAINTS ------------
     air.boundary_constraints(rap_challenges)
@@ -40,7 +40,7 @@ pub fn validate_trace<F: IsFFTField, A: AIR<Field = F>>(
             let boundary_value = constraint.value.clone();
             let trace_value = trace.get(step, col);
 
-            if boundary_value != trace_value {
+            if &boundary_value != trace_value {
                 ret = false;
                 error!("Boundary constraint inconsistency - Expected value {} in step {} and column {}, found: {}", boundary_value.representative(), step, col, trace_value.representative());
             }
@@ -57,7 +57,7 @@ pub fn validate_trace<F: IsFFTField, A: AIR<Field = F>>(
         .collect();
 
     // Iterate over trace and compute transitions
-    for step in 0..trace.n_rows() {
+    for step in 0..trace.num_steps() {
         let frame = Frame::read_from_trace(&trace, step, 1, &air.context().transition_offsets);
 
         let evaluations = air.compute_transition(&frame, rap_challenges);

provers/stark/src/examples/dummy_air.rs

@@ -24,6 +24,8 @@ impl AIR for DummyAIR {
     type RAPChallenges = ();
     type PublicInputs = ();
 
+    const STEP_SIZE: usize = 1;
+
     fn new(
         trace_length: usize,
         _pub_inputs: &Self::PublicInputs,
@@ -63,13 +65,18 @@ impl AIR for DummyAIR {
         frame: &Frame<Self::Field>,
         _rap_challenges: &Self::RAPChallenges,
     ) -> Vec<FieldElement<Self::Field>> {
-        let first_row = frame.get_row(0);
-        let second_row = frame.get_row(1);
-        let third_row = frame.get_row(2);
+        let first_step = frame.get_evaluation_step(0);
+        let second_step = frame.get_evaluation_step(1);
+        let third_step = frame.get_evaluation_step(2);
+
+        let flag = first_step.get_evaluation_element(0, 0);
+        let a0 = first_step.get_evaluation_element(0, 1);
+        let a1 = second_step.get_evaluation_element(0, 1);
+        let a2 = third_step.get_evaluation_element(0, 1);
 
-        let f_constraint = first_row[0] * (first_row[0] - FieldElement::one());
+        let f_constraint = flag * (flag - FieldElement::one());
 
-        let fib_constraint = third_row[1] - second_row[1] - first_row[1];
+        let fib_constraint = a2 - a1 - a0;
 
         vec![f_constraint, fib_constraint]
     }
@@ -118,5 +125,5 @@ pub fn dummy_trace<F: IsFFTField>(trace_length: usize) -> TraceTable<F> {
         ret.push(ret[i - 1].clone() + ret[i - 2].clone());
     }
 
-    TraceTable::from_columns(vec![vec![FieldElement::<F>::one(); trace_length], ret])
+    TraceTable::from_columns(vec![vec![FieldElement::<F>::one(); trace_length], ret], 1)
 }

provers/stark/src/examples/fibonacci_2_cols_shifted.rs

@@ -56,6 +56,8 @@ where
     type RAPChallenges = ();
     type PublicInputs = PublicInputs<Self::Field>;
 
+    const STEP_SIZE: usize = 1;
+
     fn new(
         trace_length: usize,
         pub_inputs: &Self::PublicInputs,
@@ -97,11 +99,17 @@ where
         frame: &Frame<Self::Field>,
         _rap_challenges: &Self::RAPChallenges,
     ) -> Vec<FieldElement<Self::Field>> {
-        let first_row = frame.get_row(0);
-        let second_row = frame.get_row(1);
+        let first_row = frame.get_evaluation_step(0);
+        let second_row = frame.get_evaluation_step(1);
+
+        let a0_0 = first_row.get_evaluation_element(0, 0);
+        let a0_1 = first_row.get_evaluation_element(0, 1);
+
+        let a1_0 = second_row.get_evaluation_element(0, 0);
+        let a1_1 = second_row.get_evaluation_element(0, 1);
 
-        let first_transition = &second_row[0] - &first_row[1];
-        let second_transition = &second_row[1] - &first_row[0] - &first_row[1];
+        let first_transition = a1_0 - a0_1;
+        let second_transition = a1_1 - a0_0 - a0_1;
 
         vec![first_transition, second_transition]
     }
@@ -156,7 +164,7 @@ pub fn compute_trace<F: IsFFTField>(
         col1.push(y.clone());
     }
 
-    TraceTable::from_columns(vec![col0, col1])
+    TraceTable::from_columns(vec![col0, col1], 1)
 }
 
 #[cfg(test)]

provers/stark/src/examples/fibonacci_2_columns.rs

@@ -32,6 +32,8 @@ where
     type RAPChallenges = ();
     type PublicInputs = FibonacciPublicInputs<Self::Field>;
 
+    const STEP_SIZE: usize = 1;
+
     fn new(
         trace_length: usize,
         pub_inputs: &Self::PublicInputs,
@@ -73,14 +75,19 @@ where
         frame: &Frame<Self::Field>,
         _rap_challenges: &Self::RAPChallenges,
     ) -> Vec<FieldElement<Self::Field>> {
-        let first_row = frame.get_row(0);
-        let second_row = frame.get_row(1);
+        let first_step = frame.get_evaluation_step(0);
+        let second_step = frame.get_evaluation_step(1);
 
         // constraints of Fibonacci sequence (2 terms per step):
         // s_{0, i+1} = s_{0, i} + s_{1, i}
         // s_{1, i+1} = s_{1, i} + s_{0, i+1}
-        let first_transition = &second_row[0] - &first_row[0] - &first_row[1];
-        let second_transition = &second_row[1] - &first_row[1] - &second_row[0];
+        let s0_0 = first_step.get_evaluation_element(0, 0);
+        let s0_1 = first_step.get_evaluation_element(0, 1);
+        let s1_0 = second_step.get_evaluation_element(0, 0);
+        let s1_1 = second_step.get_evaluation_element(0, 1);
+
+        let first_transition = s1_0 - s0_0 - s0_1;
+        let second_transition = s1_1 - s0_1 - s1_0;
 
         vec![first_transition, second_transition]
     }
@@ -132,5 +139,5 @@ pub fn compute_trace<F: IsFFTField>(
         ret2.push(new_val + ret2[i - 1].clone());
     }
 
-    TraceTable::from_columns(vec![ret1, ret2])
+    TraceTable::from_columns(vec![ret1, ret2], 1)
 }

provers/stark/src/examples/fibonacci_rap.rs

@@ -45,6 +45,8 @@ where
     type RAPChallenges = FieldElement<Self::Field>;
     type PublicInputs = FibonacciRAPPublicInputs<Self::Field>;
 
+    const STEP_SIZE: usize = 1;
+
     fn new(
         trace_length: usize,
         pub_inputs: &Self::PublicInputs,
@@ -92,7 +94,7 @@ where
                 aux_col.push(z_i * n_p_term.div(p_term));
             }
         }
-        TraceTable::from_columns(vec![aux_col])
+        TraceTable::from_columns(vec![aux_col], 1)
     }
 
     fn build_rap_challenges(
@@ -112,19 +114,22 @@ where
         gamma: &Self::RAPChallenges,
     ) -> Vec<FieldElement<Self::Field>> {
         // Main constraints
-        let first_row = frame.get_row(0);
-        let second_row = frame.get_row(1);
-        let third_row = frame.get_row(2);
+        let first_step = frame.get_evaluation_step(0);
+        let second_step = frame.get_evaluation_step(1);
+        let third_step = frame.get_evaluation_step(2);
+
+        let a0 = first_step.get_evaluation_element(0, 0);
+        let a1 = second_step.get_evaluation_element(0, 0);
+        let a2 = third_step.get_evaluation_element(0, 0);
 
-        let mut constraints =
-            vec![third_row[0].clone() - second_row[0].clone() - first_row[0].clone()];
+        let mut constraints = vec![a2 - a1 - a0];
 
         // Auxiliary constraints
-        let z_i = &frame.get_row(0)[2];
-        let z_i_plus_one = &frame.get_row(1)[2];
+        let z_i = first_step.get_evaluation_element(0, 2);
+        let z_i_plus_one = second_step.get_evaluation_element(0, 2);
 
-        let a_i = &frame.get_row(0)[0];
-        let b_i = &frame.get_row(0)[1];
+        let a_i = first_step.get_evaluation_element(0, 0);
+        let b_i = first_step.get_evaluation_element(0, 1);
 
         let eval = z_i_plus_one * (b_i + gamma) - z_i * (a_i + gamma);
 
@@ -186,7 +191,7 @@ pub fn fibonacci_rap_trace<F: IsFFTField>(
     let mut trace_cols = vec![fib_seq, fib_permuted];
     resize_to_next_power_of_two(&mut trace_cols);
 
-    TraceTable::from_columns(trace_cols)
+    TraceTable::from_columns(trace_cols, 1)
 }
 
 #[cfg(test)]

provers/stark/src/examples/quadratic_air.rs

@@ -36,6 +36,8 @@ where
     type RAPChallenges = ();
     type PublicInputs = QuadraticPublicInputs<Self::Field>;
 
+    const STEP_SIZE: usize = 1;
+
     fn new(
         trace_length: usize,
         pub_inputs: &Self::PublicInputs,
@@ -77,10 +79,13 @@ where
         frame: &Frame<Self::Field>,
         _rap_challenges: &Self::RAPChallenges,
     ) -> Vec<FieldElement<Self::Field>> {
-        let first_row = frame.get_row(0);
-        let second_row = frame.get_row(1);
+        let first_step = frame.get_evaluation_step(0);
+        let second_step = frame.get_evaluation_step(1);
+
+        let x = first_step.get_evaluation_element(0, 0);
+        let x_squared = second_step.get_evaluation_element(0, 0);
 
-        vec![&second_row[0] - &first_row[0] * &first_row[0]]
+        vec![x_squared - x * x]
     }
 
     fn number_auxiliary_rap_columns(&self) -> usize {
@@ -125,5 +130,5 @@ pub fn quadratic_trace<F: IsFFTField>(
         ret.push(ret[i - 1].clone() * ret[i - 1].clone());
     }
 
-    TraceTable::from_columns(vec![ret])
+    TraceTable::from_columns(vec![ret], 1)
 }

provers/stark/src/examples/simple_fibonacci.rs

@@ -37,6 +37,8 @@ where
     type RAPChallenges = ();
     type PublicInputs = FibonacciPublicInputs<Self::Field>;
 
+    const STEP_SIZE: usize = 1;
+
     fn new(
         trace_length: usize,
         pub_inputs: &Self::PublicInputs,
@@ -82,11 +84,15 @@ where
         frame: &Frame<Self::Field>,
         _rap_challenges: &Self::RAPChallenges,
     ) -> Vec<FieldElement<Self::Field>> {
-        let first_row = frame.get_row(0);
-        let second_row = frame.get_row(1);
-        let third_row = frame.get_row(2);
+        let first_step = frame.get_evaluation_step(0);
+        let second_step = frame.get_evaluation_step(1);
+        let third_step = frame.get_evaluation_step(2);
+
+        let a0 = first_step.get_evaluation_element(0, 0);
+        let a1 = second_step.get_evaluation_element(0, 0);
+        let a2 = third_step.get_evaluation_element(0, 0);
 
-        vec![third_row[0].clone() - second_row[0].clone() - first_row[0].clone()]
+        vec![a2 - a1 - a0]
     }
 
     fn boundary_constraints(
@@ -129,5 +135,5 @@ pub fn fibonacci_trace<F: IsFFTField>(
         ret.push(ret[i - 1].clone() + ret[i - 2].clone());
     }
 
-    TraceTable::from_columns(vec![ret])
+    TraceTable::from_columns(vec![ret], 1)
 }