moved constraint composition into constraint evaluator

prover/src/constraints/composition_poly.rs

@@ -3,37 +3,73 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 
-use super::ColMatrix;
-use math::{polynom::degree_of, FieldElement};
+use super::{ColMatrix, StarkDomain};
+use math::{fft, polynom::degree_of, FieldElement};
 use utils::collections::Vec;
 
-// COMPOSITION POLYNOMIAL
+// CONSTRAINT COMPOSITION POLYNOMIAL TRACE
+// ================================================================================================
+
+/// Represents merged evaluations of all constraint evaluations.
+pub struct CompositionPolyTrace<E>(Vec<E>);
+
+impl<E: FieldElement> CompositionPolyTrace<E> {
+    /// Returns a new instance of [CompositionPolyTrace] instantiated from the provided evaluations.
+    ///
+    /// # Panics
+    /// Panics if the number of evaluations is not a power of 2.
+    pub fn new(evaluations: Vec<E>) -> Self {
+        assert!(
+            evaluations.len().is_power_of_two(),
+            "length of composition polynomial trace must be a power of 2, but was {}",
+            evaluations.len(),
+        );
+
+        Self(evaluations)
+    }
+
+    /// Returns the number of evaluations in this trace.
+    pub fn num_rows(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Returns the internal vector representing this trace.
+    pub fn into_inner(self) -> Vec<E> {
+        self.0
+    }
+}
+
+// CONSTRAINT COMPOSITION POLYNOMIAL
 // ================================================================================================
 /// Represents a composition polynomial split into columns with each column being of length equal
-/// to trace_length. Thus, for example, if the composition polynomial has degree 2N - 1, where N
-/// is the trace length, it will be stored as two columns of size N (each of degree N - 1).
+/// to trace_length.
+///
+/// For example, if the composition polynomial has degree 2N - 1, where N is the trace length,
+/// it will be stored as two columns of size N (each of degree N - 1).
 pub struct CompositionPoly<E: FieldElement> {
     data: ColMatrix<E>,
 }
 
 impl<E: FieldElement> CompositionPoly<E> {
     /// Returns a new composition polynomial.
-    pub fn new(coefficients: Vec<E>, trace_length: usize, num_cols: usize) -> Self {
-        assert!(
-            coefficients.len().is_power_of_two(),
-            "size of composition polynomial must be a power of 2, but was {}",
-            coefficients.len(),
-        );
+    pub fn new(
+        composition_trace: CompositionPolyTrace<E>,
+        domain: &StarkDomain<E::BaseField>,
+        num_cols: usize,
+    ) -> Self {
         assert!(
-            trace_length.is_power_of_two(),
-            "trace length must be a power of 2, but was {trace_length}"
-        );
-        assert!(
-            trace_length < coefficients.len(),
-            "trace length must be smaller than size of composition polynomial"
+            domain.trace_length() < composition_trace.num_rows(),
+            "trace length must be smaller than length of composition polynomial trace"
         );
 
-        let polys = segment(coefficients, trace_length, num_cols);
+        let mut trace = composition_trace.into_inner();
+
+        // at this point, combined_poly contains evaluations of the combined constraint polynomial;
+        // we interpolate this polynomial to transform it into coefficient form.
+        let inv_twiddles = fft::get_inv_twiddles::<E::BaseField>(trace.len());
+        fft::interpolate_poly_with_offset(&mut trace, &inv_twiddles, domain.offset());
+
+        let polys = segment(trace, domain.trace_length(), num_cols);
 
         CompositionPoly {
             data: ColMatrix::new(polys),

prover/src/constraints/evaluation_table.rs

@@ -3,10 +3,13 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 
-use super::{CompositionPoly, ConstraintDivisor, ProverError, StarkDomain};
-use math::{batch_inversion, fft, FieldElement, StarkField};
+use super::{CompositionPolyTrace, ConstraintDivisor, StarkDomain};
+use math::{batch_inversion, FieldElement, StarkField};
 use utils::{batch_iter_mut, collections::Vec, iter_mut, uninit_vector};
 
+#[cfg(debug_assertions)]
+use math::fft;
+
 #[cfg(debug_assertions)]
 use air::TransitionConstraints;
 
@@ -96,6 +99,7 @@ impl<'a, E: FieldElement> ConstraintEvaluationTable<'a, E> {
     /// The first column always contains the value of combined transition constraint evaluations;
     /// the remaining columns contain values of assertion constraint evaluations combined based on
     /// common divisors.
+    #[allow(dead_code)]
     pub fn num_columns(&self) -> usize {
         self.evaluations.len()
     }
@@ -154,13 +158,9 @@ impl<'a, E: FieldElement> ConstraintEvaluationTable<'a, E> {
 
     // CONSTRAINT COMPOSITION
     // --------------------------------------------------------------------------------------------
-    /// Divides constraint evaluation columns by their respective divisor (in evaluation form),
-    /// combines the results into a single column, and interpolates this column into a composition
-    /// polynomial in coefficient form.
-    /// `num_cols` is the number of necessary columns (of length `trace_length`) needed to store
-    /// the coefficients of the constraint composition polynomial and is needed by
-    /// `CompositionPoly::new`.
-    pub fn into_poly(self, num_cols: usize) -> Result<CompositionPoly<E>, ProverError> {
+    /// Divides constraint evaluation columns by their respective divisor (in evaluation form) and
+    /// combines the results into a single column.
+    pub fn combine(self) -> CompositionPolyTrace<E> {
         // allocate memory for the combined polynomial
         let mut combined_poly = E::zeroed_vector(self.num_rows());
 
@@ -172,13 +172,7 @@ impl<'a, E: FieldElement> ConstraintEvaluationTable<'a, E> {
             acc_column(column, divisor, self.domain, &mut combined_poly);
         }
 
-        // at this point, combined_poly contains evaluations of the combined constraint polynomial;
-        // we interpolate this polynomial to transform it into coefficient form.
-        let inv_twiddles = fft::get_inv_twiddles::<E::BaseField>(combined_poly.len());
-        fft::interpolate_poly_with_offset(&mut combined_poly, &inv_twiddles, self.domain.offset());
-
-        let trace_length = self.domain.trace_length();
-        Ok(CompositionPoly::new(combined_poly, trace_length, num_cols))
+        CompositionPolyTrace::new(combined_poly)
     }
 
     // DEBUG HELPERS

prover/src/constraints/evaluator/default.rs

@@ -4,8 +4,8 @@
 // LICENSE file in the root directory of this source tree.
 
 use super::{
-    super::EvaluationTableFragment, BoundaryConstraints, ConstraintEvaluationTable,
-    ConstraintEvaluator, PeriodicValueTable, StarkDomain, TraceLde,
+    super::EvaluationTableFragment, BoundaryConstraints, CompositionPolyTrace,
+    ConstraintEvaluationTable, ConstraintEvaluator, PeriodicValueTable, StarkDomain, TraceLde,
 };
 use air::{
     Air, AuxTraceRandElements, ConstraintCompositionCoefficients, EvaluationFrame,
@@ -29,9 +29,8 @@ const MIN_CONCURRENT_DOMAIN_SIZE: usize = 8192;
 /// Default implementation of the [ConstraintEvaluator] trait.
 ///
 /// This implementation iterates over all evaluation frames of an extended execution trace and
-/// evaluates constraints over these frames one-by-one. Constraint evaluations for the constraints
-/// in the same domain are merged together using random linear combinations. Thus, the resulting
-/// [ConstraintEvaluationTable] will contain as many columns as there are unique constraint domains.
+/// evaluates constraints over these frames one-by-one. Constraint evaluations are merged together
+/// using random linear combinations and in the end, only a single column is returned.
 ///
 /// When `concurrent` feature is enabled, the extended execution trace is split into sets of
 /// sequential evaluation frames (called fragments), and frames in each fragment are evaluated
@@ -44,7 +43,7 @@ pub struct DefaultConstraintEvaluator<'a, A: Air, E: FieldElement<BaseField = A:
     periodic_values: PeriodicValueTable<E::BaseField>,
 }
 
-impl<'a, A, E> ConstraintEvaluator<'a, E> for DefaultConstraintEvaluator<'a, A, E>
+impl<'a, A, E> ConstraintEvaluator<E> for DefaultConstraintEvaluator<'a, A, E>
 where
     A: Air,
     E: FieldElement<BaseField = A::BaseField>,
@@ -54,8 +53,8 @@ where
     fn evaluate<T: TraceLde<E>>(
         self,
         trace: &T,
-        domain: &'a StarkDomain<<E as FieldElement>::BaseField>,
-    ) -> ConstraintEvaluationTable<'a, E> {
+        domain: &StarkDomain<<E as FieldElement>::BaseField>,
+    ) -> CompositionPolyTrace<E> {
         assert_eq!(
             trace.trace_len(),
             domain.lde_domain_size(),
@@ -108,7 +107,8 @@ where
         #[cfg(debug_assertions)]
         evaluation_table.validate_transition_degrees();
 
-        evaluation_table
+        // combine all evaluations into a single column and return
+        evaluation_table.combine()
     }
 }
 

prover/src/constraints/evaluator/mod.rs

@@ -3,7 +3,7 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 
-use super::{super::TraceLde, ConstraintEvaluationTable, StarkDomain};
+use super::{super::TraceLde, CompositionPolyTrace, ConstraintEvaluationTable, StarkDomain};
 use air::Air;
 use math::FieldElement;
 
@@ -24,17 +24,18 @@ use periodic_table::PeriodicValueTable;
 /// The logic for evaluating AIR constraints over a single evaluation frame is defined by the [Air]
 /// associated type, and the purpose of this trait is to execute this logic over all evaluation
 /// frames in an extended execution trace.
-pub trait ConstraintEvaluator<'a, E: FieldElement> {
+pub trait ConstraintEvaluator<E: FieldElement> {
     /// AIR constraints for the computation described by this evaluator.
     type Air: Air<BaseField = E::BaseField>;
 
-    /// Evaluates constraints against the provided extended execution trace.
+    /// Evaluates constraints against the provided extended execution trace, combines them into
+    /// evaluations of a single polynomial, and returns these evaluations.
     ///
     /// Constraints are evaluated over a constraint evaluation domain. This is an optimization
     /// because constraint evaluation domain can be many times smaller than the full LDE domain.
     fn evaluate<T: TraceLde<E>>(
         self,
         trace: &T,
-        domain: &'a StarkDomain<E::BaseField>,
-    ) -> ConstraintEvaluationTable<'a, E>;
+        domain: &StarkDomain<E::BaseField>,
+    ) -> CompositionPolyTrace<E>;
 }

prover/src/constraints/mod.rs

@@ -3,13 +3,13 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 
-use super::{ColMatrix, ConstraintDivisor, ProverError, RowMatrix, StarkDomain};
+use super::{ColMatrix, ConstraintDivisor, RowMatrix, StarkDomain};
 
 mod evaluator;
 pub use evaluator::{ConstraintEvaluator, DefaultConstraintEvaluator};
 
 mod composition_poly;
-pub use composition_poly::CompositionPoly;
+pub use composition_poly::{CompositionPoly, CompositionPolyTrace};
 
 mod evaluation_table;
 pub use evaluation_table::{ConstraintEvaluationTable, EvaluationTableFragment};

prover/src/lib.rs

@@ -80,7 +80,8 @@ pub use matrix::{ColMatrix, RowMatrix};
 
 mod constraints;
 pub use constraints::{
-    CompositionPoly, ConstraintCommitment, ConstraintEvaluator, DefaultConstraintEvaluator,
+    CompositionPoly, CompositionPolyTrace, ConstraintCommitment, ConstraintEvaluator,
+    DefaultConstraintEvaluator,
 };
 
 mod composer;
@@ -148,7 +149,7 @@ pub trait Prover {
         E: FieldElement<BaseField = Self::BaseField>;
 
     /// Constraints evaluator used to evaluate AIR constraints over the extended execution trace.
-    type ConstraintEvaluator<'a, E>: ConstraintEvaluator<'a, E, Air = Self::Air>
+    type ConstraintEvaluator<'a, E>: ConstraintEvaluator<E, Air = Self::Air>
     where
         E: FieldElement<BaseField = Self::BaseField>;
 
@@ -326,46 +327,29 @@ pub trait Prover {
         // 2 ----- evaluate constraints -----------------------------------------------------------
         // evaluate constraints specified by the AIR over the constraint evaluation domain, and
         // compute random linear combinations of these evaluations using coefficients drawn from
-        // the channel; this step evaluates only constraint numerators, thus, only constraints with
-        // identical denominators are merged together. the results are saved into a constraint
-        // evaluation table where each column contains merged evaluations of constraints with
-        // identical denominators.
+        // the channel
         #[cfg(feature = "std")]
         let now = Instant::now();
         let constraint_coeffs = channel.get_constraint_composition_coeffs();
         let evaluator = self.new_evaluator(&air, aux_trace_rand_elements, constraint_coeffs);
-        let constraint_evaluations = evaluator.evaluate(&trace_lde, &domain);
+        let composition_poly_trace = evaluator.evaluate(&trace_lde, &domain);
         #[cfg(feature = "std")]
         debug!(
             "Evaluated constraints over domain of 2^{} elements in {} ms",
-            constraint_evaluations.num_rows().ilog2(),
+            composition_poly_trace.num_rows().ilog2(),
             now.elapsed().as_millis()
         );
 
         // 3 ----- commit to constraint evaluations -----------------------------------------------
 
-        // first, build constraint composition polynomial from the constraint evaluation table:
-        // - divide all constraint evaluation columns by their respective divisors
-        // - combine them into a single column of evaluations,
-        // - interpolate the column into a polynomial in coefficient form
-        // - "break" the polynomial into a set of column polynomials each of degree equal to
-        //   trace_length - 1
-        #[cfg(feature = "std")]
-        let now = Instant::now();
-        let composition_poly =
-            constraint_evaluations.into_poly(air.context().num_constraint_composition_columns())?;
-        #[cfg(feature = "std")]
-        debug!(
-            "Converted constraint evaluations into {} composition polynomial columns of degree {} in {} ms",
-            composition_poly.num_columns(),
-            composition_poly.column_degree(),
-            now.elapsed().as_millis()
+        // first, build a commitment to the evaluations of the constraint composition polynomial
+        // columns
+        let (constraint_commitment, composition_poly) = self.build_constraint_commitment::<E>(
+            composition_poly_trace,
+            air.context().num_constraint_composition_columns(),
+            &domain,
         );
 
-        // then, build a commitment to the evaluations of the composition polynomial columns
-        let constraint_commitment =
-            self.build_constraint_commitment::<E>(&composition_poly, &domain);
-
         // then, commit to the evaluations of constraints by writing the root of the constraint
         // Merkle tree into the channel
         channel.commit_constraints(constraint_commitment.root());
@@ -484,23 +468,42 @@ pub trait Prover {
         Ok(proof)
     }
 
-    /// Evaluates constraint composition polynomial over the LDE domain and builds a commitment
-    /// to these evaluations.
+    /// Extends constraint composition polynomial over the LDE domain and builds a commitment to
+    /// its evaluations.
     ///
-    /// The evaluation is done by evaluating each composition polynomial column over the LDE
-    /// domain.
+    /// The extension is done by first interpolating the evaluations of the polynomial so that we
+    /// get the composition polynomial in coefficient form; then breaking the polynomial into
+    /// columns each of size equal to trace length, and finally evaluating each composition
+    /// polynomial column over the LDE domain.
     ///
     /// The commitment is computed by hashing each row in the evaluation matrix, and then building
     /// a Merkle tree from the resulting hashes.
     fn build_constraint_commitment<E>(
         &self,
-        composition_poly: &CompositionPoly<E>,
+        composition_poly_trace: CompositionPolyTrace<E>,
+        num_trace_poly_columns: usize,
         domain: &StarkDomain<Self::BaseField>,
-    ) -> ConstraintCommitment<E, Self::HashFn>
+    ) -> (ConstraintCommitment<E, Self::HashFn>, CompositionPoly<E>)
     where
         E: FieldElement<BaseField = Self::BaseField>,
     {
-        // evaluate composition polynomial columns over the LDE domain
+        // first, build constraint composition polynomial from its trace as follows:
+        // - interpolate the trace into a polynomial in coefficient form
+        // - "break" the polynomial into a set of column polynomials each of degree equal to
+        //   trace_length - 1
+        #[cfg(feature = "std")]
+        let now = Instant::now();
+        let composition_poly =
+            CompositionPoly::new(composition_poly_trace, domain, num_trace_poly_columns);
+        #[cfg(feature = "std")]
+        debug!(
+            "Converted constraint evaluations into {} composition polynomial columns of degree {} in {} ms",
+            composition_poly.num_columns(),
+            composition_poly.column_degree(),
+            now.elapsed().as_millis()
+        );
+
+        // then, evaluate composition polynomial columns over the LDE domain
         #[cfg(feature = "std")]
         let now = Instant::now();
         let composed_evaluations = RowMatrix::evaluate_polys_over::<DEFAULT_SEGMENT_WIDTH>(
@@ -515,7 +518,7 @@ pub trait Prover {
             now.elapsed().as_millis()
         );
 
-        // build constraint evaluation commitment
+        // finally, build constraint evaluation commitment
         #[cfg(feature = "std")]
         let now = Instant::now();
         let commitment = composed_evaluations.commit_to_rows();
@@ -526,6 +529,7 @@ pub trait Prover {
             constraint_commitment.tree_depth(),
             now.elapsed().as_millis()
         );
-        constraint_commitment
+
+        (constraint_commitment, composition_poly)
     }
 }