Add gen_preprocessed_columns. (#891)

crates/prover/src/constraint_framework/component.rs

@@ -80,7 +80,7 @@ impl TraceLocationAllocator {
     }
 
     /// Create a new `TraceLocationAllocator` with fixed preprocessed columns setup.
-    pub fn new_with_preproccessed_columnds(preprocessed_columns: &[PreprocessedColumn]) -> Self {
+    pub fn new_with_preproccessed_columns(preprocessed_columns: &[PreprocessedColumn]) -> Self {
         Self {
             next_tree_offsets: Default::default(),
             preprocessed_columns: preprocessed_columns
@@ -92,8 +92,19 @@ impl TraceLocationAllocator {
         }
     }
 
-    pub fn n_preprocessed_columns(&self) -> usize {
-        self.preprocessed_columns.len()
+    pub fn preprocessed_columns(&self) -> &HashMap<PreprocessedColumn, usize> {
+        &self.preprocessed_columns
+    }
+
+    // validates that `self.preprocessed_columns` is consistent with
+    // `preprocessed_columns`.
+    // I.e. preprocessed_columns[i] == self.preprocessed_columns[i].
+    pub fn validate_preprocessed_columns(&self, preprocessed_columns: &[PreprocessedColumn]) {
+        assert_eq!(preprocessed_columns.len(), self.preprocessed_columns.len());
+
+        for (column, idx) in self.preprocessed_columns.iter() {
+            assert_eq!(Some(column), preprocessed_columns.get(*idx));
+        }
     }
 }
 

crates/prover/src/constraint_framework/preprocessed_columns.rs

@@ -43,3 +43,20 @@ pub fn gen_is_step_with_offset<B: Backend>(
 
     CircleEvaluation::new(CanonicCoset::new(log_size).circle_domain(), col)
 }
+
+pub fn gen_preprocessed_column<B: Backend>(
+    preprocessed_column: &PreprocessedColumn,
+) -> CircleEvaluation<B, BaseField, BitReversedOrder> {
+    match preprocessed_column {
+        PreprocessedColumn::IsFirst(log_size) => gen_is_first(*log_size),
+        PreprocessedColumn::Plonk(_) | PreprocessedColumn::XorTable(..) => {
+            unimplemented!("eval_preprocessed_column: Plonk and XorTable are not supported.")
+        }
+    }
+}
+
+pub fn gen_preprocessed_columns<'a, B: Backend>(
+    columns: impl Iterator<Item = &'a PreprocessedColumn>,
+) -> Vec<CircleEvaluation<B, BaseField, BitReversedOrder>> {
+    columns.map(gen_preprocessed_column).collect()
+}

crates/prover/src/examples/blake/air.rs

@@ -169,7 +169,7 @@ impl BlakeComponents {
             .iter()
             .map(|l| PreprocessedColumn::IsFirst(log_size + l));
 
-        let tree_span_provider = &mut TraceLocationAllocator::new_with_preproccessed_columnds(
+        let tree_span_provider = &mut TraceLocationAllocator::new_with_preproccessed_columns(
             &chain!(
                 [scheduler_is_first_column],
                 blake_round_is_first_columns_iter,

crates/prover/src/examples/state_machine/mod.rs

@@ -10,7 +10,9 @@ use components::{
 use gen::{gen_interaction_trace, gen_trace};
 use itertools::{chain, Itertools};
 
-use crate::constraint_framework::preprocessed_columns::gen_is_first;
+use crate::constraint_framework::preprocessed_columns::{
+    gen_preprocessed_columns, PreprocessedColumn,
+};
 use crate::constraint_framework::TraceLocationAllocator;
 use crate::core::backend::simd::m31::LOG_N_LANES;
 use crate::core::backend::simd::SimdBackend;
@@ -54,12 +56,14 @@ pub fn prove_state_machine(
     let mut commitment_scheme =
         CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles);
 
+    let preprocessed_columns = [
+        PreprocessedColumn::IsFirst(x_axis_log_rows),
+        PreprocessedColumn::IsFirst(y_axis_log_rows),
+    ];
+
     // Preprocessed trace.
     let mut tree_builder = commitment_scheme.tree_builder();
-    tree_builder.extend_evals(vec![
-        gen_is_first(x_axis_log_rows),
-        gen_is_first(y_axis_log_rows),
-    ]);
+    tree_builder.extend_evals(gen_preprocessed_columns(preprocessed_columns.iter()));
     tree_builder.commit(channel);
 
     // Trace.
@@ -117,6 +121,9 @@ pub fn prove_state_machine(
         },
         (total_sum_op1, Some((claimed_sum_op1, y_row as usize - 1))),
     );
+
+    tree_span_provider.validate_preprocessed_columns(&preprocessed_columns);
+
     let components = StateMachineComponents {
         component0,
         component1,