Add interaction to proof. (#636)

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

@@ -1,17 +1,16 @@
-use std::iter::zip;
-
-use itertools::Itertools;
+use itertools::{zip_eq, Itertools};
 
 use super::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use super::{Air, AirProver, 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, CircleEvaluation, CirclePoly, SecureCirclePoly};
-use crate::core::poly::BitReversedOrder;
+use crate::core::pcs::{CommitmentTreeProver, TreeVec};
+use crate::core::poly::circle::{CanonicCoset, SecureCirclePoly};
 use crate::core::prover::LOG_BLOWUP_FACTOR;
-use crate::core::ComponentVec;
+use crate::core::vcs::blake2_merkle::Blake2sMerkleHasher;
+use crate::core::vcs::ops::MerkleOps;
+use crate::core::{ComponentVec, InteractionElements};
 
 pub trait AirExt: Air {
     fn composition_log_degree_bound(&self) -> u32 {
@@ -46,13 +45,15 @@ pub trait AirExt: Air {
         point: CirclePoint<SecureField>,
         mask_values: &ComponentVec<Vec<SecureField>>,
         random_coeff: SecureField,
+        interaction_elements: &InteractionElements,
     ) -> SecureField {
         let mut evaluation_accumulator = PointEvaluationAccumulator::new(random_coeff);
-        zip(self.components(), &mask_values.0).for_each(|(component, mask)| {
+        zip_eq(self.components(), &mask_values.0).for_each(|(component, mask)| {
             component.evaluate_constraint_quotients_at_point(
                 point,
                 mask,
                 &mut evaluation_accumulator,
+                interaction_elements,
             )
         });
         evaluation_accumulator.finalize()
@@ -65,31 +66,50 @@ pub trait AirExt: Air {
             .collect()
     }
 
-    fn component_traces<'a, B: Backend>(
+    fn component_traces<'a, B: Backend + MerkleOps<Blake2sMerkleHasher>>(
         &'a self,
-        polynomials: &'a [CirclePoly<B>],
-        evals: &'a [CircleEvaluation<B, BaseField, BitReversedOrder>],
+        trees: &'a [CommitmentTreeProver<B>],
     ) -> 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 polys = poly_iter.take(n_columns).collect();
-                let evals = eval_iter.take(n_columns).collect();
-                ComponentTrace::new(polys, evals)
-            })
-            .collect()
+        let poly_iter = &mut trees[0].polynomials.iter();
+        let eval_iter = &mut trees[0].evaluations.iter();
+        let mut component_traces = vec![];
+        self.components().iter().for_each(|component| {
+            let n_columns = component.trace_log_degree_bounds().len();
+            let polys = poly_iter.take(n_columns).collect_vec();
+            let evals = eval_iter.take(n_columns).collect_vec();
+
+            component_traces.push(ComponentTrace {
+                polys: TreeVec::new(vec![polys]),
+                evals: TreeVec::new(vec![evals]),
+            });
+        });
+
+        if trees.len() > 1 {
+            let poly_iter = &mut trees[1].polynomials.iter();
+            let eval_iter = &mut trees[1].evaluations.iter();
+            self.components()
+                .iter()
+                .zip_eq(&mut component_traces)
+                .for_each(|(_component, component_trace)| {
+                    // TODO(AlonH): Implement n_interaction_columns() for component.
+                    let polys = poly_iter.take(1).collect_vec();
+                    let evals = eval_iter.take(1).collect_vec();
+                    component_trace.polys.push(polys);
+                    component_trace.evals.push(evals);
+                });
+        }
+        component_traces
     }
 }
+
 impl<A: Air + ?Sized> AirExt for A {}
 
 pub trait AirProverExt<B: Backend>: AirProver<B> {
     fn compute_composition_polynomial(
         &self,
         random_coeff: SecureField,
         component_traces: &[ComponentTrace<'_, B>],
+        interaction_elements: &InteractionElements,
     ) -> SecureCirclePoly<B> {
         let total_constraints: usize = self
             .prover_components()
@@ -101,8 +121,12 @@ pub trait AirProverExt<B: Backend>: AirProver<B> {
             self.composition_log_degree_bound(),
             total_constraints,
         );
-        zip(self.prover_components(), component_traces).for_each(|(component, trace)| {
-            component.evaluate_constraint_quotients_on_domain(trace, &mut accumulator)
+        zip_eq(self.prover_components(), component_traces).for_each(|(component, trace)| {
+            component.evaluate_constraint_quotients_on_domain(
+                trace,
+                &mut accumulator,
+                interaction_elements,
+            )
         });
         accumulator.finalize()
     }

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

@@ -4,6 +4,7 @@ use super::channel::Blake2sChannel;
 use super::circle::CirclePoint;
 use super::fields::m31::BaseField;
 use super::fields::qm31::SecureField;
+use super::pcs::TreeVec;
 use super::poly::circle::{CircleEvaluation, CirclePoly};
 use super::poly::BitReversedOrder;
 use super::{ColumnVec, ComponentVec, InteractionElements};
@@ -35,7 +36,7 @@ pub trait AirTraceWriter<B: Backend>: AirTraceVerifier {
         elements: &InteractionElements,
     ) -> ComponentVec<CircleEvaluation<B, BaseField, BitReversedOrder>>;
 
-    fn to_air_prover(&self) -> &dyn AirProver<B>;
+    fn to_air_prover(&self) -> &impl AirProver<B>;
 }
 
 pub trait AirProver<B: Backend>: Air {
@@ -66,6 +67,7 @@ pub trait Component {
         point: CirclePoint<SecureField>,
         mask: &ColumnVec<Vec<SecureField>>,
         evaluation_accumulator: &mut PointEvaluationAccumulator,
+        interaction_elements: &InteractionElements,
     );
 }
 
@@ -84,23 +86,24 @@ pub trait ComponentProver<B: Backend>: Component {
         &self,
         trace: &ComponentTrace<'_, B>,
         evaluation_accumulator: &mut DomainEvaluationAccumulator<B>,
+        interaction_elements: &InteractionElements,
     );
 }
 
 /// 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> {
     /// Polynomials for each column.
-    pub polys: Vec<&'a CirclePoly<B>>,
+    pub polys: TreeVec<ColumnVec<&'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>>,
+    pub evals: TreeVec<ColumnVec<&'a CircleEvaluation<B, BaseField, BitReversedOrder>>>,
 }
 
 impl<'a, B: Backend> ComponentTrace<'a, B> {
     pub fn new(
-        polys: Vec<&'a CirclePoly<B>>,
-        evals: Vec<&'a CircleEvaluation<B, BaseField, BitReversedOrder>>,
+        polys: TreeVec<ColumnVec<&'a CirclePoly<B>>>,
+        evals: TreeVec<ColumnVec<&'a CircleEvaluation<B, BaseField, BitReversedOrder>>>,
     ) -> Self {
         Self { polys, evals }
     }

crates/prover/src/core/mod.rs

@@ -61,8 +61,19 @@ impl<T> DerefMut for ComponentVec<T> {
     }
 }
 
+#[derive(Default)]
 pub struct InteractionElements(BTreeMap<String, BaseField>);
 
+impl InteractionElements {
+    pub fn new(elements: BTreeMap<String, BaseField>) -> Self {
+        Self(elements)
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+}
+
 impl Index<&str> for InteractionElements {
     type Output = BaseField;
 

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

@@ -11,6 +11,6 @@ pub mod quotients;
 mod utils;
 mod verifier;
 
-pub use self::prover::{CommitmentSchemeProof, CommitmentSchemeProver};
+pub use self::prover::{CommitmentSchemeProof, CommitmentSchemeProver, CommitmentTreeProver};
 pub use self::utils::TreeVec;
 pub use self::verifier::CommitmentSchemeVerifier;