prover: move program columns into preprocessed trace (nexus-xyz#562)

* prover: merge program trace with preprocessed cols

to enforce the prover to commit to program trace before generating
logup challenges

* mix ad into the hasher

Author: slumber

prover/src/column.rs

@@ -537,6 +537,7 @@ pub enum Column {
 // }
 
 #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, ColumnsEnum)]
+#[column_derive(string_id)]
 pub enum ProgramColumn {
     /// Program memory content: every Pc in the program memory
     #[size = 4]

prover/src/lib.rs

@@ -29,6 +29,7 @@ pub fn verify(proof: Proof, view: &nexus_vm::emulator::View) -> Result<(), Verif
     machine::Machine::<machine::BaseComponents>::verify(
         proof,
         view.get_program_memory(),
+        view.view_associated_data().as_deref().unwrap_or_default(),
         view.get_initial_memory(),
         view.get_exit_code(),
         view.get_public_output(),

prover/src/machine.rs

@@ -5,7 +5,7 @@ use stwo_prover::{
     constraint_framework::TraceLocationAllocator,
     core::{
         backend::simd::SimdBackend,
-        channel::Blake2sChannel,
+        channel::{Blake2sChannel, Channel},
         fields::qm31::SecureField,
         pcs::{CommitmentSchemeProver, CommitmentSchemeVerifier, PcsConfig},
         poly::circle::{CanonicCoset, PolyOps},
@@ -14,9 +14,7 @@ use stwo_prover::{
     },
 };
 
-use super::trace::eval::{
-    INTERACTION_TRACE_IDX, ORIGINAL_TRACE_IDX, PREPROCESSED_TRACE_IDX, PROGRAM_TRACE_IDX,
-};
+use super::trace::eval::{INTERACTION_TRACE_IDX, ORIGINAL_TRACE_IDX, PREPROCESSED_TRACE_IDX};
 use super::trace::{
     program::iter_program_steps, program_trace::ProgramTracesBuilder, sidenote::SideNote,
     PreprocessedTraces, TracesBuilder,
@@ -35,7 +33,7 @@ use crate::{
         RangeCheckChip, RegisterMemCheckChip, SllChip, SltChip, SltuChip, SraChip, SrlChip,
         SubChip, SyscallChip, TimestampChip,
     },
-    column::PreprocessedColumn,
+    column::{PreprocessedColumn, ProgramColumn},
     components::{self, AllLookupElements},
     traits::generate_interaction_trace,
 };
@@ -112,6 +110,10 @@ impl<C: MachineChip + Sync> Machine<C> {
 
         // Setup protocol.
         let prover_channel = &mut Blake2sChannel::default();
+        for byte in view.view_associated_data().unwrap_or_default() {
+            prover_channel.mix_u64(byte.into());
+        }
+
         let mut commitment_scheme =
             CommitmentSchemeProver::<SimdBackend, Blake2sMerkleChannel>::new(config, &twiddles);
 
@@ -142,8 +144,13 @@ impl<C: MachineChip + Sync> Machine<C> {
         let finalized_program_trace = program_traces.finalize();
 
         let mut tree_builder = commitment_scheme.tree_builder();
-        let _preprocessed_trace_location =
-            tree_builder.extend_evals(preprocessed_trace.clone().into_circle_evaluation());
+        let _preprocessed_trace_location = tree_builder.extend_evals(
+            preprocessed_trace
+                .clone()
+                .into_circle_evaluation()
+                .into_iter()
+                .chain(finalized_program_trace.clone().into_circle_evaluation()),
+        );
         tree_builder.commit(prover_channel);
 
         let mut tree_builder = commitment_scheme.tree_builder();
@@ -165,12 +172,6 @@ impl<C: MachineChip + Sync> Machine<C> {
         let _interaction_trace_location = tree_builder.extend_evals(interaction_trace);
         tree_builder.commit(prover_channel);
 
-        // Fill columns of the program trace.
-        let mut tree_builder = commitment_scheme.tree_builder();
-        let _program_trace_location =
-            tree_builder.extend_evals(finalized_program_trace.into_circle_evaluation());
-        tree_builder.commit(prover_channel);
-
         let component = MachineComponent::new(
             &mut TraceLocationAllocator::default(),
             MachineEval::<C>::new(log_size, lookup_elements),
@@ -192,6 +193,7 @@ impl<C: MachineChip + Sync> Machine<C> {
     pub fn verify(
         proof: Proof,
         program_info: &ProgramInfo,
+        ad: &[u8],
         init_memory: &[MemoryInitializationEntry],
         exit_code: &[PublicOutputEntry],
         output_memory: &[PublicOutputEntry],
@@ -210,12 +212,16 @@ impl<C: MachineChip + Sync> Machine<C> {
 
         let config = PcsConfig::default();
         let verifier_channel = &mut Blake2sChannel::default();
+        for &byte in ad {
+            verifier_channel.mix_u64(byte.into());
+        }
+
         let commitment_scheme = &mut CommitmentSchemeVerifier::<Blake2sMerkleChannel>::new(config);
 
-        // simulate the prover and compute expected commitment to preprocessed and program traces
+        // simulate the prover and compute expected commitment to preprocessed trace
         {
             let config = PcsConfig::default();
-            let verifier_channel = &mut Blake2sChannel::default();
+            let verifier_channel = &mut verifier_channel.clone();
             let twiddles = SimdBackend::precompute_twiddles(
                 CanonicCoset::new(
                     log_size + LOG_CONSTRAINT_DEGREE + config.fri_config.log_blowup_factor,
@@ -228,18 +234,6 @@ impl<C: MachineChip + Sync> Machine<C> {
                     config, &twiddles,
                 );
             let preprocessed_trace = PreprocessedTraces::new(log_size);
-            let mut tree_builder = commitment_scheme.tree_builder();
-            let _preprocessed_trace_location =
-                tree_builder.extend_evals(preprocessed_trace.into_circle_evaluation());
-            tree_builder.commit(verifier_channel);
-
-            let preprocessed_expected = commitment_scheme.roots()[PREPROCESSED_TRACE_IDX];
-            let preprocessed = proof.commitments[PREPROCESSED_TRACE_IDX];
-            if preprocessed_expected != preprocessed {
-                return Err(VerificationError::InvalidStructure(format!("invalid commitment to preprocessed trace: \
-                                                                        expected {preprocessed_expected}, got {preprocessed}")));
-            }
-
             let program_trace = ProgramTracesBuilder::new(
                 log_size,
                 program_info,
@@ -248,14 +242,21 @@ impl<C: MachineChip + Sync> Machine<C> {
                 output_memory,
             )
             .finalize();
+
             let mut tree_builder = commitment_scheme.tree_builder();
-            tree_builder.extend_evals(program_trace.into_circle_evaluation());
+            let _preprocessed_trace_location = tree_builder.extend_evals(
+                preprocessed_trace
+                    .into_circle_evaluation()
+                    .into_iter()
+                    .chain(program_trace.into_circle_evaluation()),
+            );
             tree_builder.commit(verifier_channel);
-            let program_expected = commitment_scheme.roots()[1];
-            let program = proof.commitments[PROGRAM_TRACE_IDX];
-            if program_expected != program {
-                return Err(VerificationError::InvalidStructure(format!("invalid commitment to program trace: \
-                                                                        expected {program_expected}, got {program}")));
+
+            let preprocessed_expected = commitment_scheme.roots()[PREPROCESSED_TRACE_IDX];
+            let preprocessed = proof.commitments[PREPROCESSED_TRACE_IDX];
+            if preprocessed_expected != preprocessed {
+                return Err(VerificationError::InvalidStructure(format!("invalid commitment to preprocessed trace: \
+                                                                        expected {preprocessed_expected}, got {preprocessed}")));
             }
         }
 
@@ -270,7 +271,7 @@ impl<C: MachineChip + Sync> Machine<C> {
             .map_cols(|_| log_size);
         // use the fact that preprocessed columns are only allowed to have [0] mask
         sizes[PREPROCESSED_TRACE_IDX] = std::iter::repeat(log_size)
-            .take(PreprocessedColumn::COLUMNS_NUM)
+            .take(PreprocessedColumn::COLUMNS_NUM + ProgramColumn::COLUMNS_NUM)
             .collect();
 
         for idx in [PREPROCESSED_TRACE_IDX, ORIGINAL_TRACE_IDX] {
@@ -284,10 +285,12 @@ impl<C: MachineChip + Sync> Machine<C> {
             MachineEval::<C>::new(log_size, lookup_elements),
             claimed_sum,
         );
-        // TODO: prover must commit to the program trace before generating challenges.
-        for idx in [INTERACTION_TRACE_IDX, PROGRAM_TRACE_IDX] {
-            commitment_scheme.commit(proof.commitments[idx], &sizes[idx], verifier_channel);
-        }
+
+        commitment_scheme.commit(
+            proof.commitments[INTERACTION_TRACE_IDX],
+            &sizes[INTERACTION_TRACE_IDX],
+            verifier_channel,
+        );
 
         verify(&[&component], verifier_channel, commitment_scheme, proof)
     }
@@ -327,6 +330,7 @@ mod tests {
         Machine::<BaseComponents>::verify(
             proof,
             view.get_program_memory(),
+            &[],
             view.get_initial_memory(),
             view.get_exit_code(),
             view.get_public_output(),

prover/src/test_utils.rs

@@ -60,11 +60,19 @@ pub(crate) fn commit_traces<'a, C: MachineChip>(
     let mut commitment_scheme =
         CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, twiddles);
     let mut prover_channel = Blake2sChannel::default();
+
+    let program_trace =
+        program_traces.unwrap_or_else(|| ProgramTracesBuilder::dummy(traces.log_size()).finalize());
     // Preprocessed trace
     let preprocessed_trace = PreprocessedTraces::new(traces.log_size());
     let mut tree_builder = commitment_scheme.tree_builder();
-    let _preprocessed_trace_location =
-        tree_builder.extend_evals(preprocessed_trace.clone().into_circle_evaluation());
+    let _preprocessed_trace_location = tree_builder.extend_evals(
+        preprocessed_trace
+            .clone()
+            .into_circle_evaluation()
+            .into_iter()
+            .chain(program_trace.clone().into_circle_evaluation()),
+    );
     tree_builder.commit(&mut prover_channel);
 
     // Original trace
@@ -74,23 +82,13 @@ pub(crate) fn commit_traces<'a, C: MachineChip>(
     let mut all_elements = AllLookupElements::default();
     C::draw_lookup_elements(&mut all_elements, &mut prover_channel);
 
-    let program_trace =
-        program_traces.unwrap_or_else(|| ProgramTracesBuilder::dummy(traces.log_size()).finalize());
-
     // Interaction Trace
     let (interaction_trace, claimed_sum) =
         generate_interaction_trace::<C>(traces, &preprocessed_trace, &program_trace, &all_elements);
     let mut tree_builder = commitment_scheme.tree_builder();
     let _interaction_trace_location = tree_builder.extend_evals(interaction_trace.clone());
     tree_builder.commit(&mut prover_channel);
 
-    // Program Trace
-    let mut tree_builder = commitment_scheme.tree_builder();
-    let _program_trace_location =
-        tree_builder.extend_evals(program_trace.clone().into_circle_evaluation());
-    tree_builder.commit(&mut prover_channel);
-    // TODO: make the verifier check the program trace commitment against the expected value
-
     CommittedTraces {
         commitment_scheme,
         prover_channel,
@@ -123,10 +121,13 @@ pub(crate) fn assert_chip<C: MachineChip>(
     } = commit_traces::<C>(config, &twiddles, &finalized_trace, program_trace);
 
     let trace_evals = TreeVec::new(vec![
-        preprocessed_trace.into_circle_evaluation(),
+        [
+            preprocessed_trace.into_circle_evaluation(),
+            program_trace.into_circle_evaluation(),
+        ]
+        .concat(),
         finalized_trace.into_circle_evaluation(),
         interaction_trace,
-        program_trace.into_circle_evaluation(),
     ]);
     let trace_polys = trace_evals.map(|trace| {
         trace

prover/src/trace/eval.rs

@@ -10,13 +10,12 @@ use crate::column::{
 pub use stwo_prover::constraint_framework::{
     INTERACTION_TRACE_IDX, ORIGINAL_TRACE_IDX, PREPROCESSED_TRACE_IDX,
 };
-pub const PROGRAM_TRACE_IDX: usize = 3; // After INTERACTION_TRACE_IDX; the verifier is supposed to know the commitment of the program trace
 
 // Trace evaluation at the current row and the next row.
 pub struct TraceEval<E: EvalAtRow> {
     evals: Vec<[E::F; 2]>,
     preprocessed_evals: Vec<E::F>,
-    program_evals: Vec<[E::F; 1]>, // only the current row
+    program_evals: Vec<E::F>, // only the current row
 }
 
 impl<E: EvalAtRow> TraceEval<E> {
@@ -25,6 +24,10 @@ impl<E: EvalAtRow> TraceEval<E> {
             .iter()
             .map(|&id| eval.get_preprocessed_column(PreProcessedColumnId { id: id.to_owned() }))
             .collect();
+        let program_evals = ProgramColumn::STRING_IDS
+            .iter()
+            .map(|&id| eval.get_preprocessed_column(PreProcessedColumnId { id: id.to_owned() }))
+            .collect();
         let evals = Column::ALL_VARIANTS
             .iter()
             .flat_map(|col| std::iter::repeat(col).take(col.size()))
@@ -39,10 +42,6 @@ impl<E: EvalAtRow> TraceEval<E> {
                 }
             })
             .collect();
-        let program_evals =
-            std::iter::repeat_with(|| eval.next_interaction_mask(PROGRAM_TRACE_IDX, [0]))
-                .take(ProgramColumn::COLUMNS_NUM)
-                .collect();
         Self {
             evals,
             preprocessed_evals,
@@ -93,7 +92,7 @@ impl<E: EvalAtRow> TraceEval<E> {
         assert_eq!(col.size(), N, "column size mismatch");
         let offset = col.offset();
 
-        array::from_fn(|i| self.program_evals[offset + i][0].clone())
+        array::from_fn(|i| self.program_evals[offset + i].clone())
     }
 }
 

tests/testing-framework/src/lib.rs

@@ -203,7 +203,7 @@ mod test {
     fn test_prove_fib() {
         let elfs = compile_multi("examples/src/bin/fib", &["-C opt-level=3"], &HOME_PATH);
         let (view, execution_trace) =
-            k_trace(elfs[0].clone(), &[], &[], &[], K).expect("error generating trace");
+            k_trace(elfs[0].clone(), &[1, 2, 3], &[], &[], K).expect("error generating trace");
         let proof = prove(&execution_trace, &view).unwrap();
         verify(proof, &view).unwrap();
     }