From 8da2844f62835f46236e323bf1a5c32855fb66e4 Mon Sep 17 00:00:00 2001
From: Alon Haramati <alonh@starkware.co>
Date: Sun, 23 Jun 2024 16:57:15 +0300
Subject: [PATCH] Merge trace writer and trace generator.

---
 crates/prover/src/core/air/mod.rs             | 25 ++--------------
 crates/prover/src/core/prover/mod.rs          | 29 +++++++++++++------
 crates/prover/src/examples/fibonacci/air.rs   |  9 +++---
 .../src/examples/fibonacci/component.rs       | 17 +++++++++--
 .../src/examples/wide_fibonacci/component.rs  | 17 +++++++++--
 .../wide_fibonacci/constraint_eval.rs         |  8 ++---
 .../prover/src/examples/wide_fibonacci/mod.rs |  3 +-
 .../src/examples/wide_fibonacci/simd.rs       | 22 ++++++++++----
 crates/prover/src/trace_generation/mod.rs     | 24 ++++++++++++++-
 .../prover/src/trace_generation/registry.rs   | 10 ++++++-
 10 files changed, 109 insertions(+), 55 deletions(-)

diff --git a/crates/prover/src/core/air/mod.rs b/crates/prover/src/core/air/mod.rs
index 4b5f53c50..798d4df19 100644
--- a/crates/prover/src/core/air/mod.rs
+++ b/crates/prover/src/core/air/mod.rs
@@ -1,6 +1,7 @@
+use crate::trace_generation::AirTraceVerifier;
+
 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;
@@ -25,20 +26,6 @@ pub trait Air: AirTraceVerifier {
     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,
-    ) -> Vec<CircleEvaluation<B, BaseField, BitReversedOrder>>;
-
-    fn to_air_prover(&self) -> &impl AirProver<B>;
-}
-
 pub trait AirProver<B: Backend>: Air {
     fn prover_components(&self) -> Vec<&dyn ComponentProver<B>>;
 }
@@ -78,14 +65,6 @@ pub trait Component {
     );
 }
 
