From 915e9249c7f8140dff249e66f31be588c8495d08 Mon Sep 17 00:00:00 2001
From: Shahar Papini <spapini@starkware.co>
Date: Sun, 7 Jul 2024 11:41:10 +0300
Subject: [PATCH] ProverCommitmentScheme::commit_on_evals()

---
 crates/prover/src/core/pcs/prover.rs       | 20 +++++++++++++--
 crates/prover/src/core/prover/mod.rs       | 29 ++++++----------------
 crates/prover/src/examples/poseidon/mod.rs |  2 +-
 3 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/crates/prover/src/core/pcs/prover.rs b/crates/prover/src/core/pcs/prover.rs
index 40472df83..26a3d9026 100644
--- a/crates/prover/src/core/pcs/prover.rs
+++ b/crates/prover/src/core/pcs/prover.rs
@@ -49,11 +49,27 @@ impl<B: Backend + MerkleOps<MerkleHasher>> CommitmentSchemeProver<B> {
         channel: &mut ProofChannel,
         twiddles: &TwiddleTree<B>,
     ) {
+        let _span = span!(Level::INFO, "Commitment").entered();
         let tree =
             CommitmentTreeProver::new(polynomials, self.log_blowup_factor, channel, twiddles);
         self.trees.push(tree);
     }
 
+    pub fn commit_on_evals(
+        &mut self,
+        evals: ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
+        channel: &mut Blake2sChannel,
+        twiddles: &TwiddleTree<B>,
+    ) {
+        let span = span!(Level::INFO, "Interpolation").entered();
+        let polys = evals
+            .into_iter()
+            .map(|eval| eval.interpolate_with_twiddles(twiddles))
+            .collect();
+        span.exit();
+        self.commit(polys, channel, twiddles);
+    }
+
     pub fn roots(&self) -> TreeVec<Blake2sHash> {
         self.trees.as_ref().map(|tree| tree.commitment.root())
     }
@@ -157,7 +173,7 @@ impl<B: Backend + MerkleOps<MerkleHasher>> CommitmentTreeProver<B> {
         channel: &mut ProofChannel,
         twiddles: &TwiddleTree<B>,
     ) -> Self {
-        let span = span!(Level::INFO, "Commitment evaluation").entered();
+        let span = span!(Level::INFO, "Extension").entered();
         let evaluations = polynomials
             .iter()
             .map(|poly| {
@@ -170,7 +186,7 @@ impl<B: Backend + MerkleOps<MerkleHasher>> CommitmentTreeProver<B> {
 
         span.exit();
 
-        let _span = span!(Level::INFO, "Commitment merkle").entered();
+        let _span = span!(Level::INFO, "Merkle").entered();
         let tree = MerkleProver::commit(evaluations.iter().map(|eval| &eval.values).collect());
         channel.mix_digest(tree.root());
 
diff --git a/crates/prover/src/core/prover/mod.rs b/crates/prover/src/core/prover/mod.rs
index 72d0d1994..e08a7fc3e 100644
--- a/crates/prover/src/core/prover/mod.rs
+++ b/crates/prover/src/core/prover/mod.rs
@@ -55,30 +55,17 @@ pub fn evaluate_and_commit_on_trace<B: Backend + MerkleOps<MerkleHasher>>(
     twiddles: &TwiddleTree<B>,
     trace: ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
 ) -> Result<(CommitmentSchemeProver<B>, InteractionElements), ProvingError> {
-    let span = span!(Level::INFO, "Trace interpolation").entered();
-    // TODO(AlonH): Clone only the columns needed for interaction.
-    let trace_polys = trace
-        .clone()
-        .into_iter()
-        .map(|eval| eval.interpolate_with_twiddles(twiddles))
-        .collect();
-    span.exit();
-
     let mut commitment_scheme = CommitmentSchemeProver::new(LOG_BLOWUP_FACTOR);
-    let span = span!(Level::INFO, "Trace commitment").entered();
-    commitment_scheme.commit(trace_polys, channel, twiddles);
+    // TODO(spapini): Remove clone.
+    let span = span!(Level::INFO, "Trace").entered();
+    commitment_scheme.commit_on_evals(trace.clone(), channel, twiddles);
     span.exit();
 
     let interaction_elements = air.interaction_elements(channel);
     let interaction_trace = air.interact(&trace, &interaction_elements);
     if !interaction_trace.is_empty() {
-        let span = span!(Level::INFO, "Interaction trace interpolation").entered();
-        let interaction_trace_polys = interaction_trace
-            .into_iter()
-            .map(|eval| eval.interpolate_with_twiddles(twiddles))
-            .collect();
-        span.exit();
-        commitment_scheme.commit(interaction_trace_polys, channel, twiddles);
+        let _span = span!(Level::INFO, "Interaction").entered();
+        commitment_scheme.commit_on_evals(interaction_trace, channel, twiddles);
     }
 
     Ok((commitment_scheme, interaction_elements))
@@ -94,15 +81,15 @@ pub fn generate_proof<B: Backend + MerkleOps<MerkleHasher>>(
     // Evaluate and commit on composition polynomial.
     let random_coeff = channel.draw_felt();
 
-    let span = span!(Level::INFO, "Composition generation").entered();
+    let span = span!(Level::INFO, "Composition").entered();
+    let span1 = span!(Level::INFO, "Generation").entered();
     let composition_polynomial_poly = air.compute_composition_polynomial(
         random_coeff,
         &air.component_traces(&commitment_scheme.trees),
         interaction_elements,
     );
-    span.exit();
+    span1.exit();
 
-    let span = span!(Level::INFO, "Composition commitment").entered();
     commitment_scheme.commit(composition_polynomial_poly.to_vec(), channel, twiddles);
     span.exit();
 
diff --git a/crates/prover/src/examples/poseidon/mod.rs b/crates/prover/src/examples/poseidon/mod.rs
index 49fd132f7..60c67c5be 100644
--- a/crates/prover/src/examples/poseidon/mod.rs
+++ b/crates/prover/src/examples/poseidon/mod.rs
@@ -347,7 +347,7 @@ impl ComponentProver<SimdBackend> for PoseidonComponent {
         let eval_domain = CanonicCoset::new(self.log_column_size() + LOG_EXPAND).circle_domain();
 
         // Create a new evaluation.
-        let span = span!(Level::INFO, "Deg8 eval").entered();
+        let span = span!(Level::INFO, "Deg4 eval").entered();
         let twiddles = SimdBackend::precompute_twiddles(
             CanonicCoset::new(self.max_constraint_log_degree_bound())
                 .circle_domain()