Add n_preprocessed_columns to Components. (#874)

crates/prover/src/core/air/components.rs

@@ -9,11 +9,14 @@ use crate::core::pcs::TreeVec;
 use crate::core::poly::circle::SecureCirclePoly;
 use crate::core::ColumnVec;
 
-pub struct Components<'a>(pub Vec<&'a dyn Component>);
+pub struct Components<'a> {
+    pub components: Vec<&'a dyn Component>,
+    pub n_preprocessed_columns: usize,
+}
 
 impl<'a> Components<'a> {
     pub fn composition_log_degree_bound(&self) -> u32 {
-        self.0
+        self.components
             .iter()
             .map(|component| component.max_constraint_log_degree_bound())
             .max()
@@ -24,7 +27,11 @@ impl<'a> Components<'a> {
         &self,
         point: CirclePoint<SecureField>,
     ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
-        TreeVec::concat_cols(self.0.iter().map(|component| component.mask_points(point)))
+        TreeVec::concat_cols(
+            self.components
+                .iter()
+                .map(|component| component.mask_points(point)),
+        )
     }
 
     pub fn eval_composition_polynomial_at_point(
@@ -34,7 +41,7 @@ impl<'a> Components<'a> {
         random_coeff: SecureField,
     ) -> SecureField {
         let mut evaluation_accumulator = PointEvaluationAccumulator::new(random_coeff);
-        for component in &self.0 {
+        for component in &self.components {
             component.evaluate_constraint_quotients_at_point(
                 point,
                 mask_values,
@@ -46,31 +53,41 @@ impl<'a> Components<'a> {
 
     pub fn column_log_sizes(&self) -> TreeVec<ColumnVec<u32>> {
         TreeVec::concat_cols(
-            self.0
+            self.components
                 .iter()
                 .map(|component| component.trace_log_degree_bounds()),
         )
     }
 }
 
-pub struct ComponentProvers<'a, B: Backend>(pub Vec<&'a dyn ComponentProver<B>>);
+pub struct ComponentProvers<'a, B: Backend> {
+    pub components: Vec<&'a dyn ComponentProver<B>>,
+    pub n_preprocessed_columns: usize,
+}
 
 impl<'a, B: Backend> ComponentProvers<'a, B> {
     pub fn components(&self) -> Components<'_> {
-        Components(self.0.iter().map(|c| *c as &dyn Component).collect_vec())
+        Components {
+            components: self
+                .components
+                .iter()
+                .map(|c| *c as &dyn Component)
+                .collect_vec(),
+            n_preprocessed_columns: self.n_preprocessed_columns,
+        }
     }
     pub fn compute_composition_polynomial(
         &self,
         random_coeff: SecureField,
         trace: &Trace<'_, B>,
     ) -> SecureCirclePoly<B> {
-        let total_constraints: usize = self.0.iter().map(|c| c.n_constraints()).sum();
+        let total_constraints: usize = self.components.iter().map(|c| c.n_constraints()).sum();
         let mut accumulator = DomainEvaluationAccumulator::new(
             random_coeff,
             self.components().composition_log_degree_bound(),
             total_constraints,
         );
-        for component in &self.0 {
+        for component in &self.components {
             component.evaluate_constraint_quotients_on_domain(trace, &mut accumulator)
         }
         accumulator.finalize()

crates/prover/src/core/prover/mod.rs

@@ -11,6 +11,7 @@ use super::fields::secure_column::SECURE_EXTENSION_DEGREE;
 use super::fri::FriVerificationError;
 use super::pcs::{CommitmentSchemeProof, TreeVec};
 use super::vcs::ops::MerkleHasher;
+use crate::constraint_framework::PREPROCESSED_TRACE_IDX;
 use crate::core::channel::Channel;
 use crate::core::circle::CirclePoint;
 use crate::core::fields::m31::BaseField;
@@ -33,7 +34,13 @@ pub fn prove<B: BackendForChannel<MC>, MC: MerkleChannel>(
     channel: &mut MC::C,
     commitment_scheme: &mut CommitmentSchemeProver<'_, B, MC>,
 ) -> Result<StarkProof<MC::H>, ProvingError> {
-    let component_provers = ComponentProvers(components.to_vec());
+    let n_preprocessed_columns = commitment_scheme.trees[PREPROCESSED_TRACE_IDX]
+        .polynomials
+        .len();
+    let component_provers = ComponentProvers {
+        components: components.to_vec(),
+        n_preprocessed_columns,
+    };
     let trace = commitment_scheme.trace();
 
     // Evaluate and commit on composition polynomial.
@@ -88,7 +95,14 @@ pub fn verify<MC: MerkleChannel>(
     commitment_scheme: &mut CommitmentSchemeVerifier<MC>,
     proof: StarkProof<MC::H>,
 ) -> Result<(), VerificationError> {
-    let components = Components(components.to_vec());
+    let n_preprocessed_columns = commitment_scheme.trees[PREPROCESSED_TRACE_IDX]
+        .column_log_sizes
+        .len();
+
+    let components = Components {
+        components: components.to_vec(),
+        n_preprocessed_columns,
+    };
     let random_coeff = channel.draw_felt();
 
     // Read composition polynomial commitment.

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

@@ -807,14 +807,18 @@ mod tests {
         let proof = prove(components, channel, commitment_scheme).unwrap();
 
         // Verify.
-        let components = Components(components.iter().map(|&c| c as &dyn Component).collect());
+        let components = Components {
+            components: components.iter().map(|&c| c as &dyn Component).collect(),
+            n_preprocessed_columns: 0,
+        };
+
         let log_sizes = components.column_log_sizes();
         let channel = &mut Blake2sChannel::default();
         let commitment_scheme = &mut CommitmentSchemeVerifier::<Blake2sMerkleChannel>::new(config);
         commitment_scheme.commit(proof.commitments[0], &[], channel);
         commitment_scheme.commit(proof.commitments[1], &log_sizes[1], channel);
         commitment_scheme.commit(proof.commitments[2], &log_sizes[2], channel);
-        verify(&components.0, channel, commitment_scheme, proof)
+        verify(&components.components, channel, commitment_scheme, proof)
     }
 
     #[test]
@@ -891,14 +895,18 @@ mod tests {
             claim,
             MLE_EVAL_TRACE,
         );
-        let components = Components(vec![&mle_coeffs_col_component, &mle_eval_component]);
+        let components = Components {
+            components: vec![&mle_coeffs_col_component, &mle_eval_component],
+            n_preprocessed_columns: 0,
+        };
+
         let log_sizes = components.column_log_sizes();
         let channel = &mut Blake2sChannel::default();
         let commitment_scheme = &mut CommitmentSchemeVerifier::<Blake2sMerkleChannel>::new(config);
         commitment_scheme.commit(proof.commitments[0], &[], channel);
         commitment_scheme.commit(proof.commitments[1], &log_sizes[1], channel);
         commitment_scheme.commit(proof.commitments[2], &log_sizes[2], channel);
-        verify(&components.0, channel, commitment_scheme, proof)
+        verify(&components.components, channel, commitment_scheme, proof)
     }
 
     #[test]