add the program segment to the preprocessed trace

Author: ohad-nir-starkware

stwo_cairo_prover/crates/cairo-air/src/lib.rs

@@ -23,15 +23,19 @@ pub enum PreProcessedTraceVariant {
     Canonical,
     CanonicalWithoutPedersen,
     CanonicalSmall,
+    CanonicalWithProgram,
 }
 impl PreProcessedTraceVariant {
-    pub fn to_preprocessed_trace(&self) -> PreProcessedTrace {
+    pub fn to_preprocessed_trace(&self, program: &[(u32, [u32; 8])]) -> PreProcessedTrace {
         match self {
             PreProcessedTraceVariant::Canonical => PreProcessedTrace::canonical(),
             PreProcessedTraceVariant::CanonicalWithoutPedersen => {
                 PreProcessedTrace::canonical_without_pedersen()
             }
             PreProcessedTraceVariant::CanonicalSmall => PreProcessedTrace::canonical_small(),
+            PreProcessedTraceVariant::CanonicalWithProgram => {
+                PreProcessedTrace::canonical_with_program(program)
+            }
         }
     }
 }

stwo_cairo_prover/crates/cairo-air/src/verifier.rs

@@ -346,10 +346,11 @@ pub fn verify_cairo_ex<MC: MerkleChannel>(
     pcs_config.mix_into(channel);
     let commitment_scheme_verifier = &mut CommitmentSchemeVerifier::<MC>::new(pcs_config);
 
+    let preprocessed_trace =
+        preprocessed_trace_variant.to_preprocessed_trace(&claim.public_data.public_memory.program);
+
     let mut log_sizes = claim.log_sizes();
-    log_sizes[PREPROCESSED_TRACE_IDX] = preprocessed_trace_variant
-        .to_preprocessed_trace()
-        .log_sizes();
+    log_sizes[PREPROCESSED_TRACE_IDX] = preprocessed_trace.log_sizes();
 
     // Preproccessed trace.
     commitment_scheme_verifier.commit(stark_proof.commitments[0], &log_sizes[0], channel);
@@ -375,7 +376,7 @@ pub fn verify_cairo_ex<MC: MerkleChannel>(
         &claim,
         &interaction_elements,
         &interaction_claim,
-        &preprocessed_trace_variant.to_preprocessed_trace().ids(),
+        &preprocessed_trace.ids(),
     );
     let components = component_generator.components();
 

stwo_cairo_prover/crates/common/src/preprocessed_columns/mod.rs

@@ -6,5 +6,6 @@ pub mod poseidon;
 pub mod poseidon_round_keys;
 pub mod preprocessed_trace;
 pub mod preprocessed_utils;
+pub mod program;
 #[cfg(feature = "prover")]
 pub mod simd_prelude;

stwo_cairo_prover/crates/common/src/preprocessed_columns/preprocessed_trace.rs

@@ -11,6 +11,7 @@ use super::pedersen::{PedersenPoints, PEDERSEN_TABLE_N_COLUMNS};
 use super::poseidon::{PoseidonRoundKeys, N_WORDS as POSEIDON_N_WORDS};
 #[cfg(feature = "prover")]
 use super::simd_prelude::*;
+use crate::preprocessed_columns::program;
 
 // Size to initialize the preprocessed trace with for `PreprocessedColumn::BitwiseXor`.
 const XOR_N_BITS: [u32; 5] = [4, 7, 8, 9, 10];
@@ -151,6 +152,23 @@ impl PreProcessedTrace {
         Self::from_columns(columns)
     }
 
+    pub fn canonical_with_program(program: &[(u32, [u32; 8])]) -> Self {
+        let canonical_without_program = Self::canonical().columns;
+        let curr_program_columns = (0..program::PROGRAM_N_COLUMNS).map(|x| {
+            Box::new(program::ProgramColumn::new(x, program)) as Box<dyn PreProcessedColumn>
+        });
+        let columns = chain!(canonical_without_program, curr_program_columns)
+            .sorted_by_key(|column| column.log_size())
+            .collect_vec();
+
+        assert!(
+            columns.iter().map(|col| 1 << col.log_size()).sum::<u32>() == CANONICAL_SIZE,
+            "Canonical preprocessed trace has unexpected size"
+        );
+
+        Self::from_columns(columns)
+    }
+
     pub fn log_sizes(&self) -> Vec<u32> {
         self.columns.iter().map(|c| c.log_size()).collect()
     }

stwo_cairo_prover/crates/common/src/preprocessed_columns/program.rs

@@ -0,0 +1,175 @@
+use stwo::core::fields::m31::M31;
+use stwo_constraint_framework::preprocessed_columns::PreProcessedColumnId;
+
+use super::preprocessed_trace::PreProcessedColumn;
+#[cfg(feature = "prover")]
+use super::simd_prelude::*;
+use crate::prover_types::cpu::FELT252_N_WORDS;
+
+pub const PROGRAM_N_COLUMNS: usize = FELT252_N_WORDS;
+
+/// Extracts a 9-bit limb from a 256-bit value stored as `[u32; 8]` little-endian limbs.
+fn extract_9bit_limb(value: &[u32; 8], limb_index: usize) -> u32 {
+    let bit_offset = limb_index * 9;
+    let word_index = bit_offset / 32;
+    let bit_shift = bit_offset % 32;
+    let mut result = value[word_index] >> bit_shift;
+    if bit_shift + 9 > 32 {
+        result |= value[word_index + 1] << (32 - bit_shift);
+    }
+    result & 0x1FF
+}
+
+#[derive(Debug)]
+pub struct ProgramColumn {
+    col_index: usize,
+    column_data: Vec<M31>,
+}
+impl ProgramColumn {
+    pub fn new(col_index: usize, program: &[(u32, [u32; 8])]) -> Self {
+        let padded_len = program.len().next_power_of_two();
+        let column_data = (0..padded_len)
+            .map(|i| {
+                if i < program.len() {
+                    M31::from_u32_unchecked(extract_9bit_limb(&program[i].1, col_index))
+                } else {
+                    M31(0)
+                }
+            })
+            .collect();
+        Self {
+            col_index,
+            column_data,
+        }
+    }
+
+    pub fn get_data(&self) -> &Vec<M31> {
+        &self.column_data
+    }
+}
+
+impl PreProcessedColumn for ProgramColumn {
+    fn log_size(&self) -> u32 {
+        self.column_data.len().ilog2()
+    }
+
+    fn id(&self) -> PreProcessedColumnId {
+        PreProcessedColumnId {
+            id: format!("curr_program_{}", self.col_index),
+        }
+    }
+
+    #[cfg(feature = "prover")]
+    fn packed_at(&self, vec_row: usize) -> PackedM31 {
+        let array = self.get_data()[(vec_row * N_LANES)..((vec_row + 1) * N_LANES)]
+            .try_into()
+            .unwrap();
+        PackedM31::from_array(array)
+    }
+
+    #[cfg(feature = "prover")]
+    fn gen_column_simd(&self) -> CircleEvaluation<SimdBackend, BaseField, BitReversedOrder> {
+        CircleEvaluation::new(
+            CanonicCoset::new(self.log_size()).circle_domain(),
+            BaseColumn::from_cpu(self.get_data()),
+        )
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    use super::*;
+
+    #[test]
+    fn test_extract_9bit_limb() {
+        // Value: 0b_000_111_010_110_001_011_101 = 0x1D6B5 (in 7 limbs of 9 bits)
+        // limb 0: 101 = 0b_000_101_101 ... let's use a concrete example.
+        // value[0] = 0x01FF_03FE = bits: limb0=0x1FE(510), limb1=0x1FF(511), ...
+        // Actually let's just test specific values.
+
+        // All zeros.
+        let value = [0u32; 8];
+        for i in 0..7 {
+            assert_eq!(extract_9bit_limb(&value, i), 0);
+        }
+
+        // limb 0 = lower 9 bits of value[0].
+        let value = [0x1AB, 0, 0, 0, 0, 0, 0, 0]; // 0x1AB = 0b110101011 = 427
+        assert_eq!(extract_9bit_limb(&value, 0), 0x1AB);
+
+        // limb 3 spans value[0] bits 27-31 and value[1] bits 0-3.
+        // bit_offset = 27, word_index = 0, bit_shift = 27.
+        // Place 0b_101010111 at bits 27-35:
+        // value[0] bits 27-31 = lower 5 bits of limb = 0b10111 => value[0] |= 0b10111 << 27
+        // value[1] bits 0-3 = upper 4 bits of limb = 0b1010 => value[1] |= 0b1010
+        let value = [0b10111 << 27, 0b1010, 0, 0, 0, 0, 0, 0];
+        assert_eq!(extract_9bit_limb(&value, 3), 0b101010111);
+    }
+
+    #[test]
+    fn test_program_column_new() {
+        let program = vec![
+            (0u32, [0x1FFu32, 0, 0, 0, 0, 0, 0, 0]), // limb 0 = 0x1FF (511)
+            (1, [0, 0, 0, 0, 0, 0, 0, 0]),           // all zeros
+        ];
+
+        let col0 = ProgramColumn::new(0, &program);
+        // Padded to next power of 2 = 2.
+        assert_eq!(col0.column_data.len(), 2);
+        assert_eq!(col0.column_data[0], M31(511));
+        assert_eq!(col0.column_data[1], M31(0));
+
+        let col1 = ProgramColumn::new(1, &program);
+        // 0x1FF in bits 0-8, limb 1 = bits 9-17 = 0.
+        assert_eq!(col1.column_data[0], M31(0));
+    }
+
+    #[test]
+    fn test_program_column_pads_to_power_of_two() {
+        let program = vec![
+            (0u32, [1u32, 0, 0, 0, 0, 0, 0, 0]),
+            (1, [2, 0, 0, 0, 0, 0, 0, 0]),
+            (2, [3, 0, 0, 0, 0, 0, 0, 0]),
+        ];
+
+        let col = ProgramColumn::new(0, &program);
+        // 3 entries padded to 4.
+        assert_eq!(col.column_data.len(), 4);
+        assert_eq!(col.column_data[0], M31(1));
+        assert_eq!(col.column_data[1], M31(2));
+        assert_eq!(col.column_data[2], M31(3));
+        assert_eq!(col.column_data[3], M31(0));
+    }
+
+    #[test]
+    fn test_program_column_log_size() {
+        let program = vec![(0u32, [0u32; 8]), (1, [0; 8]), (2, [0; 8])];
+        let col = ProgramColumn::new(0, &program);
+        // 3 entries padded to 4 = 2^2.
+        assert_eq!(col.log_size(), 2);
+    }
+
+    #[test]
+    fn test_program_column_id() {
+        let program = vec![(0u32, [0u32; 8])];
+        let col = ProgramColumn::new(3, &program);
+        assert_eq!(col.id().id, "curr_program_3");
+    }
+
+    #[cfg(feature = "prover")]
+    #[test]
+    fn test_packed_at() {
+        // Create a program with N_LANES entries so packed_at(0) covers them all.
+        let program: Vec<(u32, [u32; 8])> = (0..N_LANES)
+            .map(|i| (i as u32, [i as u32, 0, 0, 0, 0, 0, 0, 0]))
+            .collect();
+
+        let col = ProgramColumn::new(0, &program);
+        let packed = col.packed_at(0);
+        let array = packed.to_array();
+
+        for (i, val) in array.iter().enumerate() {
+            assert_eq!(*val, M31(i as u32));
+        }
+    }
+}

stwo_cairo_prover/crates/prover/src/prover.rs

@@ -92,7 +92,11 @@ where
             .half_coset,
     );
 
-    let preprocessed_trace = Arc::new(prover_params.preprocessed_trace.to_preprocessed_trace());
+    let preprocessed_trace = Arc::new(
+        prover_params
+            .preprocessed_trace
+            .to_preprocessed_trace(&input.program),
+    );
     let preprocessed_trace_polys =
         SimdBackend::interpolate_columns(gen_trace(preprocessed_trace.clone()), &twiddles);
 

stwo_cairo_prover/crates/prover/src/witness/components/memory_id_to_big.rs

@@ -672,7 +672,7 @@ mod tests {
         // Preprocessed trace.
         let mut tree_builder = commitment_scheme.tree_builder();
         tree_builder.extend_evals(gen_trace(Arc::new(
-            PreProcessedTraceVariant::CanonicalWithoutPedersen.to_preprocessed_trace(),
+            PreProcessedTraceVariant::CanonicalWithoutPedersen.to_preprocessed_trace(&[]),
         )));
         tree_builder.finalize_interaction();
 

stwo_cairo_prover/crates/prover/src/witness/preprocessed_trace.rs

@@ -22,11 +22,12 @@ pub fn generate_preprocessed_commitment_root<MC: MerkleChannel>(
     log_blowup_factor: u32,
     preprocessed_trace: PreProcessedTraceVariant,
     lifting_log_size: Option<u32>,
+    program: &[(u32, [u32; 8])],
 ) -> <<MC as MerkleChannel>::H as MerkleHasherLifted>::Hash
 where
     SimdBackend: BackendForChannel<MC>,
 {
-    let preprocessed_trace = Arc::new(preprocessed_trace.to_preprocessed_trace());
+    let preprocessed_trace = Arc::new(preprocessed_trace.to_preprocessed_trace(program));
 
     // Precompute twiddles for the commitment scheme.
     let mut max_log_size = preprocessed_trace.log_sizes().into_iter().max().unwrap();
@@ -80,6 +81,7 @@ fn test_canonical_preprocessed_root_regression() {
         log_blowup_factor,
         PreProcessedTraceVariant::Canonical,
         None,
+        &[],
     );
 
     assert_eq!(root, expected);
@@ -101,6 +103,7 @@ fn test_small_canonical_preprocessed_root_regression() {
         log_blowup_factor,
         PreProcessedTraceVariant::CanonicalSmall,
         None,
+        &[],
     );
 
     assert_eq!(root, expected);

stwo_cairo_prover/crates/prover/src/witness/utils.rs

@@ -136,13 +136,19 @@ fn get_preprocessed_roots<MC: MerkleChannel>(
     max_log_blowup_factor: u32,
     preprocessed_trace: PreProcessedTraceVariant,
     lifting_log_size: Option<u32>,
+    program: &[(u32, [u32; 8])],
 ) -> Vec<<MC::H as MerkleHasherLifted>::Hash>
 where
     stwo::prover::backend::simd::SimdBackend: BackendForChannel<MC>,
 {
     (1..=max_log_blowup_factor)
         .map(|i| {
-            generate_preprocessed_commitment_root::<MC>(i, preprocessed_trace, lifting_log_size)
+            generate_preprocessed_commitment_root::<MC>(
+                i,
+                preprocessed_trace,
+                lifting_log_size,
+                program,
+            )
         })
         .collect_vec()
 }
@@ -158,6 +164,7 @@ pub fn export_preprocessed_roots() {
         max_log_blowup_factor,
         PreProcessedTraceVariant::Canonical,
         None,
+        &[],
     );
     blake_roots.iter().enumerate().for_each(|(i, root)| {
         let root_bytes = root.0;
@@ -179,6 +186,7 @@ pub fn export_preprocessed_roots() {
             max_log_blowup_factor,
             PreProcessedTraceVariant::CanonicalWithoutPedersen,
             None,
+            &[],
         )
         .into_iter()
         .enumerate()
@@ -203,6 +211,7 @@ pub fn export_circuit_cairo_verifier_preprocessed_roots() {
         max_log_blowup_factor,
         PreProcessedTraceVariant::CanonicalSmall,
         Some(lifting_log_size),
+        &[],
     );
     blake_m31_small_roots
         .iter()