From e1cc99438c0f46477a7801dc5cfa57f0cc84ac1d Mon Sep 17 00:00:00 2001 From: Gali Michlevich Date: Tue, 14 Jan 2025 13:11:44 +0200 Subject: [PATCH] Struct for PreProcessed Columns ID --- .../prover/src/constraint_framework/component.rs | 12 ++++++------ .../src/constraint_framework/expr/evaluator.rs | 5 +++-- crates/prover/src/constraint_framework/info.rs | 12 ++++++++---- crates/prover/src/constraint_framework/mod.rs | 3 ++- .../constraint_framework/preprocessed_columns.rs | 13 +++++++++++-- crates/prover/src/examples/blake/air.rs | 6 +++--- .../src/examples/blake/preprocessed_columns.rs | 14 ++++++++------ crates/prover/src/examples/plonk/mod.rs | 8 +++++--- .../xor/gkr_lookups/preprocessed_columns.rs | 13 ++++++++----- 9 files changed, 54 insertions(+), 32 deletions(-) diff --git a/crates/prover/src/constraint_framework/component.rs b/crates/prover/src/constraint_framework/component.rs index 5b1885d7c..01044da03 100644 --- a/crates/prover/src/constraint_framework/component.rs +++ b/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, /// Mapping of preprocessed columns to their index. - // TODO(Gali): Change Vec type to struct PreProcessedColumnId {pub id: String}. - preprocessed_columns: Vec, + preprocessed_columns: Vec, /// 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,14 @@ impl TraceLocationAllocator { } } - pub const fn preprocessed_columns(&self) -> &Vec { + pub const fn preprocessed_columns(&self) -> &Vec { &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]) { + pub fn validate_preprocessed_columns(&self, preprocessed_columns: &[PreProcessedColumnId]) { assert_eq!(self.preprocessed_columns, preprocessed_columns); } } @@ -142,7 +142,7 @@ impl FrameworkComponent { if let Some(pos) = location_allocator .preprocessed_columns .iter() - .position(|x| x == col) + .position(|x| x.id == col.id) { pos } else { diff --git a/crates/prover/src/constraint_framework/expr/evaluator.rs b/crates/prover/src/constraint_framework/expr/evaluator.rs index 10dce37d4..404ef600e 100644 --- a/crates/prover/src/constraint_framework/expr/evaluator.rs +++ b/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!(); diff --git a/crates/prover/src/constraint_framework/info.rs b/crates/prover/src/constraint_framework/info.rs index 02e3de8d2..6bb8462f3 100644 --- a/crates/prover/src/constraint_framework/info.rs +++ b/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>>, pub n_constraints: usize, - // TODO(Gali): Change Vec type to struct PreProcessedColumnId {pub id: String}. - pub preprocessed_columns: Vec, + pub preprocessed_columns: Vec, pub logup: LogupAtRow, pub arithmetic_counts: ArithmeticCounts, } impl InfoEvaluator { - pub fn new(log_size: u32, preprocessed_columns: Vec, logup_sums: LogupSums) -> Self { + pub fn new( + log_size: u32, + preprocessed_columns: Vec, + 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() } diff --git a/crates/prover/src/constraint_framework/mod.rs b/crates/prover/src/constraint_framework/mod.rs index 88a2908c6..ca037ec18 100644 --- a/crates/prover/src/constraint_framework/mod.rs +++ b/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 } diff --git a/crates/prover/src/constraint_framework/preprocessed_columns.rs b/crates/prover/src/constraint_framework/preprocessed_columns.rs index 01906ad7b..dc9ceeb88 100644 --- a/crates/prover/src/constraint_framework/preprocessed_columns.rs +++ b/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(), + } } } diff --git a/crates/prover/src/examples/blake/air.rs b/crates/prover/src/examples/blake/air.rs index 312f5c4ba..7bc044ca3 100644 --- a/crates/prover/src/examples/blake/air.rs +++ b/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 = ROUND_LOG_SPLIT + let blake_round_is_first_columns_iter: Vec = ROUND_LOG_SPLIT .iter() .map(|l| IsFirst::new(log_size + l).id()) .collect_vec(); diff --git a/crates/prover/src/examples/blake/preprocessed_columns.rs b/crates/prover/src/examples/blake/preprocessed_columns.rs index d30813b01..4c554adf0 100644 --- a/crates/prover/src/examples/blake/preprocessed_columns.rs +++ b/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 { diff --git a/crates/prover/src/examples/plonk/mod.rs b/crates/prover/src/examples/plonk/mod.rs index 8728b1e94..b2bcdbcb3 100644 --- a/crates/prover/src/examples/plonk/mod.rs +++ b/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), + } } } diff --git a/crates/prover/src/examples/xor/gkr_lookups/preprocessed_columns.rs b/crates/prover/src/examples/xor/gkr_lookups/preprocessed_columns.rs index b94781254..8329fd078 100644 --- a/crates/prover/src/examples/xor/gkr_lookups/preprocessed_columns.rs +++ b/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 + ), + } } }