Remove PreProcessedColumn Enum

crates/prover/src/constraint_framework/component.rs

@@ -1,9 +1,7 @@
 use std::borrow::Cow;
-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")]
@@ -12,7 +10,6 @@ use tracing::{span, Level};
 
 use super::cpu_domain::CpuDomainEvaluator;
 use super::logup::LogupSums;
-use super::preprocessed_columns::PreprocessedColumn;
 use super::{
     EvalAtRow, InfoEvaluator, PointEvaluator, SimdDomainEvaluator, PREPROCESSED_TRACE_IDX,
 };
@@ -50,7 +47,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<Rc<dyn PreprocessedColumn>, usize>,
+    preprocessed_columns: Vec<String>,
     /// Controls whether the preprocessed columns are dynamic or static (default=Dynamic).
     preprocessed_columns_allocation_mode: PreprocessedColumnsAllocationMode,
 }
@@ -82,38 +79,23 @@ impl TraceLocationAllocator {
     }
 
     /// Create a new `TraceLocationAllocator` with fixed preprocessed columns setup.
-    pub fn new_with_preproccessed_columns(
-        preprocessed_columns: &[Rc<dyn PreprocessedColumn>],
-    ) -> Self {
+    pub fn new_with_preproccessed_columns(preprocessed_columns: &[String]) -> Self {
         Self {
             next_tree_offsets: Default::default(),
-            preprocessed_columns: preprocessed_columns
-                .iter()
-                .enumerate()
-                .map(|(i, col)| (col.clone(), i))
-                .collect(),
+            preprocessed_columns: preprocessed_columns.to_vec(),
             preprocessed_columns_allocation_mode: PreprocessedColumnsAllocationMode::Static,
         }
     }
 
-    pub const fn preprocessed_columns(&self) -> &HashMap<Rc<dyn PreprocessedColumn>, usize> {
+    pub const fn preprocessed_columns(&self) -> &Vec<String> {
         &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: &[Rc<dyn PreprocessedColumn>],
-    ) {
-        assert_eq!(preprocessed_columns.len(), self.preprocessed_columns.len());
-        for (column, idx) in self.preprocessed_columns.iter() {
-            let preprocessed_column = preprocessed_columns
-                .get(*idx)
-                .expect("Preprocessed column is missing from preprocessed_columns");
-            assert_eq!(column.id(), preprocessed_column.id());
-        }
+    pub fn validate_preprocessed_columns(&self, preprocessed_columns: &[String]) {
+        assert_eq!(self.preprocessed_columns, preprocessed_columns);
     }
 }
 
@@ -152,22 +134,25 @@ impl<E: FrameworkEval> FrameworkComponent<E> {
             .iter()
             .map(|col| {
                 let next_column = location_allocator.preprocessed_columns.len();
-                *location_allocator
+                if let Some(pos) = location_allocator
                     .preprocessed_columns
-                    .entry(col.clone())
-                    .or_insert_with(|| {
-                        if matches!(
-                            location_allocator.preprocessed_columns_allocation_mode,
-                            PreprocessedColumnsAllocationMode::Static
-                        ) {
-                            panic!(
-                                "Preprocessed column {:?} is missing from static alloction",
-                                col
-                            );
-                        }
-
-                        next_column
-                    })
+                    .iter()
+                    .position(|x| x == col)
+                {
+                    pos
+                } else {
+                    if matches!(
+                        location_allocator.preprocessed_columns_allocation_mode,
+                        PreprocessedColumnsAllocationMode::Static
+                    ) {
+                        panic!(
+                            "Preprocessed column {:?} is missing from static allocation",
+                            col
+                        );
+                    }
+                    location_allocator.preprocessed_columns.push(col.clone());
+                    next_column
+                }
             })
             .collect();
         Self {

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

@@ -1,10 +1,7 @@
-use std::rc::Rc;
-
 use num_traits::Zero;
 
 use super::{BaseExpr, ExtExpr};
 use crate::constraint_framework::expr::ColumnExpr;
-use crate::constraint_framework::preprocessed_columns::PreprocessedColumn;
 use crate::constraint_framework::{EvalAtRow, Relation, RelationEntry, INTERACTION_TRACE_IDX};
 use crate::core::fields::m31;
 use crate::core::lookups::utils::Fraction;
@@ -176,8 +173,8 @@ impl EvalAtRow for ExprEvaluator {
         intermediate
     }
 
-    fn get_preprocessed_column(&mut self, column: Rc<dyn PreprocessedColumn>) -> Self::F {
-        BaseExpr::Param(column.name().to_string())
+    fn get_preprocessed_column(&mut self, column: String) -> Self::F {
+        BaseExpr::Param(column)
     }
 
     crate::constraint_framework::logup_proxy!();
@@ -210,7 +207,7 @@ mod tests {
 \
         let constraint_1 = (QM31Impl::from_partial_evals([trace_2_column_3_offset_0, trace_2_column_4_offset_0, trace_2_column_5_offset_0, trace_2_column_6_offset_0]) \
             - (QM31Impl::from_partial_evals([trace_2_column_3_offset_neg_1, trace_2_column_4_offset_neg_1, trace_2_column_5_offset_neg_1, trace_2_column_6_offset_neg_1]) \
-                - ((total_sum) * (preprocessed_is_first)))) \
+                - ((total_sum) * (preprocessed_is_first_16)))) \
             * (intermediate1) \
             - (qm31(1, 0, 0, 0));"
             .to_string();

crates/prover/src/constraint_framework/info.rs

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

crates/prover/src/constraint_framework/mod.rs

@@ -13,14 +13,12 @@ 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};
 pub use info::InfoEvaluator;
 use num_traits::{One, Zero};
 pub use point::PointEvaluator;
-use preprocessed_columns::PreprocessedColumn;
 pub use simd_domain::SimdDomainEvaluator;
 
 use crate::core::fields::m31::BaseField;
@@ -88,7 +86,7 @@ pub trait EvalAtRow {
         mask_item
     }
 
-    fn get_preprocessed_column(&mut self, _column: Rc<dyn PreprocessedColumn>) -> Self::F {
+    fn get_preprocessed_column(&mut self, _column: String) -> Self::F {
         let [mask_item] = self.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
         mask_item
     }
@@ -173,11 +171,12 @@ 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(std::rc::Rc::new(
+                self.logup.is_first = self.get_preprocessed_column(
                     crate::constraint_framework::preprocessed_columns::IsFirst::new(
                         self.logup.log_size,
-                    ),
-                ));
+                    )
+                    .id(),
+                );
                 self.logup.is_finalized = false;
             }
             self.logup.fracs.push(fraction.clone());