Add get_preprocessed_column api. (#867)

crates/prover/src/constraint_framework/mod.rs

@@ -18,6 +18,7 @@ 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;
@@ -75,6 +76,11 @@ pub trait EvalAtRow {
         mask_item
     }
 
+    fn get_preprocessed_column(&mut self, _column: PreprocessedColumn) -> Self::F {
+        let [mask_item] = self.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
+        mask_item
+    }
+
     /// Returns the mask values of the given offsets for the next column in the interaction.
     fn next_interaction_mask<const N: usize>(
         &mut self,

crates/prover/src/constraint_framework/preprocessed_columns.rs

@@ -6,6 +6,14 @@ use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
 use crate::core::poly::BitReversedOrder;
 use crate::core::utils::{bit_reverse_index, coset_index_to_circle_domain_index};
 
+// TODO(ilya): Where should this enum be placed?
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum PreprocessedColumn {
+    XorTable(u32, u32, usize),
+    IsFirst(u32),
+    Plonk(usize),
+}
+
 /// Generates a column with a single one at the first position, and zeros elsewhere.
 pub fn gen_is_first<B: Backend>(log_size: u32) -> CircleEvaluation<B, BaseField, BitReversedOrder> {
     let mut col = Col::<B, BaseField>::zeros(1 << log_size);

crates/prover/src/examples/blake/round/mod.rs

@@ -6,9 +6,8 @@ use num_traits::Zero;
 
 use super::{BlakeXorElements, N_ROUND_INPUT_FELTS};
 use crate::constraint_framework::logup::{LogupAtRow, LookupElements};
-use crate::constraint_framework::{
-    EvalAtRow, FrameworkComponent, FrameworkEval, InfoEvaluator, PREPROCESSED_TRACE_IDX,
-};
+use crate::constraint_framework::preprocessed_columns::PreprocessedColumn;
+use crate::constraint_framework::{EvalAtRow, FrameworkComponent, FrameworkEval, InfoEvaluator};
 use crate::core::fields::qm31::SecureField;
 
 pub type BlakeRoundComponent = FrameworkComponent<BlakeRoundEval>;
@@ -32,7 +31,7 @@ impl FrameworkEval for BlakeRoundEval {
         self.log_size + 1
     }
     fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
-        let [is_first] = eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
+        let is_first = eval.get_preprocessed_column(PreprocessedColumn::IsFirst(self.log_size()));
         let blake_eval = constraints::BlakeRoundEval {
             eval,
             xor_lookup_elements: &self.xor_lookup_elements,

crates/prover/src/examples/blake/scheduler/mod.rs

@@ -8,9 +8,9 @@ use num_traits::Zero;
 use super::round::RoundElements;
 use super::N_ROUND_INPUT_FELTS;
 use crate::constraint_framework::logup::{LogupAtRow, LookupElements};
+use crate::constraint_framework::preprocessed_columns::PreprocessedColumn;
 use crate::constraint_framework::{
     EvalAtRow, FrameworkComponent, FrameworkEval, InfoEvaluator, INTERACTION_TRACE_IDX,
-    PREPROCESSED_TRACE_IDX,
 };
 use crate::core::fields::qm31::SecureField;
 
@@ -32,7 +32,7 @@ impl FrameworkEval for BlakeSchedulerEval {
         self.log_size + 1
     }
     fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
-        let [is_first] = eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
+        let is_first = eval.get_preprocessed_column(PreprocessedColumn::IsFirst(self.log_size()));
         eval_blake_scheduler_constraints(
             &mut eval,
             &self.blake_lookup_elements,

crates/prover/src/examples/blake/xor_table/constraints.rs

@@ -2,7 +2,8 @@ use itertools::Itertools;
 
 use super::{limb_bits, XorElements};
 use crate::constraint_framework::logup::{LogupAtRow, LookupElements};
-use crate::constraint_framework::{EvalAtRow, PREPROCESSED_TRACE_IDX};
+use crate::constraint_framework::preprocessed_columns::PreprocessedColumn;
+use crate::constraint_framework::EvalAtRow;
 use crate::core::fields::m31::BaseField;
 use crate::core::lookups::utils::Fraction;
 
@@ -19,9 +20,17 @@ impl<'a, E: EvalAtRow, const ELEM_BITS: u32, const EXPAND_BITS: u32>
         // al, bl are the constant columns for the inputs: All pairs of elements in [0,
         // 2^LIMB_BITS).
         // cl is the constant column for the xor: al ^ bl.
-        let [al] = self.eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
-        let [bl] = self.eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
-        let [cl] = self.eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
+        let al = self
+            .eval
+            .get_preprocessed_column(PreprocessedColumn::XorTable(ELEM_BITS, EXPAND_BITS, 0));
+
+        let bl = self
+            .eval
+            .get_preprocessed_column(PreprocessedColumn::XorTable(ELEM_BITS, EXPAND_BITS, 1));
+
+        let cl = self
+            .eval
+            .get_preprocessed_column(PreprocessedColumn::XorTable(ELEM_BITS, EXPAND_BITS, 2));
 
         let frac_chunks = (0..(1 << (2 * EXPAND_BITS)))
             .map(|i| {

crates/prover/src/examples/blake/xor_table/mod.rs

@@ -20,6 +20,7 @@ use num_traits::Zero;
 pub use r#gen::{generate_constant_trace, generate_interaction_trace, generate_trace};
 
 use crate::constraint_framework::logup::{LogupAtRow, LookupElements};
+use crate::constraint_framework::preprocessed_columns::PreprocessedColumn;
 use crate::constraint_framework::{
     EvalAtRow, FrameworkComponent, FrameworkEval, InfoEvaluator, INTERACTION_TRACE_IDX,
     PREPROCESSED_TRACE_IDX,
@@ -106,7 +107,7 @@ impl<const ELEM_BITS: u32, const EXPAND_BITS: u32> FrameworkEval
         column_bits::<ELEM_BITS, EXPAND_BITS>() + 1
     }
     fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
-        let [is_first] = eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
+        let is_first = eval.get_preprocessed_column(PreprocessedColumn::IsFirst(self.log_size()));
         let xor_eval = constraints::XorTableEval::<'_, _, ELEM_BITS, EXPAND_BITS> {
             eval,
             lookup_elements: &self.lookup_elements,

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

@@ -5,10 +5,10 @@ use tracing::{span, Level};
 use crate::constraint_framework::logup::{
     ClaimedPrefixSum, LogupAtRow, LogupTraceGenerator, LookupElements,
 };
-use crate::constraint_framework::preprocessed_columns::gen_is_first;
+use crate::constraint_framework::preprocessed_columns::{gen_is_first, PreprocessedColumn};
 use crate::constraint_framework::{
     assert_constraints, EvalAtRow, FrameworkComponent, FrameworkEval, TraceLocationAllocator,
-    INTERACTION_TRACE_IDX, PREPROCESSED_TRACE_IDX,
+    INTERACTION_TRACE_IDX,
 };
 use crate::core::backend::simd::column::BaseColumn;
 use crate::core::backend::simd::m31::LOG_N_LANES;
@@ -49,20 +49,20 @@ impl FrameworkEval for PlonkEval {
     }
 
     fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
-        let [is_first] = eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
+        let is_first = eval.get_preprocessed_column(PreprocessedColumn::IsFirst(self.log_size()));
         let mut logup = LogupAtRow::<_>::new(
             INTERACTION_TRACE_IDX,
             self.total_sum,
             Some(self.claimed_sum),
             is_first,
         );
 
-        let [a_wire] = eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
-        let [b_wire] = eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
+        let a_wire = eval.get_preprocessed_column(PreprocessedColumn::Plonk(0));
+        let b_wire = eval.get_preprocessed_column(PreprocessedColumn::Plonk(1));
         // Note: c_wire could also be implicit: (self.eval.point() - M31_CIRCLE_GEN.into_ef()).x.
         //   A constant column is easier though.
-        let [c_wire] = eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
-        let [op] = eval.next_interaction_mask(PREPROCESSED_TRACE_IDX, [0]);
+        let c_wire = eval.get_preprocessed_column(PreprocessedColumn::Plonk(2));
+        let op = eval.get_preprocessed_column(PreprocessedColumn::Plonk(3));
 
         let mult = eval.next_trace_mask();
         let a_val = eval.next_trace_mask();

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

@@ -1,6 +1,7 @@
 use num_traits::{One, Zero};
 
 use crate::constraint_framework::logup::{ClaimedPrefixSum, LogupAtRow, LookupElements};
+use crate::constraint_framework::preprocessed_columns::PreprocessedColumn;
 use crate::constraint_framework::{
     EvalAtRow, FrameworkComponent, FrameworkEval, InfoEvaluator, INTERACTION_TRACE_IDX,
 };
@@ -41,7 +42,7 @@ impl<const COORDINATE: usize> FrameworkEval for StateTransitionEval<COORDINATE>
         self.log_n_rows + LOG_CONSTRAINT_DEGREE
     }
     fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
-        let [is_first] = eval.next_interaction_mask(2, [0]);
+        let is_first = eval.get_preprocessed_column(PreprocessedColumn::IsFirst(self.log_size()));
         let mut logup: LogupAtRow<E> = LogupAtRow::new(
             INTERACTION_TRACE_IDX,
             self.total_sum,