Reuse commitment evaluation (#495)

src/core/air/air_ext.rs

@@ -1,13 +1,16 @@
 use std::iter::zip;
 
 use itertools::Itertools;
+use tracing::{span, Level};
 
 use super::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use super::{Air, ComponentTrace};
 use crate::core::backend::Backend;
 use crate::core::circle::CirclePoint;
+use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
-use crate::core::poly::circle::{CanonicCoset, CirclePoly, SecureCirclePoly};
+use crate::core::poly::circle::{CanonicCoset, CircleEvaluation, CirclePoly, SecureCirclePoly};
+use crate::core::poly::BitReversedOrder;
 use crate::core::prover::LOG_BLOWUP_FACTOR;
 use crate::core::ComponentVec;
 
@@ -52,6 +55,7 @@ pub trait AirExt<B: Backend>: Air<B> {
         ComponentVec<Vec<CirclePoint<SecureField>>>,
         ComponentVec<Vec<SecureField>>,
     ) {
+        let _span = span!(Level::INFO, "Evaluate columns out of domain").entered();
         let mut component_points = ComponentVec(Vec::new());
         let mut component_values = ComponentVec(Vec::new());
         zip(self.components(), component_traces).for_each(|(component, trace)| {
@@ -101,14 +105,17 @@ pub trait AirExt<B: Backend>: Air<B> {
     fn component_traces<'a>(
         &'a self,
         polynomials: &'a [CirclePoly<B>],
+        evals: &'a [CircleEvaluation<B, BaseField, BitReversedOrder>],
     ) -> Vec<ComponentTrace<'_, B>> {
         let poly_iter = &mut polynomials.iter();
+        let eval_iter = &mut evals.iter();
         self.components()
             .iter()
             .map(|component| {
                 let n_columns = component.trace_log_degree_bounds().len();
-                let columns = poly_iter.take(n_columns).collect();
-                ComponentTrace::new(columns)
+                let polys = poly_iter.take(n_columns).collect();
+                let evals = eval_iter.take(n_columns).collect();
+                ComponentTrace::new(polys, evals)
             })
             .collect()
     }

src/core/air/mod.rs

@@ -3,8 +3,10 @@ use std::iter::zip;
 use self::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use super::backend::Backend;
 use super::circle::CirclePoint;
+use super::fields::m31::BaseField;
 use super::fields::qm31::SecureField;
-use super::poly::circle::CirclePoly;
+use super::poly::circle::{CircleEvaluation, CirclePoly};
+use super::poly::BitReversedOrder;
 use super::ColumnVec;
 
 pub mod accumulation;
@@ -60,7 +62,7 @@ pub trait Component<B: Backend> {
         ColumnVec<Vec<SecureField>>,
     ) {
         let points = self.mask_points(point);
-        let values = zip(&points, &trace.columns)
+        let values = zip(&points, &trace.polys)
             .map(|(col_points, col)| {
                 col_points
                     .iter()
@@ -83,12 +85,21 @@ pub trait Component<B: Backend> {
     // TODO(spapini): Extra functions for FRI and decommitment.
 }
 
+/// A component trace is a set of polynomials for each column on that component.
+/// Each polynomial is stored both in a coefficients, and evaluations form (for efficiency)
 pub struct ComponentTrace<'a, B: Backend> {
-    pub columns: Vec<&'a CirclePoly<B>>,
+    /// Polynomials for each column.
+    pub polys: Vec<&'a CirclePoly<B>>,
+    /// Evaluations for each column. The evaluation domain is the commitment domain for that column
+    /// obtained from [AirExt::trace_commitment_domains()].
+    pub evals: Vec<&'a CircleEvaluation<B, BaseField, BitReversedOrder>>,
 }
 
 impl<'a, B: Backend> ComponentTrace<'a, B> {
-    pub fn new(columns: Vec<&'a CirclePoly<B>>) -> Self {
-        Self { columns }
+    pub fn new(
+        polys: Vec<&'a CirclePoly<B>>,
+        evals: Vec<&'a CircleEvaluation<B, BaseField, BitReversedOrder>>,
+    ) -> Self {
+        Self { polys, evals }
     }
 }

src/core/prover/mod.rs

@@ -98,7 +98,10 @@ pub fn prove<B: Backend + MerkleOps<MerkleHasher>>(
     let span = span!(Level::INFO, "Composition generation").entered();
     let composition_polynomial_poly = air.compute_composition_polynomial(
         random_coeff,
-        &air.component_traces(&commitment_scheme.trees[0].polynomials),
+        &air.component_traces(
+            &commitment_scheme.trees[0].polynomials,
+            &commitment_scheme.trees[0].evaluations,
+        ),
     );
     span.exit();
 

src/examples/fibonacci/component.rs

@@ -88,7 +88,7 @@ impl Component<CPUBackend> for FibonacciComponent {
         trace: &ComponentTrace<'_, CPUBackend>,
         evaluation_accumulator: &mut DomainEvaluationAccumulator<CPUBackend>,
     ) {
-        let poly = &trace.columns[0];
+        let poly = &trace.polys[0];
         let trace_domain = CanonicCoset::new(self.log_size);
         let trace_eval_domain = CanonicCoset::new(self.log_size + 1).circle_domain();
         let trace_eval = poly.evaluate(trace_eval_domain).bit_reverse();

src/examples/fibonacci/mod.rs

@@ -143,7 +143,9 @@ mod tests {
         let fib = Fibonacci::new(5, m31!(443693538));
         let trace = fib.get_trace();
         let trace_poly = trace.interpolate();
-        let trace = ComponentTrace::new(vec![&trace_poly]);
+        let trace_eval =
+            trace_poly.evaluate(CanonicCoset::new(trace_poly.log_size() + 1).circle_domain());
+        let trace = ComponentTrace::new(vec![&trace_poly], vec![&trace_eval]);
 
         // TODO(ShaharS), Change to a channel implementation to retrieve the random
         // coefficients from extension field.

src/examples/wide_fibonacci/avx.rs

@@ -70,16 +70,10 @@ impl Component<AVX512Backend> for WideFibComponent {
         trace: &ComponentTrace<'_, AVX512Backend>,
         evaluation_accumulator: &mut DomainEvaluationAccumulator<AVX512Backend>,
     ) {
-        let span = span!(Level::INFO, "Constraint eval extension").entered();
-        assert_eq!(trace.columns.len(), N_COLS);
+        assert_eq!(trace.polys.len(), N_COLS);
         // TODO(spapini): Steal evaluation from commitment.
         let eval_domain = CanonicCoset::new(self.log_size + 1).circle_domain();
-        let trace_eval = trace
-            .columns
-            .iter()
-            .map(|poly| poly.evaluate(eval_domain))
-            .collect_vec();
-        span.exit();
+        let trace_eval = &trace.evals;
 
         // Denoms.
         let span = span!(Level::INFO, "Constraint eval denominators").entered();

src/examples/wide_fibonacci/constraint_eval.rs

@@ -43,7 +43,7 @@ impl Component<CPUBackend> for WideFibComponent {
         let trace_eval_domain = CanonicCoset::new(max_constraint_degree).circle_domain();
         let mut trace_evals = vec![];
         for poly_index in 0..256 {
-            let poly = &trace.columns[poly_index];
+            let poly = &trace.polys[poly_index];
             trace_evals.push(poly.evaluate(trace_eval_domain));
         }
         let zero_domain = CanonicCoset::new(self.log_size).coset;

src/examples/wide_fibonacci/mod.rs

@@ -81,9 +81,15 @@ mod tests {
             .into_iter()
             .map(|eval| eval.interpolate())
             .collect_vec();
+        let eval_domain = CanonicCoset::new(wide_fib.log_size + 1).circle_domain();
+        let trace_evals = trace_polys
+            .iter()
+            .map(|poly| poly.evaluate(eval_domain))
+            .collect_vec();
 
         let trace = ComponentTrace {
-            columns: trace_polys.iter().collect(),
+            polys: trace_polys.iter().collect(),
+            evals: trace_evals.iter().collect(),
         };
 
         wide_fib.evaluate_constraint_quotients_on_domain(&trace, &mut acc);