Struct for PreProcessed Columns ID

crates/prover/src/constraint_framework/component.rs

@@ -10,6 +10,7 @@ use tracing::{span, Level};
 
 use super::cpu_domain::CpuDomainEvaluator;
 use super::logup::LogupSums;
+use super::preprocessed_columns::PreProcessedColumnId;
 use super::{
     EvalAtRow, InfoEvaluator, PointEvaluator, SimdDomainEvaluator, PREPROCESSED_TRACE_IDX,
 };
@@ -47,8 +48,7 @@ pub struct TraceLocationAllocator {
     /// Mapping of tree index to next available column offset.
     next_tree_offsets: TreeVec<usize>,
     /// Mapping of preprocessed columns to their index.
-    // TODO(Gali): Change Vec type to struct PreProcessedColumnId {pub id: String}.
-    preprocessed_columns: Vec<String>,
+    preprocessed_columns: Vec<PreProcessedColumnId>,
     /// Controls whether the preprocessed columns are dynamic or static (default=Dynamic).
     preprocessed_columns_allocation_mode: PreprocessedColumnsAllocationMode,
 }
@@ -80,7 +80,7 @@ impl TraceLocationAllocator {
     }
 
     /// Create a new `TraceLocationAllocator` with fixed preprocessed columns setup.