-pub trait ComponentTraceWriter<B: Backend> {
-    fn write_interaction_trace(
-        &self,
-        trace: &ColumnVec<&CircleEvaluation<B, BaseField, BitReversedOrder>>,
-        elements: &InteractionElements,
-    ) -> ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>;
-}
-
 pub trait ComponentProver<B: Backend>: Component {
     /// Evaluates the constraint quotients of the component on the evaluation domain.
     /// Accumulates quotients in `evaluation_accumulator`.
diff --git a/crates/prover/src/core/prover/mod.rs b/crates/prover/src/core/prover/mod.rs
index 4e18fb43c..8457e4baa 100644
--- a/crates/prover/src/core/prover/mod.rs
+++ b/crates/prover/src/core/prover/mod.rs
@@ -2,7 +2,7 @@ use itertools::Itertools;
 use thiserror::Error;
 use tracing::{span, Level};
 
-use super::air::{AirProver, AirTraceWriter};
+use super::air::AirProver;
 use super::backend::Backend;
 use super::fields::secure_column::SECURE_EXTENSION_DEGREE;
 use super::fri::FriVerificationError;
@@ -26,6 +26,7 @@ use crate::core::vcs::hasher::Hasher;
 use crate::core::vcs::ops::MerkleOps;
 use crate::core::vcs::verifier::MerkleVerificationError;
 use crate::core::ComponentVec;
+use crate::trace_generation::AirTraceGenerator;
 
 type Channel = Blake2sChannel;
 type ChannelHasher = Blake2sHasher;
@@ -55,7 +56,7 @@ pub struct AdditionalProofData {
 }
 
 pub fn evaluate_and_commit_on_trace<B: Backend + MerkleOps<MerkleHasher>>(
-    air: &impl AirTraceWriter<B>,
+    air: &impl AirTraceGenerator<B>,
     channel: &mut Channel,
     twiddles: &TwiddleTree<B>,
     trace: ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
@@ -149,7 +150,7 @@ pub fn generate_proof<B: Backend + MerkleOps<MerkleHasher>>(
 }
 
 pub fn prove<B: Backend + MerkleOps<MerkleHasher>>(
-    air: &impl AirTraceWriter<B>,
+    air: &impl AirTraceGenerator<B>,
     channel: &mut Channel,
     trace: ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
 ) -> Result<StarkProof, ProvingError> {
@@ -355,10 +356,7 @@ mod tests {
     use num_traits::Zero;
 
     use crate::core::air::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
-    use crate::core::air::{
-        Air, AirProver, AirTraceVerifier, AirTraceWriter, Component, ComponentProver,
-        ComponentTrace, ComponentTraceWriter,
-    };
+    use crate::core::air::{Air, AirProver, Component, ComponentProver, ComponentTrace};
     use crate::core::backend::cpu::CpuCircleEvaluation;
     use crate::core::backend::CpuBackend;
     use crate::core::channel::Blake2sChannel;
@@ -374,6 +372,8 @@ mod tests {
     use crate::core::test_utils::test_channel;
     use crate::core::{ColumnVec, InteractionElements, LookupValues};
     use crate::qm31;
+    use crate::trace_generation::registry::ComponentRegistry;
+    use crate::trace_generation::{AirTraceGenerator, AirTraceVerifier, TraceGenerator};
 
     struct TestAir<C: ComponentProver<CpuBackend>> {
         component: C,
@@ -391,7 +391,7 @@ mod tests {
         }
     }
 
-    impl AirTraceWriter<CpuBackend> for TestAir<TestComponent> {
+    impl AirTraceGenerator<CpuBackend> for TestAir<TestComponent> {
         fn interact(
             &self,
             _trace: &ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>,
@@ -456,7 +456,18 @@ mod tests {
         }
     }
 
-    impl ComponentTraceWriter<CpuBackend> for TestComponent {
+    impl TraceGenerator<CpuBackend> for TestComponent {
+        type ComponentInputs = ();
+
+        fn add_inputs(&mut self, _inputs: &Self::ComponentInputs) {}
+
+        fn write_trace(
+            _component_id: &str,
+            _registry: &mut ComponentRegistry,
+        ) -> ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>> {
+            vec![]
+        }
+
         fn write_interaction_trace(
             &self,
             _trace: &ColumnVec<&CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>,
diff --git a/crates/prover/src/examples/fibonacci/air.rs b/crates/prover/src/examples/fibonacci/air.rs
index c8c1b8c39..40c8eea6f 100644
--- a/crates/prover/src/examples/fibonacci/air.rs
+++ b/crates/prover/src/examples/fibonacci/air.rs
@@ -1,15 +1,14 @@
 use itertools::{zip_eq, Itertools};
 
 use super::component::FibonacciComponent;
-use crate::core::air::{
-    Air, AirProver, AirTraceVerifier, AirTraceWriter, Component, ComponentProver,
-};
+use crate::core::air::{Air, AirProver, Component, ComponentProver};
 use crate::core::backend::CpuBackend;
 use crate::core::channel::Blake2sChannel;
 use crate::core::fields::m31::BaseField;
 use crate::core::poly::circle::CircleEvaluation;
 use crate::core::poly::BitReversedOrder;
 use crate::core::{ColumnVec, InteractionElements};
+use crate::trace_generation::{AirTraceGenerator, AirTraceVerifier};
 
 pub struct FibonacciAir {
     pub component: FibonacciComponent,
@@ -33,7 +32,7 @@ impl AirTraceVerifier for FibonacciAir {
     }
 }
 
-impl AirTraceWriter<CpuBackend> for FibonacciAir {
+impl AirTraceGenerator<CpuBackend> for FibonacciAir {
     fn interact(
         &self,
         _trace: &ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>,
@@ -82,7 +81,7 @@ impl AirTraceVerifier for MultiFibonacciAir {
     }
 }
 
-impl AirTraceWriter<CpuBackend> for MultiFibonacciAir {
+impl AirTraceGenerator<CpuBackend> for MultiFibonacciAir {
     fn interact(
         &self,
         _trace: &ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>,
diff --git a/crates/prover/src/examples/fibonacci/component.rs b/crates/prover/src/examples/fibonacci/component.rs
index 7045f2fd9..ccdfd88b0 100644
--- a/crates/prover/src/examples/fibonacci/component.rs
+++ b/crates/prover/src/examples/fibonacci/component.rs
@@ -4,7 +4,7 @@ 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, ComponentTraceWriter};
+use crate::core::air::{Component, ComponentProver, ComponentTrace};
 use crate::core::backend::CpuBackend;
 use crate::core::circle::{CirclePoint, Coset};
 use crate::core::constraints::{coset_vanishing, pair_vanishing};
@@ -17,6 +17,8 @@ use crate::core::poly::BitReversedOrder;
 use crate::core::prover::BASE_TRACE;
 use crate::core::utils::bit_reverse_index;
 use crate::core::{ColumnVec, InteractionElements, LookupValues};
+use crate::trace_generation::registry::ComponentRegistry;
+use crate::trace_generation::TraceGenerator;
 
 pub struct FibonacciComponent {
     pub log_size: u32,
@@ -128,7 +130,18 @@ impl Component for FibonacciComponent {
     }
 }
 
-impl ComponentTraceWriter<CpuBackend> for FibonacciComponent {
+impl TraceGenerator<CpuBackend> for FibonacciComponent {
+    type ComponentInputs = ();
+
+    fn add_inputs(&mut self, _inputs: &Self::ComponentInputs) {}
+
+    fn write_trace(
+        _component_id: &str,
+        _registry: &mut ComponentRegistry,
+    ) -> ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>> {
+        vec![]
+    }
+
     fn write_interaction_trace(
         &self,
         _trace: &ColumnVec<&CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>,
diff --git a/crates/prover/src/examples/wide_fibonacci/component.rs b/crates/prover/src/examples/wide_fibonacci/component.rs
index e57d163ea..e2226a037 100644
--- a/crates/prover/src/examples/wide_fibonacci/component.rs
+++ b/crates/prover/src/examples/wide_fibonacci/component.rs
@@ -2,7 +2,7 @@ use itertools::Itertools;
 
 use crate::core::air::accumulation::PointEvaluationAccumulator;
 use crate::core::air::mask::fixed_mask_points;
-use crate::core::air::{Air, Component, ComponentTraceWriter};
+use crate::core::air::{Air, Component};
 use crate::core::backend::cpu::CpuCircleEvaluation;
 use crate::core::backend::CpuBackend;
 use crate::core::circle::{CirclePoint, Coset};
@@ -17,6 +17,8 @@ use crate::core::poly::BitReversedOrder;
 use crate::core::utils::shifted_secure_combination;
 use crate::core::{ColumnVec, InteractionElements, LookupValues};
 use crate::examples::wide_fibonacci::trace_gen::write_lookup_column;
+use crate::trace_generation::registry::ComponentRegistry;
+use crate::trace_generation::TraceGenerator;
 
 pub const LOG_N_COLUMNS: usize = 8;
 pub const N_COLUMNS: usize = 1 << LOG_N_COLUMNS;
@@ -252,7 +254,18 @@ impl Component for WideFibComponent {
     }
 }
 
-impl ComponentTraceWriter<CpuBackend> for WideFibComponent {
+impl TraceGenerator<CpuBackend> for WideFibComponent {
+    type ComponentInputs = ();
+
+    fn add_inputs(&mut self, _inputs: &Self::ComponentInputs) {}
+
+    fn write_trace(
+        _component_id: &str,
+        _registry: &mut ComponentRegistry,
+    ) -> ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>> {
+        vec![]
+    }
+
     fn write_interaction_trace(
         &self,
         trace: &ColumnVec<&CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>,
diff --git a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs
index 664a17804..a67f2b299 100644
--- a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs
+++ b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs
@@ -9,10 +9,7 @@ use super::component::{
 };
 use super::trace_gen::write_trace_row;
 use crate::core::air::accumulation::{ColumnAccumulator, DomainEvaluationAccumulator};
-use crate::core::air::{
-    AirProver, AirTraceVerifier, AirTraceWriter, Component, ComponentProver, ComponentTrace,
-    ComponentTraceWriter,
-};
+use crate::core::air::{AirProver, Component, ComponentProver, ComponentTrace};
 use crate::core::backend::CpuBackend;
 use crate::core::channel::{Blake2sChannel, Channel};
 use crate::core::circle::Coset;
@@ -29,6 +26,7 @@ use crate::core::utils::{
 };
 use crate::core::{ColumnVec, InteractionElements, LookupValues};
 use crate::examples::wide_fibonacci::component::LOG_N_COLUMNS;
+use crate::trace_generation::{AirTraceGenerator, AirTraceVerifier, TraceGenerator};
 
 // TODO(AlonH): Rename file to `cpu.rs`.
 
@@ -40,7 +38,7 @@ impl AirTraceVerifier for WideFibAir {
     }
 }
 
-impl AirTraceWriter<CpuBackend> for WideFibAir {
+impl AirTraceGenerator<CpuBackend> for WideFibAir {
     fn interact(
         &self,
         trace: &ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>,
diff --git a/crates/prover/src/examples/wide_fibonacci/mod.rs b/crates/prover/src/examples/wide_fibonacci/mod.rs
index befc6ec0e..89fb78450 100644
--- a/crates/prover/src/examples/wide_fibonacci/mod.rs
+++ b/crates/prover/src/examples/wide_fibonacci/mod.rs
@@ -13,7 +13,7 @@ mod tests {
     use super::component::{Input, WideFibAir, WideFibComponent, LOG_N_COLUMNS};
     use super::constraint_eval::gen_trace;
     use crate::core::air::accumulation::DomainEvaluationAccumulator;
-    use crate::core::air::{Component, ComponentProver, ComponentTrace, ComponentTraceWriter};
+    use crate::core::air::{Component, ComponentProver, ComponentTrace};
     use crate::core::backend::cpu::CpuCircleEvaluation;
     use crate::core::backend::CpuBackend;
     use crate::core::channel::{Blake2sChannel, Channel};
@@ -30,6 +30,7 @@ mod tests {
     use crate::core::vcs::hasher::Hasher;
     use crate::core::InteractionElements;
     use crate::examples::wide_fibonacci::trace_gen::write_lookup_column;
+    use crate::trace_generation::TraceGenerator;
     use crate::{m31, qm31};
 
     pub fn assert_constraints_on_row(row: &[BaseField]) {
diff --git a/crates/prover/src/examples/wide_fibonacci/simd.rs b/crates/prover/src/examples/wide_fibonacci/simd.rs
index d2759eb4c..e88dffd7b 100644
--- a/crates/prover/src/examples/wide_fibonacci/simd.rs
+++ b/crates/prover/src/examples/wide_fibonacci/simd.rs
@@ -5,10 +5,7 @@ use tracing::{span, Level};
 use super::component::LOG_N_COLUMNS;
 use crate::core::air::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use crate::core::air::mask::fixed_mask_points;
-use crate::core::air::{
-    Air, AirProver, AirTraceVerifier, AirTraceWriter, Component, ComponentProver, ComponentTrace,
-    ComponentTraceWriter,
-};
+use crate::core::air::{Air, AirProver, Component, ComponentProver, ComponentTrace};
 use crate::core::backend::simd::column::BaseFieldVec;
 use crate::core::backend::simd::m31::{PackedBaseField, LOG_N_LANES};
 use crate::core::backend::simd::qm31::PackedSecureField;
@@ -26,6 +23,8 @@ use crate::core::poly::BitReversedOrder;
 use crate::core::prover::BASE_TRACE;
 use crate::core::{ColumnVec, InteractionElements, LookupValues};
 use crate::examples::wide_fibonacci::component::{ALPHA_ID, N_COLUMNS, Z_ID};
+use crate::trace_generation::registry::ComponentRegistry;
+use crate::trace_generation::{AirTraceGenerator, AirTraceVerifier, TraceGenerator};
 
 // TODO(AlonH): Remove this once the Cpu and Simd implementations are aligned.
 pub struct SimdWideFibComponent {
@@ -66,7 +65,7 @@ impl AirTraceVerifier for SimdWideFibAir {
     }
 }
 
-impl AirTraceWriter<SimdBackend> for SimdWideFibAir {
+impl AirTraceGenerator<SimdBackend> for SimdWideFibAir {
     fn interact(
         &self,
         _trace: &ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>>,
@@ -162,7 +161,18 @@ pub fn gen_trace(
 }
 
 // TODO(AlonH): Implement.
-impl ComponentTraceWriter<SimdBackend> for SimdWideFibComponent {
+impl TraceGenerator<SimdBackend> for SimdWideFibComponent {
+    type ComponentInputs = ();
+
+    fn add_inputs(&mut self, _inputs: &Self::ComponentInputs) {}
+
+    fn write_trace(
+        _component_id: &str,
+        _registry: &mut ComponentRegistry,
+    ) -> ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>> {
+        vec![]
+    }
+
     fn write_interaction_trace(
         &self,
         _trace: &ColumnVec<&CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>>,
diff --git a/crates/prover/src/trace_generation/mod.rs b/crates/prover/src/trace_generation/mod.rs
index b05f4d201..ab1e2e0f8 100644
--- a/crates/prover/src/trace_generation/mod.rs
+++ b/crates/prover/src/trace_generation/mod.rs
@@ -3,11 +3,13 @@ pub mod registry;
 use downcast_rs::{impl_downcast, Downcast};
 use registry::ComponentRegistry;
 
+use crate::core::air::AirProver;
 use crate::core::backend::Backend;
+use crate::core::channel::Blake2sChannel;
 use crate::core::fields::m31::BaseField;
 use crate::core::poly::circle::CircleEvaluation;
 use crate::core::poly::BitReversedOrder;
-use crate::core::ColumnVec;
+use crate::core::{ColumnVec, InteractionElements};
 
 pub trait ComponentGen: Downcast {}
 impl_downcast!(ComponentGen);
@@ -30,4 +32,24 @@ pub trait TraceGenerator<B: Backend> {
         component_id: &str,
         registry: &mut ComponentRegistry,
     ) -> ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>;
+
+    fn write_interaction_trace(
+        &self,
+        trace: &ColumnVec<&CircleEvaluation<B, BaseField, BitReversedOrder>>,
+        elements: &InteractionElements,
+    ) -> ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>;
+}
+
+pub trait AirTraceVerifier {
+    fn interaction_elements(&self, channel: &mut Blake2sChannel) -> InteractionElements;
+}
+
+pub trait AirTraceGenerator<B: Backend>: AirTraceVerifier {
+    fn interact(
+        &self,
+        trace: &ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
+        elements: &InteractionElements,
+    ) -> Vec<CircleEvaluation<B, BaseField, BitReversedOrder>>;
+
+    fn to_air_prover(&self) -> &impl AirProver<B>;
 }
diff --git a/crates/prover/src/trace_generation/registry.rs b/crates/prover/src/trace_generation/registry.rs
index 7da8a79a8..5f26a5abc 100644
--- a/crates/prover/src/trace_generation/registry.rs
+++ b/crates/prover/src/trace_generation/registry.rs
@@ -37,7 +37,7 @@ mod tests {
     use crate::core::fields::m31::BaseField;
     use crate::core::poly::circle::CircleEvaluation;
     use crate::core::poly::BitReversedOrder;
-    use crate::core::ColumnVec;
+    use crate::core::{ColumnVec, InteractionElements};
     use crate::trace_generation::TraceGenerator;
 
     #[derive(Default)]
@@ -60,6 +60,14 @@ mod tests {
         ) -> ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>> {
             unimplemented!("TestTraceGenerator::write_trace")
         }
+
+        fn write_interaction_trace(
+            &self,
+            _trace: &ColumnVec<&CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>,
+            _elements: &InteractionElements,
+        ) -> ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>> {
+            unimplemented!("TestTraceGenerator::write_interaction_trace")
+        }
     }
 
     #[test]