Use PreProcessedColumn Trait Instead of Enum

crates/prover/src/constraint_framework/component.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use std::fmt::{self, Display, Formatter};
 use std::iter::zip;
 use std::ops::Deref;
+use std::rc::Rc;
 
 use itertools::Itertools;
 #[cfg(feature = "parallel")]
@@ -49,7 +50,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.
-    preprocessed_columns: HashMap<PreprocessedColumn, usize>,
+    preprocessed_columns: HashMap<Rc<dyn PreprocessedColumn>, usize>,
     /// Controls whether the preprocessed columns are dynamic or static (default=Dynamic).
     preprocessed_columns_allocation_mode: PreprocessedColumnsAllocationMode,
 }
@@ -81,30 +82,37 @@ impl TraceLocationAllocator {
     }
 
     /// Create a new `TraceLocationAllocator` with fixed preprocessed columns setup.
-    pub fn new_with_preproccessed_columns(preprocessed_columns: &[PreprocessedColumn]) -> Self {
+    pub fn new_with_preproccessed_columns(
+        preprocessed_columns: &[Rc<dyn PreprocessedColumn>],
+    ) -> Self {
         Self {
             next_tree_offsets: Default::default(),
             preprocessed_columns: preprocessed_columns
                 .iter()
                 .enumerate()
-                .map(|(i, &col)| (col, i))
+                .map(|(i, col)| (col.clone(), i))
                 .collect(),
             preprocessed_columns_allocation_mode: PreprocessedColumnsAllocationMode::Static,
         }
     }
 
-    pub const fn preprocessed_columns(&self) -> &HashMap<PreprocessedColumn, usize> {
+    pub const fn preprocessed_columns(&self) -> &HashMap<Rc<dyn 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]) {
+    pub fn validate_preprocessed_columns(
+        &self,
+        preprocessed_columns: &[Rc<dyn 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));
+            let preprocessed_column = preprocessed_columns
+                .get(*idx)
+                .expect("Preprocessed column is missing from preprocessed_columns");
+            assert_eq!(column.id(), preprocessed_column.id());
         }
     }
 }
