updated changelog, added comments, and minor refactoring

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 0.7.0 (TBD)
 * [BREAKING] replaced the `TraceLde` struct with a trait (#207).
 * [BREAKING] refactored `RandomCoin` trait (#214).
+* [BREAKING] replaced the `ConstraintEvaluator` struct with a trait (#217).
 
 ## 0.6.5 (2023-08-09) - math crate only
 * Added conditional support for serde on field elements (#209)

prover/src/constraints/evaluation_table.rs

@@ -91,10 +91,11 @@ impl<'a, E: FieldElement> ConstraintEvaluationTable<'a, E> {
         self.evaluations[0].len()
     }
 
-    /// Returns number of columns in this table. 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)]
+    /// Returns number of columns in this table.
+    ///
+    /// 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.
     pub fn num_columns(&self) -> usize {
         self.evaluations.len()
     }

prover/src/constraints/evaluator.rs → prover/src/constraints/evaluator/default.rs

@@ -4,14 +4,13 @@
 // LICENSE file in the root directory of this source tree.
 
 use super::{
-    super::TraceLde, evaluation_table::EvaluationTableFragment, BoundaryConstraints,
-    ConstraintEvaluationTable, PeriodicValueTable, StarkDomain,
+    super::EvaluationTableFragment, BoundaryConstraints, ConstraintEvaluationTable,
+    ConstraintEvaluator, PeriodicValueTable, StarkDomain, TraceLde,
 };
 use air::{
     Air, AuxTraceRandElements, ConstraintCompositionCoefficients, EvaluationFrame,
     TransitionConstraints,
 };
-
 use math::FieldElement;
 use utils::iter_mut;
 
@@ -24,22 +23,19 @@ use utils::{iterators::*, rayon};
 #[cfg(feature = "concurrent")]
 const MIN_CONCURRENT_DOMAIN_SIZE: usize = 8192;
 
-// CONSTRAINT EVALUATOR TRAIT
-// ================================================================================================
-
-pub trait ConstraintEvaluator<'a, E: FieldElement> {
-    type Air: Air<BaseField = E::BaseField>;
-
-    fn evaluate<T: TraceLde<BaseField = E::BaseField, ExtensionField = E>>(
-        self,
-        trace: &T,
-        domain: &'a StarkDomain<E::BaseField>,
-    ) -> ConstraintEvaluationTable<'a, E>;
-}
-
-// CONSTRAINT EVALUATOR DEFAULT IMPLEMENTATION
+// DEFAULT CONSTRAINT EVALUATOR
 // ================================================================================================
 
+/// 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.
+///
+/// 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
+/// in separate threads.
 pub struct DefaultConstraintEvaluator<'a, A: Air, E: FieldElement<BaseField = A::BaseField>> {
     air: &'a A,
     boundary_constraints: BoundaryConstraints<E>,
@@ -48,15 +44,14 @@ pub struct DefaultConstraintEvaluator<'a, A: Air, E: FieldElement<BaseField = A:
     periodic_values: PeriodicValueTable<E::BaseField>,
 }
 
-impl<'a, A: Air, E: FieldElement<BaseField = A::BaseField>> ConstraintEvaluator<'a, E>
-    for DefaultConstraintEvaluator<'a, A, E>
+impl<'a, A, E> ConstraintEvaluator<'a, E> for DefaultConstraintEvaluator<'a, A, E>
+where
+    A: Air,
+    E: FieldElement<BaseField = A::BaseField>,
 {
     type Air = A;
 
-    /// Evaluates constraints against the provided extended execution trace. 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<BaseField = E::BaseField, ExtensionField = E>>(
+    fn evaluate<T: TraceLde<E>>(
         self,
         trace: &T,
         domain: &'a StarkDomain<<E as FieldElement>::BaseField>,
@@ -116,7 +111,12 @@ impl<'a, A: Air, E: FieldElement<BaseField = A::BaseField>> ConstraintEvaluator<
         evaluation_table
     }
 }
-impl<'a, A: Air, E: FieldElement<BaseField = A::BaseField>> DefaultConstraintEvaluator<'a, A, E> {
+
+impl<'a, A, E> DefaultConstraintEvaluator<'a, A, E>
+where
+    A: Air,
+    E: FieldElement<BaseField = A::BaseField>,
+{
     // CONSTRUCTOR
     // --------------------------------------------------------------------------------------------
     /// Returns a new evaluator which can be used to evaluate transition and boundary constraints
@@ -154,7 +154,7 @@ impl<'a, A: Air, E: FieldElement<BaseField = A::BaseField>> DefaultConstraintEva
     /// Evaluates constraints for a single fragment of the evaluation table.
     ///
     /// This evaluates constraints only over the main segment of the execution trace.
-    fn evaluate_fragment_main<T: TraceLde<BaseField = E::BaseField, ExtensionField = E>>(
+    fn evaluate_fragment_main<T: TraceLde<E>>(
         &self,
         trace: &T,
         domain: &StarkDomain<A::BaseField>,
@@ -205,7 +205,7 @@ impl<'a, A: Air, E: FieldElement<BaseField = A::BaseField>> DefaultConstraintEva
     ///
     /// This evaluates constraints only over all segments of the execution trace (i.e. main segment
     /// and all auxiliary segments).
-    fn evaluate_fragment_full<T: TraceLde<BaseField = E::BaseField, ExtensionField = E>>(
+    fn evaluate_fragment_full<T: TraceLde<E>>(
         &self,
         trace: &T,
         domain: &StarkDomain<A::BaseField>,

prover/src/constraints/evaluator/mod.rs

@@ -0,0 +1,40 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+//
+// 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 air::Air;
+use math::FieldElement;
+
+mod default;
+pub use default::DefaultConstraintEvaluator;
+
+mod boundary;
+use boundary::BoundaryConstraints;
+
+mod periodic_table;
+use periodic_table::PeriodicValueTable;
+
+// CONSTRAINT EVALUATOR TRAIT
+// ================================================================================================
+
+/// Contains logic for evaluating AIR constraints over an extended execution trace.
+///
+/// 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> {
+    /// AIR constraints for the computation described by this evaluator.
+    type Air: Air<BaseField = E::BaseField>;
+
+    /// Evaluates constraints against the provided extended execution trace.
+    ///
+    /// 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>;
+}

prover/src/constraints/mod.rs

@@ -5,20 +5,14 @@
 
 use super::{ColMatrix, ConstraintDivisor, ProverError, RowMatrix, StarkDomain};
 
-mod boundary;
-use boundary::BoundaryConstraints;
-
-mod periodic_table;
-use periodic_table::PeriodicValueTable;
-
 mod evaluator;
 pub use evaluator::{ConstraintEvaluator, DefaultConstraintEvaluator};
 
 mod composition_poly;
 pub use composition_poly::CompositionPoly;
 
 mod evaluation_table;
-pub use evaluation_table::ConstraintEvaluationTable;
+pub use evaluation_table::{ConstraintEvaluationTable, EvaluationTableFragment};
 
 mod commitment;
 pub use commitment::ConstraintCommitment;

prover/src/lib.rs

@@ -109,8 +109,8 @@ const DEFAULT_SEGMENT_WIDTH: usize = 8;
 /// A STARK prover can be used to generate STARK proofs. The prover contains definitions of a
 /// computation's AIR (specified via [Air](Prover::Air) associated type), execution trace
 /// (specified via [Trace](Prover::Trace) associated type) and hash function to be used (specified
-/// via [HashFn](Prover::HashFn) associated type), and exposes [prove()](Prover::prove) method which can
-/// be used to build STARK proofs for provided execution traces.
+/// via [HashFn](Prover::HashFn) associated type), and exposes [prove()](Prover::prove) method which
+/// can be used to build STARK proofs for provided execution traces.
 ///
 /// Thus, once a prover is defined and instantiated, generating a STARK proof consists of two
 /// steps:
@@ -120,6 +120,12 @@ const DEFAULT_SEGMENT_WIDTH: usize = 8;
 ///
 /// The generated proof is built using protocol parameters defined by the [ProofOptions] struct
 /// return from [Prover::options] method.
+///
+/// To further customize the prover, implementers can specify custom implementations of the
+/// [RandomCoin], [TraceLde], and [ConstraintEvaluator] associated types (default implementations
+/// of these types are provided with the prover). For example, providing custom implementations
+/// of [TraceLde] and/or [ConstraintEvaluator] can be beneficial when some steps of proof
+/// generation can be delegated to non-CPU hardware (e.g., GPUs).
 pub trait Prover {
     /// Base field for the computation described by this prover.
     type BaseField: StarkField + ExtensibleField<2> + ExtensibleField<3>;
@@ -137,11 +143,9 @@ pub trait Prover {
     type RandomCoin: RandomCoin<BaseField = Self::BaseField, Hasher = Self::HashFn>;
 
     /// Trace low-degree extension for building the LDEs of trace segments and their commitments.
-    type TraceLde<E: FieldElement<BaseField = Self::BaseField>>: TraceLde<
-        BaseField = Self::BaseField,
-        ExtensionField = E,
-        HashFn = Self::HashFn,
-    >;
+    type TraceLde<E>: TraceLde<E, HashFn = Self::HashFn>
+    where
+        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>

prover/src/trace/trace_lde/default/mod.rs

@@ -65,11 +65,11 @@ impl<E: FieldElement, H: ElementHasher<BaseField = E::BaseField>> DefaultTraceLd
     }
 }
 
-impl<E: FieldElement, H: ElementHasher<BaseField = E::BaseField>> TraceLde
-    for DefaultTraceLde<E, H>
+impl<E, H> TraceLde<E> for DefaultTraceLde<E, H>
+where
+    E: FieldElement,
+    H: ElementHasher<BaseField = E::BaseField>,
 {
-    type BaseField = E::BaseField;
-    type ExtensionField = E;
     type HashFn = H;
 
     /// Takes the main trace segment columns as input, interpolates them into polynomials in
@@ -81,12 +81,12 @@ impl<E: FieldElement, H: ElementHasher<BaseField = E::BaseField>> TraceLde
     /// segment and the new [DefaultTraceLde].
     fn new(
         trace_info: &TraceInfo,
-        main_trace: &ColMatrix<Self::BaseField>,
-        domain: &StarkDomain<Self::BaseField>,
-    ) -> (TracePolyTable<Self::ExtensionField>, Self) {
+        main_trace: &ColMatrix<E::BaseField>,
+        domain: &StarkDomain<E::BaseField>,
+    ) -> (TracePolyTable<E>, Self) {
         // extend the main execution trace and build a Merkle tree from the extended trace
         let (main_segment_lde, main_segment_tree, main_segment_polys) =
-            build_trace_commitment::<E, Self::BaseField, H>(main_trace, domain);
+            build_trace_commitment::<E, E::BaseField, H>(main_trace, domain);
 
         let trace_poly_table = TracePolyTable::new(main_segment_polys);
         let trace_lde = DefaultTraceLde {
@@ -122,7 +122,7 @@ impl<E: FieldElement, H: ElementHasher<BaseField = E::BaseField>> TraceLde
     fn add_aux_segment(
         &mut self,
         aux_trace: &ColMatrix<E>,
-        domain: &StarkDomain<Self::BaseField>,
+        domain: &StarkDomain<E::BaseField>,
     ) -> (ColMatrix<E>, <Self::HashFn as Hasher>::Digest) {
         // extend the auxiliary trace segment and build a Merkle tree from the extended trace
         let (aux_segment_lde, aux_segment_tree, aux_segment_polys) =
@@ -151,7 +151,7 @@ impl<E: FieldElement, H: ElementHasher<BaseField = E::BaseField>> TraceLde
     fn read_main_trace_frame_into(
         &self,
         lde_step: usize,
-        frame: &mut EvaluationFrame<Self::BaseField>,
+        frame: &mut EvaluationFrame<E::BaseField>,
     ) {
         // at the end of the trace, next state wraps around and we read the first step again
         let next_lde_step = (lde_step + self.blowup()) % self.trace_len();

prover/src/trace/trace_lde/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::{ColMatrix, EvaluationFrame, FieldElement, StarkField, TracePolyTable};
+use super::{ColMatrix, EvaluationFrame, FieldElement, TracePolyTable};
 use crate::StarkDomain;
 use air::{proof::Queries, TraceInfo, TraceLayout};
 use crypto::{ElementHasher, Hasher};
@@ -21,13 +21,9 @@ pub use default::DefaultTraceLde;
 ///   will always be elements in the base field (even when an extension field is used).
 /// - Auxiliary segments: a list of 0 or more segments for traces generated after the prover
 ///   commits to the first trace segment. Currently, at most 1 auxiliary segment is possible.
-pub trait TraceLde: Sync {
-    /// The base field, used for computation on the main trace segment.
-    type BaseField: StarkField;
-    /// The extension field, used for computation on auxiliary trace segments.
-    type ExtensionField: FieldElement<BaseField = Self::BaseField>;
+pub trait TraceLde<E: FieldElement>: Sync {
     /// The hash function used for building the Merkle tree commitments to trace segment LDEs.
-    type HashFn: ElementHasher<BaseField = Self::BaseField>;
+    type HashFn: ElementHasher<BaseField = E::BaseField>;
 
     /// Takes the main trace segment columns as input, interpolates them into polynomials in
     /// coefficient form, and evaluates the polynomials over the LDE domain.
@@ -36,9 +32,9 @@ pub trait TraceLde: Sync {
     /// and a new [TraceLde] instance from which the LDE and trace commitments can be obtained.
     fn new(
         trace_info: &TraceInfo,
-        main_trace: &ColMatrix<Self::BaseField>,
-        domain: &StarkDomain<Self::BaseField>,
-    ) -> (TracePolyTable<Self::ExtensionField>, Self);
+        main_trace: &ColMatrix<E::BaseField>,
+        domain: &StarkDomain<E::BaseField>,
+    ) -> (TracePolyTable<E>, Self);
 
     /// Returns the commitment to the low-degree extension of the main trace segment.
     fn get_main_trace_commitment(&self) -> <Self::HashFn as Hasher>::Digest;
@@ -57,26 +53,19 @@ pub trait TraceLde: Sync {
     /// - this segment would exceed the number of segments specified by the trace layout.
     fn add_aux_segment(
         &mut self,
-        aux_trace: &ColMatrix<Self::ExtensionField>,
-        domain: &StarkDomain<Self::BaseField>,
-    ) -> (
-        ColMatrix<Self::ExtensionField>,
-        <Self::HashFn as Hasher>::Digest,
-    );
+        aux_trace: &ColMatrix<E>,
+        domain: &StarkDomain<E::BaseField>,
+    ) -> (ColMatrix<E>, <Self::HashFn as Hasher>::Digest);
 
     /// Reads current and next rows from the main trace segment into the specified frame.
     fn read_main_trace_frame_into(
         &self,
         lde_step: usize,
-        frame: &mut EvaluationFrame<Self::BaseField>,
+        frame: &mut EvaluationFrame<E::BaseField>,
     );
 
     /// Reads current and next rows from the auxiliary trace segment into the specified frame.
-    fn read_aux_trace_frame_into(
-        &self,
-        lde_step: usize,
-        frame: &mut EvaluationFrame<Self::ExtensionField>,
-    );
+    fn read_aux_trace_frame_into(&self, lde_step: usize, frame: &mut EvaluationFrame<E>);
 
     /// Returns trace table rows at the specified positions along with Merkle authentication paths
     /// from the commitment root to these rows.