Add column for if is ROW is padding or execution

poseidon2-starky/src/plonky2/columns.rs

@@ -6,28 +6,31 @@ pub(crate) const SBOX_DEGREE: usize = 7;
 pub(crate) const ROUNDS_F: usize = 8;
 pub(crate) const ROUNDS_P: usize = 22;
 
+/// If the row is execution or padding.
+pub(crate) const COL_IS_EXE: usize = 0;
+
 /// The starting point of the column input
-pub(crate) const COL_INPUT_START: usize = 0;
+pub(crate) const COL_INPUT_START: usize = 1;
 
 /// The starting point of the state after each 1st full round
-pub(crate) const COL_1ST_FULLROUND_STATE_START: usize = COL_INPUT_START + STATE_SIZE; // 12
+pub(crate) const COL_1ST_FULLROUND_STATE_START: usize = COL_INPUT_START + STATE_SIZE; // 1 + 12
 
 /// The value of state[0] after each partial round
 pub(crate) const COL_PARTIAL_ROUND_STATE_START: usize =
-    COL_1ST_FULLROUND_STATE_START + STATE_SIZE * (ROUNDS_F / 2); // 12 + 48
+    COL_1ST_FULLROUND_STATE_START + STATE_SIZE * (ROUNDS_F / 2); // 1 + 12 + 48
 
 /// The starting point of the state after the partial round
 pub(crate) const COL_PARTIAL_ROUND_END_STATE_START: usize =
-    COL_PARTIAL_ROUND_STATE_START + ROUNDS_P - 1; // 12 + 48 + 22
+    COL_PARTIAL_ROUND_STATE_START + ROUNDS_P - 1; // 1 + 12 + 48 + 22
 
 /// The starting point of the state after each 2nd full round
 pub(crate) const COL_2ND_FULLROUND_STATE_START: usize =
-    COL_PARTIAL_ROUND_END_STATE_START + STATE_SIZE; // 12 + 48 + 22 + 12
+    COL_PARTIAL_ROUND_END_STATE_START + STATE_SIZE; // 1 + 12 + 48 + 22 + 12
 
 /// The starting point of the column output
 /// This is the same as the last state after the 2nd full round
 pub(crate) const COL_OUTPUT_START: usize =
-    COL_2ND_FULLROUND_STATE_START + STATE_SIZE * ((ROUNDS_F / 2) - 1); // 12 + 48 + 22 + 12 + 36
+    COL_2ND_FULLROUND_STATE_START + STATE_SIZE * ((ROUNDS_F / 2) - 1); // 1 + 12 + 48 + 22 + 12 + 36
 
 /// The total number of columns
-pub(crate) const NUM_COLS: usize = COL_2ND_FULLROUND_STATE_START + STATE_SIZE * (ROUNDS_F / 2); // 12 + 96 + 22 + 12 + 36 + 12
+pub(crate) const NUM_COLS: usize = COL_2ND_FULLROUND_STATE_START + STATE_SIZE * (ROUNDS_F / 2); // 1 + 12 + 96 + 22 + 12 + 36 + 12

poseidon2-starky/src/plonky2/generation.rs

@@ -8,6 +8,8 @@ use plonky2::hash::{
     poseidon2::{Poseidon2, WIDTH},
 };
 
+use super::columns::COL_IS_EXE;
+
 // Represent a row of the preimage
 #[derive(Debug, Clone, Default)]
 pub struct Row<Field: RichField> {
@@ -17,12 +19,17 @@ pub struct Row<Field: RichField> {
 /// Pad the trace to a power of 2.
 #[must_use]
 fn pad_trace<F: RichField>(mut trace: Vec<Vec<F>>) -> Vec<Vec<F>> {
-    let ext_trace_len = trace[0].len().next_power_of_two();
+    let original_len = trace[0].len();
+    let ext_trace_len = original_len.next_power_of_two();
 
     // All columns have their last value duplicated.
     for row in &mut trace {
         row.resize(ext_trace_len, *row.last().unwrap());
     }
+    // Set COL_IS_EXE to ZERO
+    for i in original_len..ext_trace_len {
+        trace[COL_IS_EXE][i] = F::ZERO;
+    }
 
     trace
 }
@@ -93,6 +100,8 @@ pub fn generate_poseidon2_trace<F: RichField>(step_rows: &Vec<Row<F>>) -> [Vec<F
     let mut trace: Vec<Vec<F>> = vec![vec![F::ZERO; trace_len]; NUM_COLS];
 
     for (i, row) in step_rows.iter().enumerate() {
+        trace[COL_IS_EXE][i] = F::ONE;
+
         for j in 0..STATE_SIZE {
             trace[COL_INPUT_START + j][i] = row.preimage[j];
         }

poseidon2-starky/src/plonky2/stark.rs

@@ -16,6 +16,8 @@ use starky::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 use std::fmt::Display;
 use std::marker::PhantomData;
 
+use super::columns::{COL_INPUT_START, COL_IS_EXE};
+
 // degree: 1
 fn add_rc_constraints<
     F: RichField + Extendable<D>,
@@ -209,8 +211,15 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for Poseidon2_12S
         P: PackedField<Scalar = FE>,
     {
         let lv = vars.local_values;
-        let mut state: [P; STATE_SIZE] =
-            matmul_external12_constraints(lv[0..STATE_SIZE].try_into().unwrap());
+
+        // row can be execution or padding.
+        yield_constr.constraint(lv[COL_IS_EXE] * (lv[COL_IS_EXE] - P::ONES));
+
+        let mut state: [P; STATE_SIZE] = matmul_external12_constraints(
+            lv[COL_INPUT_START..(COL_INPUT_START + STATE_SIZE)]
+                .try_into()
+                .unwrap(),
+        );
         // first full rounds
         for r in 0..(ROUNDS_F / 2) {
             state = add_rc_constraints(&state, r);