Create ComponentTraceWriter trait.

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

@@ -1,6 +1,6 @@
 use std::iter::zip;
 
-use itertools::Itertools;
+use itertools::{zip_eq, Itertools};
 
 use super::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use super::{Air, AirProver, ComponentTrace};
@@ -12,7 +12,7 @@ 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::ComponentVec;
+use crate::core::{ColumnVec, ComponentVec};
 
 pub trait AirExt: Air {
     fn composition_log_degree_bound(&self) -> u32 {
@@ -100,6 +100,23 @@ 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: &[Vec<BaseField>],
+    ) -> ComponentVec<CircleEvaluation<B, BaseField, BitReversedOrder>> {
+        let trace_iter = &mut trace.iter();
+        ComponentVec(
+            zip_eq(self.prover_components().iter(), elements.iter())
+                .map(|(component, component_elements)| {
+                    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, component_elements)
+                })
+                .collect(),
+        )
+    }
+
     fn compute_composition_polynomial(
         &self,
         random_coeff: SecureField,
@@ -121,4 +138,5 @@ pub trait AirProverExt<B: Backend>: AirProver<B> {
         accumulator.finalize()
     }
 }
+
 impl<B: Backend, A: AirProver<B>> AirProverExt<B> for A {}

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

@@ -53,7 +53,15 @@ pub trait Component {
     );
 }
 