@@ -146,7 +154,7 @@ impl<E: FrameworkEval> FrameworkComponent<E> {
                 let next_column = location_allocator.preprocessed_columns.len();
                 *location_allocator
                     .preprocessed_columns
-                    .entry(*col)
+                    .entry(col.clone())
                     .or_insert_with(|| {
                         if matches!(
                             location_allocator.preprocessed_columns_allocation_mode,

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

@@ -1,3 +1,5 @@
+use std::rc::Rc;
+
 use num_traits::Zero;
 
 use super::{BaseExpr, ExtExpr};
@@ -174,7 +176,7 @@ impl EvalAtRow for ExprEvaluator {
         intermediate
     }
 
-    fn get_preprocessed_column(&mut self, column: PreprocessedColumn) -> Self::F {
+    fn get_preprocessed_column(&mut self, column: Rc<dyn PreprocessedColumn>) -> Self::F {
         BaseExpr::Param(column.name().to_string())
     }
 

crates/prover/src/constraint_framework/info.rs

@@ -22,14 +22,14 @@ use crate::core::pcs::TreeVec;
 pub struct InfoEvaluator {
     pub mask_offsets: TreeVec<Vec<Vec<isize>>>,
     pub n_constraints: usize,
-    pub preprocessed_columns: Vec<PreprocessedColumn>,
+    pub preprocessed_columns: Vec<Rc<dyn PreprocessedColumn>>,
     pub logup: LogupAtRow<Self>,
     pub arithmetic_counts: ArithmeticCounts,
 }
 impl InfoEvaluator {
     pub fn new(
         log_size: u32,
-        preprocessed_columns: Vec<PreprocessedColumn>,
+        preprocessed_columns: Vec<Rc<dyn PreprocessedColumn>>,
         logup_sums: LogupSums,
     ) -> Self {
         Self {
@@ -70,7 +70,7 @@ impl EvalAtRow for InfoEvaluator {
         array::from_fn(|_| FieldCounter::one())
     }
 
-    fn get_preprocessed_column(&mut self, column: PreprocessedColumn) -> Self::F {
+    fn get_preprocessed_column(&mut self, column: Rc<dyn PreprocessedColumn>) -> Self::F {
         self.preprocessed_columns.push(column);
         FieldCounter::one()
     }

crates/prover/src/constraint_framework/logup.rs

@@ -52,7 +52,7 @@ pub struct LogupAtRow<E: EvalAtRow> {
     pub fracs: Vec<Fraction<E::EF, E::EF>>,
     pub is_finalized: bool,
     /// The value of the `is_first` constant column at current row.
-    /// See [`super::preprocessed_columns::gen_is_first()`].
+    /// See [`super::preprocessed_columns::IsFirst`].
     pub is_first: E::F,
     pub log_size: u32,
 }

crates/prover/src/constraint_framework/mod.rs

@@ -13,6 +13,7 @@ mod simd_domain;
 use std::array;
 use std::fmt::Debug;
 use std::ops::{Add, AddAssign, Mul, Neg, Sub};
+use std::rc::Rc;
 
 pub use assert::{assert_constraints, AssertEvaluator};
 pub use component::{FrameworkComponent, FrameworkEval, TraceLocationAllocator};
@@ -87,7 +88,7 @@ pub trait EvalAtRow {
         mask_item
     }
 
-    fn get_preprocessed_column(&mut self, _column: PreprocessedColumn) -> Self::F {
+    fn get_preprocessed_column(&mut self, _column: Rc<dyn PreprocessedColumn>) -> Self::F {
         let [mask_item] = self.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
         mask_item
     }
@@ -172,11 +173,11 @@ macro_rules! logup_proxy {
     () => {
         fn write_logup_frac(&mut self, fraction: Fraction<Self::EF, Self::EF>) {
             if self.logup.fracs.is_empty() {
-                self.logup.is_first = self.get_preprocessed_column(
-                    crate::constraint_framework::preprocessed_columns::PreprocessedColumn::IsFirst(
+                self.logup.is_first = self.get_preprocessed_column(std::rc::Rc::new(
+                    crate::constraint_framework::preprocessed_columns::IsFirst::new(
                         self.logup.log_size,
                     ),
-                );
+                ));
                 self.logup.is_finalized = false;
             }
             self.logup.fracs.push(fraction.clone());

crates/prover/src/constraint_framework/preprocessed_columns.rs

@@ -18,28 +18,27 @@ const SIMD_ENUMERATION_0: PackedM31 = unsafe {
     ]))
 };
 
-// TODO(Gali): Rename to PrerocessedColumn.
-pub trait PreprocessedColumnTrait: Debug {
+pub trait PreprocessedColumn: Debug {
     fn name(&self) -> &'static str;
     /// Used for comparing preprocessed columns.
     /// Column IDs must be unique in a given context.
     fn id(&self) -> String;
     fn log_size(&self) -> u32;
 }
-impl PartialEq for dyn PreprocessedColumnTrait {
+impl PartialEq for dyn PreprocessedColumn {
     fn eq(&self, other: &Self) -> bool {
         self.id() == other.id()
     }
 }
-impl Eq for dyn PreprocessedColumnTrait {}
-impl Hash for dyn PreprocessedColumnTrait {
+impl Eq for dyn PreprocessedColumn {}
+impl Hash for dyn PreprocessedColumn {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
         self.id().hash(state);
     }
 }
 
 /// A column with `1` at the first position, and `0` elsewhere.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct IsFirst {
     pub log_size: u32,
 }
@@ -71,7 +70,7 @@ impl IsFirst {
         CircleEvaluation::new(CanonicCoset::new(self.log_size).circle_domain(), col)
     }
 }
-impl PreprocessedColumnTrait for IsFirst {
+impl PreprocessedColumn for IsFirst {
     fn name(&self) -> &'static str {
         "preprocessed_is_first"
     }
@@ -85,10 +84,9 @@ impl PreprocessedColumnTrait for IsFirst {
     }
 }
 
-// TODO(ilya): Where should this enum be placed?
-// TODO(Gali): Add documentation for the rest of the variants.
+// TODO(Gali): Remove Enum.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub enum PreprocessedColumn {
+pub enum PreprocessedColumnEnum {
     /// A column with `1` at the first position, and `0` elsewhere.
     IsFirst(u32),
     Plonk(usize),
@@ -97,29 +95,29 @@ pub enum PreprocessedColumn {
     XorTable(u32, u32, usize),
 }
 
-impl PreprocessedColumn {
+impl PreprocessedColumnEnum {
     pub const fn name(&self) -> &'static str {
         match self {
-            PreprocessedColumn::IsFirst(_) => "preprocessed_is_first",
-            PreprocessedColumn::Plonk(_) => "preprocessed_plonk",
-            PreprocessedColumn::Seq(_) => "preprocessed_seq",
-            PreprocessedColumn::XorTable(..) => "preprocessed_xor_table",
+            PreprocessedColumnEnum::IsFirst(_) => "preprocessed_is_first",
+            PreprocessedColumnEnum::Plonk(_) => "preprocessed_plonk",
+            PreprocessedColumnEnum::Seq(_) => "preprocessed_seq",
+            PreprocessedColumnEnum::XorTable(..) => "preprocessed_xor_table",
         }
     }
 
     pub fn log_size(&self) -> u32 {
         match self {
-            PreprocessedColumn::IsFirst(log_size) => *log_size,
-            PreprocessedColumn::Seq(log_size) => *log_size,
-            PreprocessedColumn::XorTable(log_size, ..) => *log_size,
-            PreprocessedColumn::Plonk(_) => unimplemented!(),
+            PreprocessedColumnEnum::IsFirst(log_size) => *log_size,
+            PreprocessedColumnEnum::Seq(log_size) => *log_size,
+            PreprocessedColumnEnum::XorTable(log_size, ..) => *log_size,
+            PreprocessedColumnEnum::Plonk(_) => unimplemented!(),
         }
     }
 
     /// Returns the values of the column at the given row.
     pub fn packed_at(&self, vec_row: usize) -> PackedM31 {
         match self {
-            PreprocessedColumn::IsFirst(log_size) => {
+            PreprocessedColumnEnum::IsFirst(log_size) => {
                 assert!(vec_row < (1 << log_size) / N_LANES);
                 if vec_row == 0 {
                     unsafe {
@@ -135,7 +133,7 @@ impl PreprocessedColumn {
                     PackedM31::zero()
                 }
             }
-            PreprocessedColumn::Seq(log_size) => {
+            PreprocessedColumnEnum::Seq(log_size) => {
                 assert!(vec_row < (1 << log_size) / N_LANES);
                 PackedM31::broadcast(M31::from(vec_row * N_LANES)) + SIMD_ENUMERATION_0
             }
@@ -146,14 +144,14 @@ impl PreprocessedColumn {
 
     /// Generates a column according to the preprocessed column chosen.
     pub fn gen_preprocessed_column<B: Backend>(
-        preprocessed_column: &PreprocessedColumn,
+        preprocessed_column: &PreprocessedColumnEnum,
     ) -> CircleEvaluation<B, BaseField, BitReversedOrder> {
         match preprocessed_column {
-            PreprocessedColumn::IsFirst(log_size) => gen_is_first(*log_size),
-            PreprocessedColumn::Plonk(_) | PreprocessedColumn::XorTable(..) => {
+            PreprocessedColumnEnum::IsFirst(log_size) => gen_is_first(*log_size),
+            PreprocessedColumnEnum::Plonk(_) | PreprocessedColumnEnum::XorTable(..) => {
                 unimplemented!("eval_preprocessed_column: Plonk and XorTable are not supported.")
             }
-            PreprocessedColumn::Seq(log_size) => gen_seq(*log_size),
+            PreprocessedColumnEnum::Seq(log_size) => gen_seq(*log_size),
         }
     }
 }
@@ -195,35 +193,24 @@ pub fn gen_seq<B: Backend>(log_size: u32) -> CircleEvaluation<B, BaseField, BitR
 }
 
 pub fn gen_preprocessed_columns<'a, B: Backend>(
-    columns: impl Iterator<Item = &'a PreprocessedColumn>,
+    columns: impl Iterator<Item = &'a PreprocessedColumnEnum>,
 ) -> Vec<CircleEvaluation<B, BaseField, BitReversedOrder>> {
     columns
-        .map(PreprocessedColumn::gen_preprocessed_column)
+        .map(PreprocessedColumnEnum::gen_preprocessed_column)
         .collect()
 }
 
 #[cfg(test)]
 mod tests {
+    use super::IsFirst;
     use crate::core::backend::simd::m31::N_LANES;
-    use crate::core::backend::simd::SimdBackend;
-    use crate::core::backend::Column;
-    use crate::core::fields::m31::{BaseField, M31};
     const LOG_SIZE: u32 = 8;
 
-    #[test]
-    fn test_gen_seq() {
-        let seq = super::gen_seq::<SimdBackend>(LOG_SIZE);
-
-        for i in 0..(1 << LOG_SIZE) {
-            assert_eq!(seq.at(i), BaseField::from_u32_unchecked(i as u32));
-        }
-    }
-
     // TODO(Gali): Add packed_at tests for xor_table and plonk.
     #[test]
     fn test_packed_at_is_first() {
-        let is_first = super::PreprocessedColumn::IsFirst(LOG_SIZE);
-        let expected_is_first = super::gen_is_first::<SimdBackend>(LOG_SIZE).to_cpu();
+        let is_first = IsFirst::new(LOG_SIZE);
+        let expected_is_first = is_first.gen_column_simd().to_cpu();
 
         for i in 0..(1 << LOG_SIZE) / N_LANES {
             assert_eq!(
@@ -232,17 +219,4 @@ mod tests {
             );
         }
     }
-
-    #[test]
-    fn test_packed_at_seq() {
-        let seq = super::PreprocessedColumn::Seq(LOG_SIZE);
-        let expected_seq: [_; 1 << LOG_SIZE] = std::array::from_fn(|i| M31::from(i as u32));
-
-        let packed_seq = std::array::from_fn::<_, { (1 << LOG_SIZE) / N_LANES }, _>(|i| {
-            seq.packed_at(i).to_array()
-        })
-        .concat();
-
-        assert_eq!(packed_seq, expected_seq);
-    }
 }

crates/prover/src/constraint_framework/relation_tracker.rs

@@ -5,7 +5,6 @@ use itertools::Itertools;
 use num_traits::Zero;
 
 use super::logup::LogupSums;
-use super::preprocessed_columns::PreprocessedColumn;
 use super::{
     Batching, EvalAtRow, FrameworkEval, InfoEvaluator, Relation, RelationEntry,
     TraceLocationAllocator, INTERACTION_TRACE_IDX,
@@ -146,10 +145,6 @@ impl EvalAtRow for RelationTrackerEvaluator<'_> {
         })
     }
 
-    fn get_preprocessed_column(&mut self, column: PreprocessedColumn) -> Self::F {
-        column.packed_at(self.vec_row)
-    }
-
     fn add_constraint<G>(&mut self, _constraint: G) {}
 
     fn combine_ef(_values: [Self::F; SECURE_EXTENSION_DEGREE]) -> Self::EF {