-    pub fn new_with_preproccessed_columns(preprocessed_columns: &[String]) -> Self {
+    pub fn new_with_preproccessed_columns(preprocessed_columns: &[PreProcessedColumnId]) -> Self {
         assert!(
             preprocessed_columns.iter().all_unique(),
             "Duplicate preprocessed columns are not allowed!"
@@ -92,14 +92,16 @@ impl TraceLocationAllocator {
         }
     }
 
-    pub const fn preprocessed_columns(&self) -> &Vec<String> {
+    pub const fn preprocessed_columns(&self) -> &Vec<PreProcessedColumnId> {
         &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: &[String]) {
+    // TODO(Gali): Change to only validating that this is a permutation (not necessarily the same
+    // order).
+    pub fn validate_preprocessed_columns(&self, preprocessed_columns: &[PreProcessedColumnId]) {
         assert_eq!(self.preprocessed_columns, preprocessed_columns);
     }
 }
@@ -142,7 +144,7 @@ impl<E: FrameworkEval> FrameworkComponent<E> {
                 if let Some(pos) = location_allocator
                     .preprocessed_columns
                     .iter()
-                    .position(|x| x == col)
+                    .position(|x| x.id == col.id)
                 {
                     pos
                 } else {

crates/prover/src/constraint_framework/expr/evaluator.rs

@@ -2,6 +2,7 @@ use num_traits::Zero;
 
 use super::{BaseExpr, ExtExpr};
 use crate::constraint_framework::expr::ColumnExpr;
+use crate::constraint_framework::preprocessed_columns::PreProcessedColumnId;
 use crate::constraint_framework::{EvalAtRow, Relation, RelationEntry, INTERACTION_TRACE_IDX};
 use crate::core::fields::m31;
 use crate::core::lookups::utils::Fraction;
@@ -173,8 +174,8 @@ impl EvalAtRow for ExprEvaluator {
         intermediate
     }
 
-    fn get_preprocessed_column(&mut self, column: String) -> Self::F {
-        BaseExpr::Param(column)
+    fn get_preprocessed_column(&mut self, column: PreProcessedColumnId) -> Self::F {
+        BaseExpr::Param(column.id)
     }
 
     crate::constraint_framework::logup_proxy!();

crates/prover/src/constraint_framework/info.rs

@@ -6,6 +6,7 @@ use std::rc::Rc;
 use num_traits::{One, Zero};
 
 use super::logup::{LogupAtRow, LogupSums};
+use super::preprocessed_columns::PreProcessedColumnId;
 use super::{EvalAtRow, INTERACTION_TRACE_IDX};
 use crate::constraint_framework::PREPROCESSED_TRACE_IDX;
 use crate::core::fields::m31::BaseField;
@@ -21,13 +22,16 @@ use crate::core::pcs::TreeVec;
 pub struct InfoEvaluator {
     pub mask_offsets: TreeVec<Vec<Vec<isize>>>,
     pub n_constraints: usize,
-    // TODO(Gali): Change Vec type to struct PreProcessedColumnId {pub id: String}.
-    pub preprocessed_columns: Vec<String>,
+    pub preprocessed_columns: Vec<PreProcessedColumnId>,
     pub logup: LogupAtRow<Self>,
     pub arithmetic_counts: ArithmeticCounts,
 }
 impl InfoEvaluator {
-    pub fn new(log_size: u32, preprocessed_columns: Vec<String>, logup_sums: LogupSums) -> Self {
+    pub fn new(
+        log_size: u32,
+        preprocessed_columns: Vec<PreProcessedColumnId>,
+        logup_sums: LogupSums,
+    ) -> Self {
         Self {
             mask_offsets: Default::default(),
             n_constraints: Default::default(),
@@ -66,7 +70,7 @@ impl EvalAtRow for InfoEvaluator {
         array::from_fn(|_| FieldCounter::one())
     }
 
-    fn get_preprocessed_column(&mut self, column: String) -> Self::F {
+    fn get_preprocessed_column(&mut self, column: PreProcessedColumnId) -> Self::F {
         self.preprocessed_columns.push(column);
         FieldCounter::one()
     }

crates/prover/src/constraint_framework/mod.rs

@@ -19,6 +19,7 @@ pub use component::{FrameworkComponent, FrameworkEval, TraceLocationAllocator};
 pub use info::InfoEvaluator;
 use num_traits::{One, Zero};
 pub use point::PointEvaluator;
+use preprocessed_columns::PreProcessedColumnId;
 pub use simd_domain::SimdDomainEvaluator;
 
 use crate::core::fields::m31::BaseField;
@@ -86,7 +87,7 @@ pub trait EvalAtRow {
         mask_item
     }
 
-    fn get_preprocessed_column(&mut self, _column: String) -> Self::F {
+    fn get_preprocessed_column(&mut self, _column: PreProcessedColumnId) -> Self::F {
         let [mask_item] = self.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
         mask_item
     }

crates/prover/src/constraint_framework/preprocessed_columns.rs

@@ -9,6 +9,13 @@ use crate::core::fields::m31::BaseField;
 use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
 use crate::core::poly::BitReversedOrder;
 
+/// Used for comparing preprocessed columns.
+/// Column IDs must be unique in a given context.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct PreProcessedColumnId {
+    pub id: String,
+}
+
 /// A column with `1` at the first position, and `0` elsewhere.
 #[derive(Debug, Clone)]
 pub struct IsFirst {
@@ -42,8 +49,10 @@ impl IsFirst {
         CircleEvaluation::new(CanonicCoset::new(self.log_size).circle_domain(), col)
     }
 
-    pub fn id(&self) -> String {
-        format!("preprocessed_is_first_{}", self.log_size).to_string()
+    pub fn id(&self) -> PreProcessedColumnId {
+        PreProcessedColumnId {
+            id: format!("preprocessed_is_first_{}", self.log_size).to_string(),
+        }
     }
 }
 

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

@@ -9,7 +9,7 @@ use super::preprocessed_columns::XorTable;
 use super::round::{blake_round_info, BlakeRoundComponent, BlakeRoundEval};
 use super::scheduler::{BlakeSchedulerComponent, BlakeSchedulerEval};
 use super::xor_table::{xor12, xor4, xor7, xor8, xor9};
-use crate::constraint_framework::preprocessed_columns::IsFirst;
+use crate::constraint_framework::preprocessed_columns::{IsFirst, PreProcessedColumnId};
 use crate::constraint_framework::{TraceLocationAllocator, PREPROCESSED_TRACE_IDX};
 use crate::core::air::{Component, ComponentProver};
 use crate::core::backend::simd::m31::LOG_N_LANES;
@@ -27,7 +27,7 @@ use crate::examples::blake::{
     round, xor_table, BlakeXorElements, XorAccums, N_ROUNDS, ROUND_LOG_SPLIT,
 };
 
-fn preprocessed_xor_columns() -> [String; 20] {
+fn preprocessed_xor_columns() -> [PreProcessedColumnId; 20] {
     [
         XorTable::new(12, 4, 0).id(),
         XorTable::new(12, 4, 1).id(),
@@ -188,7 +188,7 @@ impl BlakeComponents {
         let log_size = stmt0.log_size;
 
         let scheduler_is_first_column = IsFirst::new(log_size).id();
-        let blake_round_is_first_columns_iter: Vec<String> = ROUND_LOG_SPLIT
+        let blake_round_is_first_columns_iter: Vec<PreProcessedColumnId> = ROUND_LOG_SPLIT
             .iter()
             .map(|l| IsFirst::new(log_size + l).id())
             .collect_vec();

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

@@ -1,6 +1,6 @@
 use tracing::{span, Level};
 
-use crate::constraint_framework::preprocessed_columns::IsFirst;
+use crate::constraint_framework::preprocessed_columns::{IsFirst, PreProcessedColumnId};
 use crate::core::backend::simd::column::BaseColumn;
 use crate::core::backend::simd::SimdBackend;
 use crate::core::fields::m31::BaseField;
@@ -24,11 +24,13 @@ impl XorTable {
         }
     }
 
-    pub fn id(&self) -> String {
-        format!(
-            "preprocessed_xor_table_{}_{}_{}",
-            self.n_bits, self.n_expand_bits, self.index_in_table
-        )
+    pub fn id(&self) -> PreProcessedColumnId {
+        PreProcessedColumnId {
+            id: format!(
+                "preprocessed_xor_table_{}_{}_{}",
+                self.n_bits, self.n_expand_bits, self.index_in_table
+            ),
+        }
     }
 
     pub const fn limb_bits(&self) -> u32 {

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

@@ -3,7 +3,7 @@ use num_traits::One;
 use tracing::{span, Level};
 
 use crate::constraint_framework::logup::{ClaimedPrefixSum, LogupTraceGenerator, LookupElements};
-use crate::constraint_framework::preprocessed_columns::IsFirst;
+use crate::constraint_framework::preprocessed_columns::{IsFirst, PreProcessedColumnId};
 use crate::constraint_framework::{
     assert_constraints, relation, EvalAtRow, FrameworkComponent, FrameworkEval, RelationEntry,
     TraceLocationAllocator,
@@ -280,8 +280,10 @@ impl Plonk {
         Self { name }
     }
 
-    pub fn id(&self) -> String {
-        format!("preprocessed_plonk_{}", self.name)
+    pub fn id(&self) -> PreProcessedColumnId {
+        PreProcessedColumnId {
+            id: format!("preprocessed_plonk_{}", self.name),
+        }
     }
 }
 

crates/prover/src/examples/xor/gkr_lookups/preprocessed_columns.rs

@@ -1,5 +1,6 @@
 use num_traits::One;
 
+use crate::constraint_framework::preprocessed_columns::PreProcessedColumnId;
 use crate::core::backend::simd::SimdBackend;
 use crate::core::backend::{Col, Column};
 use crate::core::fields::m31::BaseField;
@@ -38,10 +39,12 @@ impl IsStepWithOffset {
         CircleEvaluation::new(CanonicCoset::new(self.log_size).circle_domain(), col)
     }
 
-    pub fn id(&self) -> String {
-        format!(
-            "preprocessed_is_step_with_offset_{}_{}_{}",
-            self.log_size, self.log_step, self.offset
-        )
+    pub fn id(&self) -> PreProcessedColumnId {
+        PreProcessedColumnId {
+            id: format!(
+                "preprocessed_is_step_with_offset_{}_{}_{}",
+                self.log_size, self.log_step, self.offset
+            ),
+        }
     }
 }