-pub trait ComponentProver<B: Backend>: Component {
+pub trait ComponentTraceWriter<B: Backend> {
+    fn write_interaction_trace(
+        &self,
+        trace: &ColumnVec<&CircleEvaluation<B, BaseField, BitReversedOrder>>,
+        elements: &[BaseField],
+    ) -> ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>;
+}
+
+pub trait ComponentProver<B: Backend>: Component + ComponentTraceWriter<B> {
     /// Evaluates the constraint quotients of the component on the evaluation domain.
     /// Accumulates quotients in `evaluation_accumulator`.
     fn evaluate_constraint_quotients_on_domain(

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

@@ -269,15 +269,21 @@ mod tests {
     use num_traits::Zero;
 
     use crate::core::air::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
-    use crate::core::air::{Air, AirProver, Component, ComponentProver, ComponentTrace};
+    use crate::core::air::{
+        Air, AirProver, Component, ComponentProver, ComponentTrace, ComponentTraceWriter,
+    };
     use crate::core::backend::cpu::CPUCircleEvaluation;
     use crate::core::backend::CPUBackend;
     use crate::core::circle::{CirclePoint, CirclePointIndex, Coset};
     use crate::core::fields::m31::BaseField;
     use crate::core::fields::qm31::SecureField;
-    use crate::core::poly::circle::{CanonicCoset, CircleDomain, MAX_CIRCLE_DOMAIN_LOG_SIZE};
+    use crate::core::poly::circle::{
+        CanonicCoset, CircleDomain, CircleEvaluation, MAX_CIRCLE_DOMAIN_LOG_SIZE,
+    };
+    use crate::core::poly::BitReversedOrder;
     use crate::core::prover::{prove, ProvingError};
     use crate::core::test_utils::test_channel;
+    use crate::core::ColumnVec;
     use crate::qm31;
 
     struct TestAir<C: ComponentProver<CPUBackend>> {
@@ -335,6 +341,16 @@ mod tests {
         }
     }
 
+    impl ComponentTraceWriter<CPUBackend> for TestComponent {
+        fn write_interaction_trace(
+            &self,
+            _trace: &ColumnVec<&CircleEvaluation<CPUBackend, BaseField, BitReversedOrder>>,
+            _elements: &[BaseField],
+        ) -> ColumnVec<CircleEvaluation<CPUBackend, BaseField, BitReversedOrder>> {
+            vec![]
+        }
+    }
+
     impl ComponentProver<CPUBackend> for TestComponent {
         fn evaluate_constraint_quotients_on_domain(
             &self,

crates/prover/src/examples/fibonacci/component.rs

@@ -4,14 +4,15 @@ use num_traits::One;
 
 use crate::core::air::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use crate::core::air::mask::shifted_mask_points;
-use crate::core::air::{Component, ComponentProver, ComponentTrace};
+use crate::core::air::{Component, ComponentProver, ComponentTrace, ComponentTraceWriter};
 use crate::core::backend::CPUBackend;
 use crate::core::circle::{CirclePoint, Coset};
 use crate::core::constraints::{coset_vanishing, pair_vanishing};
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::{ExtensionOf, FieldExpOps};
-use crate::core::poly::circle::CanonicCoset;
+use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
+use crate::core::poly::BitReversedOrder;
 use crate::core::utils::bit_reverse_index;
 use crate::core::ColumnVec;
 
@@ -116,6 +117,16 @@ impl Component for FibonacciComponent {
     }
 }
 
+impl ComponentTraceWriter<CPUBackend> for FibonacciComponent {
+    fn write_interaction_trace(
+        &self,
+        _trace: &ColumnVec<&CircleEvaluation<CPUBackend, BaseField, BitReversedOrder>>,
+        _elements: &[BaseField],
+    ) -> ColumnVec<CircleEvaluation<CPUBackend, BaseField, BitReversedOrder>> {
+        vec![]
+    }
+}
+
 impl ComponentProver<CPUBackend> for FibonacciComponent {
     fn evaluate_constraint_quotients_on_domain(
         &self,

crates/prover/src/examples/wide_fibonacci/avx.rs

@@ -4,7 +4,9 @@ use tracing::{span, Level};
 
 use super::component::{WideFibAir, WideFibComponent};
 use crate::core::air::accumulation::DomainEvaluationAccumulator;
-use crate::core::air::{AirProver, Component, ComponentProver, ComponentTrace};
+use crate::core::air::{
+    AirProver, Component, ComponentProver, ComponentTrace, ComponentTraceWriter,
+};
 use crate::core::backend::avx512::qm31::PackedSecureField;
 use crate::core::backend::avx512::{AVX512Backend, BaseFieldVec, PackedBaseField, VECS_LOG_SIZE};
 use crate::core::backend::{Col, Column, ColumnOps};
@@ -48,6 +50,16 @@ pub fn gen_trace(
         .collect_vec()
 }
 
+impl ComponentTraceWriter<AVX512Backend> for WideFibComponent {
+    fn write_interaction_trace(
+        &self,
+        _trace: &ColumnVec<&CircleEvaluation<AVX512Backend, BaseField, BitReversedOrder>>,
+        _elements: &[BaseField],
+    ) -> ColumnVec<CircleEvaluation<AVX512Backend, BaseField, BitReversedOrder>> {
+        vec![]
+    }
+}
+
 impl ComponentProver<AVX512Backend> for WideFibComponent {
     fn evaluate_constraint_quotients_on_domain(
         &self,

crates/prover/src/examples/wide_fibonacci/component.rs

@@ -1,13 +1,18 @@
+use itertools::Itertools;
+
 use crate::core::air::accumulation::PointEvaluationAccumulator;
 use crate::core::air::mask::fixed_mask_points;
-use crate::core::air::{Air, Component};
+use crate::core::air::{Air, Component, ComponentTraceWriter};
+use crate::core::backend::CPUBackend;
 use crate::core::circle::CirclePoint;
 use crate::core::constraints::coset_vanishing;
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::FieldExpOps;
-use crate::core::poly::circle::CanonicCoset;
+use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
+use crate::core::poly::BitReversedOrder;
 use crate::core::ColumnVec;
+use crate::examples::wide_fibonacci::trace_gen::write_lookup_column;
 
 pub const LOG_N_COLUMNS: usize = 8;
 pub const N_COLUMNS: usize = 1 << LOG_N_COLUMNS;
@@ -90,6 +95,21 @@ impl Component for WideFibComponent {
     }
 }
 
+impl ComponentTraceWriter<CPUBackend> for WideFibComponent {
+    fn write_interaction_trace(
+        &self,
+        trace: &ColumnVec<&CircleEvaluation<CPUBackend, BaseField, BitReversedOrder>>,
+        elements: &[BaseField],
+    ) -> ColumnVec<CircleEvaluation<CPUBackend, BaseField, BitReversedOrder>> {
+        let domain = trace[0].domain;
+        let input_trace = trace.iter().map(|eval| &eval.values).collect_vec();
+        let (alpha, z) = (elements[0], elements[1]);
+        let values = write_lookup_column(&input_trace, alpha, z);
+        let eval = CircleEvaluation::new(domain, values);
+        vec![eval]
+    }
+}
+
 // Input for the fibonacci claim.
 #[derive(Debug, Clone, Copy)]
 pub struct Input {

crates/prover/src/examples/wide_fibonacci/mod.rs

@@ -111,7 +111,8 @@ mod tests {
         let alpha = m31!(7);
         let z = m31!(11);
         let trace = gen_trace(&wide_fib, vec![input]);
-        let lookup_column = write_lookup_column(&trace, alpha, z);
+        let input_trace = trace.iter().collect_vec();
+        let lookup_column = write_lookup_column(&input_trace, alpha, z);
 
         assert_constraints_on_lookup_column(&lookup_column, &trace, alpha, z)
     }

crates/prover/src/examples/wide_fibonacci/trace_gen.rs

@@ -23,7 +23,7 @@ pub fn write_trace_row(
 }
 
 pub fn write_lookup_column(
-    input_trace: &[Vec<BaseField>],
+    input_trace: &[&Vec<BaseField>],
     // TODO(AlonH): Change alpha and z to SecureField.
     alpha: BaseField,
     z: BaseField,