Separate componenet and air trace writers.

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

@@ -1,19 +1,17 @@
-use std::collections::{BTreeMap, BTreeSet};
 use std::iter::zip;
 
-use itertools::{zip_eq, Itertools};
+use itertools::Itertools;
 
 use super::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use super::{Air, AirProver, ComponentTrace};
 use crate::core::backend::Backend;
-use crate::core::channel::{Blake2sChannel, Channel};
 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::prover::LOG_BLOWUP_FACTOR;
-use crate::core::{ColumnVec, ComponentVec, InteractionElements};
+use crate::core::ComponentVec;
 
 pub trait AirExt: Air {
     fn composition_log_degree_bound(&self) -> u32 {
@@ -43,15 +41,6 @@ pub trait AirExt: Air {
         component_points
     }
 
-    fn interaction_elements(&self, channel: &mut Blake2sChannel) -> InteractionElements {
-        let mut ids = BTreeSet::new();
-        for component in self.components() {
-            ids.extend(component.interaction_element_ids());
-        }
-        let elements = channel.draw_felts(ids.len()).into_iter().map(|e| e.0 .0);
-        InteractionElements(BTreeMap::from_iter(zip_eq(ids, elements)))
-    }
-
     fn eval_composition_polynomial_at_point(
         &self,
         point: CirclePoint<SecureField>,
@@ -97,24 +86,6 @@ pub trait AirExt: Air {
 impl<A: Air + ?Sized> AirExt for A {}
 
 pub trait AirProverExt<B: Backend>: AirProver<B> {
-    fn interact(
-        &self,
-        trace: &ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
-        elements: &InteractionElements,
-    ) -> ComponentVec<CircleEvaluation<B, BaseField, BitReversedOrder>> {
-        let trace_iter = &mut trace.iter();
-        ComponentVec(
-            self.prover_components()
-                .iter()
-                .map(|component| {
-                    let n_columns = component.trace_log_degree_bounds().len();
-                    let trace_columns = trace_iter.take(n_columns).collect_vec();
-                    component.write_interaction_trace(&trace_columns, elements)
-                })
-                .collect(),
-        )
-    }
-
     fn compute_composition_polynomial(
         &self,
         random_coeff: SecureField,

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

@@ -1,11 +1,12 @@
 use self::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use super::backend::Backend;
+use super::channel::Blake2sChannel;
 use super::circle::CirclePoint;
 use super::fields::m31::BaseField;
 use super::fields::qm31::SecureField;
 use super::poly::circle::{CircleEvaluation, CirclePoly};
 use super::poly::BitReversedOrder;
-use super::{ColumnVec, InteractionElements};
+use super::{ColumnVec, ComponentVec, InteractionElements};
 
 pub mod accumulation;
 mod air_ext;
@@ -22,6 +23,21 @@ pub use air_ext::{AirExt, AirProverExt};
 pub trait Air {
     fn components(&self) -> Vec<&dyn Component>;
 }
+
+pub trait AirTraceVerifier {
+    fn interaction_elements(&self, channel: &mut Blake2sChannel) -> InteractionElements;
+}
+
+pub trait AirTraceWriter<B: Backend>: AirTraceVerifier {
+    fn interact(
+        &self,
+        trace: &ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
+        elements: &InteractionElements,
+    ) -> ComponentVec<CircleEvaluation<B, BaseField, BitReversedOrder>>;
+
+    fn to_air_prover(&self) -> &dyn AirProver<B>;
+}
+
 pub trait AirProver<B: Backend>: Air {
     fn prover_components(&self) -> Vec<&dyn ComponentProver<B>>;
 }
@@ -61,8 +77,7 @@ pub trait ComponentTraceWriter<B: Backend> {
     ) -> ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>;
 }
 
-// TODO(AlonH): Rethink this trait.
-pub trait ComponentProver<B: Backend>: Component + ComponentTraceWriter<B> {
+pub trait ComponentProver<B: Backend>: Component {
     /// Evaluates the constraint quotients of the component on the evaluation domain.
     /// Accumulates quotients in `evaluation_accumulator`.
     fn evaluate_constraint_quotients_on_domain(