From 3e0fefffc44d1b5872e0319b341cb68f958d76a9 Mon Sep 17 00:00:00 2001
From: Ian Joiner <14581281+iajoiner@users.noreply.github.com>
Date: Wed, 11 Dec 2024 16:14:54 -0500
Subject: [PATCH] feat: allow `FirstRoundBuilder` to produce MLEs

---
 .../src/sql/proof/first_round_builder.rs      |  74 ++++++-
 .../src/sql/proof/first_round_builder_test.rs |  73 +++++++
 crates/proof-of-sql/src/sql/proof/mod.rs      |   2 +
 .../proof-of-sql/src/sql/proof/proof_plan.rs  |   2 +-
 .../proof-of-sql/src/sql/proof/query_proof.rs | 106 +++++++---
 .../src/sql/proof/query_proof_test.rs         | 193 +++++++++++++++++-
 .../sql/proof/verifiable_query_result_test.rs |   2 +-
 .../verifiable_query_result_test_utility.rs   |  19 +-
 .../src/sql/proof_gadgets/range_check_test.rs |   2 +-
 .../src/sql/proof_plans/demo_mock_plan.rs     |   2 +-
 .../src/sql/proof_plans/empty_exec.rs         |   2 +-
 .../src/sql/proof_plans/filter_exec.rs        |   2 +-
 .../filter_exec_test_dishonest_prover.rs      |   2 +-
 .../src/sql/proof_plans/group_by_exec.rs      |   2 +-
 .../src/sql/proof_plans/projection_exec.rs    |   2 +-
 .../src/sql/proof_plans/slice_exec.rs         |   2 +-
 .../src/sql/proof_plans/table_exec.rs         |   2 +-
 .../src/sql/proof_plans/union_exec.rs         |   2 +-
 18 files changed, 431 insertions(+), 60 deletions(-)
 create mode 100644 crates/proof-of-sql/src/sql/proof/first_round_builder_test.rs

diff --git a/crates/proof-of-sql/src/sql/proof/first_round_builder.rs b/crates/proof-of-sql/src/sql/proof/first_round_builder.rs
index 15989adf4..432c04123 100644
--- a/crates/proof-of-sql/src/sql/proof/first_round_builder.rs
+++ b/crates/proof-of-sql/src/sql/proof/first_round_builder.rs
@@ -1,6 +1,16 @@
-use alloc::vec::Vec;
+use crate::{
+    base::{
+        commitment::{Commitment, CommittableColumn, VecCommitmentExt},
+        polynomial::MultilinearExtension,
+        scalar::Scalar,
+    },
+    utils::log,
+};
+use alloc::{boxed::Box, vec::Vec};
 /// Track the result created by a query
-pub struct FirstRoundBuilder {
+pub struct FirstRoundBuilder<'a, S> {
+    commitment_descriptor: Vec<CommittableColumn<'a>>,
+    pcs_proof_mles: Vec<Box<dyn MultilinearExtension<S> + 'a>>,
     /// The number of challenges used in the proof.
     /// Specifically, these are the challenges that the verifier sends to
     /// the prover after the prover sends the result, but before the prover
@@ -10,20 +20,26 @@ pub struct FirstRoundBuilder {
     one_evaluation_lengths: Vec<usize>,
 }
 
-impl Default for FirstRoundBuilder {
+impl<'a, S: Scalar> Default for FirstRoundBuilder<'a, S> {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl FirstRoundBuilder {
+impl<'a, S: Scalar> FirstRoundBuilder<'a, S> {
     pub fn new() -> Self {
         Self {
+            commitment_descriptor: Vec::new(),
+            pcs_proof_mles: Vec::new(),
             num_post_result_challenges: 0,
             one_evaluation_lengths: Vec::new(),
         }
     }
 
+    pub fn pcs_proof_mles(&self) -> &[Box<dyn MultilinearExtension<S> + 'a>] {
+        &self.pcs_proof_mles
+    }
+
     /// Get the one evaluation lengths used in the proof.
     pub(crate) fn one_evaluation_lengths(&self) -> &[usize] {
         &self.one_evaluation_lengths
@@ -34,6 +50,56 @@ impl FirstRoundBuilder {
         self.one_evaluation_lengths.push(length);
     }
 
+    /// Produce an MLE for a intermediate computed column that we can reference in sumcheck.
+    ///
+    /// Because the verifier doesn't have access to the MLE's commitment, we will need to
+    /// commit to the MLE before we form the sumcheck polynomial.
+    pub fn produce_intermediate_mle(
+        &mut self,
+        data: impl MultilinearExtension<S> + Into<CommittableColumn<'a>> + Copy + 'a,
+    ) {
+        self.commitment_descriptor.push(data.into());
+        self.pcs_proof_mles.push(Box::new(data));
+    }
+
+    /// Compute commitments of all the interemdiate MLEs used in sumcheck
+    #[tracing::instrument(
+        name = "FirstRoundBuilder::commit_intermediate_mles",
+        level = "debug",
+        skip_all
+    )]
+    pub fn commit_intermediate_mles<C: Commitment>(
+        &self,
+        offset_generators: usize,
+        setup: &C::PublicSetup<'_>,
+    ) -> Vec<C> {
+        Vec::from_commitable_columns_with_offset(
+            &self.commitment_descriptor,
+            offset_generators,
+            setup,
+        )
+    }
+
+    /// Given the evaluation vector, compute evaluations of all the MLEs used in sumcheck except
+    /// for those that correspond to result columns sent to the verifier.
+    #[tracing::instrument(
+        name = "FirstRoundBuilder::evaluate_pcs_proof_mles",
+        level = "debug",
+        skip_all
+    )]
+    pub fn evaluate_pcs_proof_mles(&self, evaluation_vec: &[S]) -> Vec<S> {
+        log::log_memory_usage("Start");
+
+        let mut res = Vec::with_capacity(self.pcs_proof_mles.len());
+        for evaluator in &self.pcs_proof_mles {
+            res.push(evaluator.inner_product(evaluation_vec));
+        }
+
+        log::log_memory_usage("End");
+
+        res
+    }
+
     /// The number of challenges used in the proof.
     /// Specifically, these are the challenges that the verifier sends to
     /// the prover after the prover sends the result, but before the prover
diff --git a/crates/proof-of-sql/src/sql/proof/first_round_builder_test.rs b/crates/proof-of-sql/src/sql/proof/first_round_builder_test.rs
new file mode 100644
index 000000000..dc2720ae1
--- /dev/null
+++ b/crates/proof-of-sql/src/sql/proof/first_round_builder_test.rs
@@ -0,0 +1,73 @@
+use super::FirstRoundBuilder;
+use crate::base::{
+    commitment::{Commitment, CommittableColumn},
+    scalar::Curve25519Scalar,
+};
+use curve25519_dalek::RistrettoPoint;
+
+#[test]
+fn we_can_compute_commitments_for_intermediate_mles_using_a_zero_offset() {
+    let mle1 = [1, 2];
+    let mle2 = [10i64, 20];
+    let mut builder = FirstRoundBuilder::<Curve25519Scalar>::new();
+    builder.produce_intermediate_mle(&mle1[..]);
+    builder.produce_intermediate_mle(&mle2[..]);
+    let offset_generators = 0_usize;
+    let commitments: Vec<RistrettoPoint> = builder.commit_intermediate_mles(offset_generators, &());
+    assert_eq!(
+        commitments,
+        [RistrettoPoint::compute_commitments(
+            &[CommittableColumn::from(&mle2[..])],
+            offset_generators,
+            &()
+        )[0]]
+    );
+}
+
+#[test]
+fn we_can_compute_commitments_for_intermediate_mles_using_a_non_zero_offset() {
+    let mle1 = [1, 2];
+    let mle2 = [10i64, 20];
+    let mut builder = FirstRoundBuilder::<Curve25519Scalar>::new();
+    builder.produce_intermediate_mle(&mle1[..]);
+    builder.produce_intermediate_mle(&mle2[..]);
+    let offset_generators = 123_usize;
+    let commitments: Vec<RistrettoPoint> = builder.commit_intermediate_mles(offset_generators, &());
+    assert_eq!(
+        commitments,
+        [RistrettoPoint::compute_commitments(
+            &[CommittableColumn::from(&mle2[..])],
+            offset_generators,
+            &()
+        )[0]]
+    );
+}
+
+#[test]
+fn we_can_evaluate_pcs_proof_mles() {
+    let mle1 = [1, 2];
+    let mle2 = [10i64, 20];
+    let mut builder = FirstRoundBuilder::<Curve25519Scalar>::new();
+    builder.produce_intermediate_mle(&mle1[..]);
+    builder.produce_intermediate_mle(&mle2[..]);
+    let evaluation_vec = [
+        Curve25519Scalar::from(100u64),
+        Curve25519Scalar::from(10u64),
+    ];
+    let evals = builder.evaluate_pcs_proof_mles(&evaluation_vec);
+    let expected_evals = [
+        Curve25519Scalar::from(120u64),
+        Curve25519Scalar::from(1200u64),
+    ];
+    assert_eq!(evals, expected_evals);
+}
+
+#[test]
+fn we_can_add_post_result_challenges() {
+    let mut builder = FirstRoundBuilder::<Curve25519Scalar>::new();
+    assert_eq!(builder.num_post_result_challenges(), 0);
+    builder.request_post_result_challenges(1);
+    assert_eq!(builder.num_post_result_challenges(), 1);
+    builder.request_post_result_challenges(2);
+    assert_eq!(builder.num_post_result_challenges(), 3);
+}
diff --git a/crates/proof-of-sql/src/sql/proof/mod.rs b/crates/proof-of-sql/src/sql/proof/mod.rs
index 96dfc5874..7b64a99d5 100644
--- a/crates/proof-of-sql/src/sql/proof/mod.rs
+++ b/crates/proof-of-sql/src/sql/proof/mod.rs
@@ -63,6 +63,8 @@ pub(crate) use result_element_serialization::{
 
 mod first_round_builder;
 pub(crate) use first_round_builder::FirstRoundBuilder;
+#[cfg(all(test, feature = "blitzar"))]
+mod first_round_builder_test;
 
 #[cfg(all(test, feature = "arrow"))]
 mod provable_query_result_test;
diff --git a/crates/proof-of-sql/src/sql/proof/proof_plan.rs b/crates/proof-of-sql/src/sql/proof/proof_plan.rs
index a84dcdab7..f5c5c5fbe 100644
--- a/crates/proof-of-sql/src/sql/proof/proof_plan.rs
+++ b/crates/proof-of-sql/src/sql/proof/proof_plan.rs
@@ -36,7 +36,7 @@ pub trait ProverEvaluate {
     /// Evaluate the query, modify `FirstRoundBuilder` and return the result.
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S>;
diff --git a/crates/proof-of-sql/src/sql/proof/query_proof.rs b/crates/proof-of-sql/src/sql/proof/query_proof.rs
index 708b330d1..e95707d99 100644
--- a/crates/proof-of-sql/src/sql/proof/query_proof.rs
+++ b/crates/proof-of-sql/src/sql/proof/query_proof.rs
@@ -57,12 +57,16 @@ pub(super) struct QueryProof<CP: CommitmentEvaluationProof> {
     pub bit_distributions: Vec<BitDistribution>,
     /// One evaluation lengths
     pub one_evaluation_lengths: Vec<usize>,
-    /// Commitments
-    pub commitments: Vec<CP::Commitment>,
+    /// First Round Commitments
+    pub first_round_commitments: Vec<CP::Commitment>,
+    /// Final Round Commitments
+    pub final_round_commitments: Vec<CP::Commitment>,
     /// Sumcheck Proof
     pub sumcheck_proof: SumcheckProof<CP::Scalar>,
-    /// MLEs used in sumcheck except for the result columns
-    pub pcs_proof_evaluations: Vec<CP::Scalar>,
+    /// MLEs used in first round sumcheck except for the result columns
+    pub first_round_pcs_proof_evaluations: Vec<CP::Scalar>,
+    /// MLEs used in final round sumcheck except for the result columns
+    pub final_round_pcs_proof_evaluations: Vec<CP::Scalar>,
     /// Inner product proof of the MLEs' evaluations
     pub evaluation_proof: CP,
     /// Length of the range of generators we use
@@ -100,7 +104,7 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
             .collect();
 
         // Prover First Round: Evaluate the query && get the right number of post result challenges
-        let mut first_round_builder = FirstRoundBuilder::new();
+        let mut first_round_builder = FirstRoundBuilder::<CP::Scalar>::new();
         let query_result = expr.first_round_evaluate(&mut first_round_builder, &alloc, &table_map);
         let owned_table_result = OwnedTable::from(&query_result);
         let provable_result = query_result.into();
@@ -117,6 +121,10 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
         assert!(num_sumcheck_variables > 0);
         let post_result_challenge_count = first_round_builder.num_post_result_challenges();
 
+        // commit to any intermediate MLEs
+        let first_round_commitments =
+            first_round_builder.commit_intermediate_mles(min_row_num, setup);
+
         // construct a transcript for the proof
         let mut transcript: Keccak256Transcript = make_transcript(
             expr,
@@ -125,6 +133,7 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
             min_row_num,
             one_evaluation_lengths,
             post_result_challenge_count,
+            &first_round_commitments,
         );
 
         // These are the challenges that will be consumed by the proof
@@ -137,35 +146,37 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
                 .take(post_result_challenge_count)
                 .collect();
 
-        let mut builder = FinalRoundBuilder::new(num_sumcheck_variables, post_result_challenges);
+        let mut final_round_builder =
+            FinalRoundBuilder::new(num_sumcheck_variables, post_result_challenges);
 
         for col_ref in total_col_refs {
-            builder.produce_anchored_mle(accessor.get_column(col_ref));
+            final_round_builder.produce_anchored_mle(accessor.get_column(col_ref));
         }
 
-        expr.final_round_evaluate(&mut builder, &alloc, &table_map);
+        expr.final_round_evaluate(&mut final_round_builder, &alloc, &table_map);
 
-        let num_sumcheck_variables = builder.num_sumcheck_variables();
+        let num_sumcheck_variables = final_round_builder.num_sumcheck_variables();
 
         // commit to any intermediate MLEs
-        let commitments = builder.commit_intermediate_mles(min_row_num, setup);
+        let final_round_commitments =
+            final_round_builder.commit_intermediate_mles(min_row_num, setup);
 
         // add the commitments, bit distributions and one evaluation lengths to the proof
         extend_transcript_with_commitments(
             &mut transcript,
-            &commitments,
-            builder.bit_distributions(),
+            &final_round_commitments,
+            final_round_builder.bit_distributions(),
         );
 
         // construct the sumcheck polynomial
-        let subpolynomial_constraint_count = builder.num_sumcheck_subpolynomials();
+        let subpolynomial_constraint_count = final_round_builder.num_sumcheck_subpolynomials();
         let num_random_scalars = num_sumcheck_variables + subpolynomial_constraint_count;
         let random_scalars: Vec<_> =
             core::iter::repeat_with(|| transcript.scalar_challenge_as_be())
                 .take(num_random_scalars)
                 .collect();
         let state = make_sumcheck_prover_state(
-            builder.sumcheck_subpolynomials(),
+            final_round_builder.sumcheck_subpolynomials(),
             num_sumcheck_variables,
             &SumcheckRandomScalars::new(&random_scalars, range_length, num_sumcheck_variables),
         );
@@ -177,21 +188,36 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
         // evaluate the MLEs used in sumcheck except for the result columns
         let mut evaluation_vec = vec![Zero::zero(); range_length];
         compute_evaluation_vector(&mut evaluation_vec, &evaluation_point);
-        let pcs_proof_evaluations = builder.evaluate_pcs_proof_mles(&evaluation_vec);
+        let first_round_pcs_proof_evaluations =
+            first_round_builder.evaluate_pcs_proof_mles(&evaluation_vec);
+        let final_round_pcs_proof_evaluations =
+            final_round_builder.evaluate_pcs_proof_mles(&evaluation_vec);
 
         // commit to the MLE evaluations
-        transcript.extend_canonical_serialize_as_le(&pcs_proof_evaluations);
+        transcript.extend_canonical_serialize_as_le(&first_round_pcs_proof_evaluations);
+        transcript.extend_canonical_serialize_as_le(&final_round_pcs_proof_evaluations);
 
         // fold together the pre result MLEs -- this will form the input to an inner product proof
         // of their evaluations (fold in this context means create a random linear combination)
         let random_scalars: Vec<_> =
             core::iter::repeat_with(|| transcript.scalar_challenge_as_be())
-                .take(pcs_proof_evaluations.len())
+                .take(
+                    first_round_pcs_proof_evaluations.len()
+                        + final_round_pcs_proof_evaluations.len(),
+                )
                 .collect();
 
-        assert_eq!(random_scalars.len(), builder.pcs_proof_mles().len());
+        assert_eq!(
+            random_scalars.len(),
+            first_round_builder.pcs_proof_mles().len() + final_round_builder.pcs_proof_mles().len()
+        );
         let mut folded_mle = vec![Zero::zero(); range_length];
-        for (multiplier, evaluator) in random_scalars.iter().zip(builder.pcs_proof_mles().iter()) {
+        for (multiplier, evaluator) in random_scalars.iter().zip(
+            first_round_builder
+                .pcs_proof_mles()
+                .iter()
+                .chain(final_round_builder.pcs_proof_mles().iter()),
+        ) {
             evaluator.mul_add(&mut folded_mle, multiplier);
         }
 
@@ -205,11 +231,13 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
         );
 
         let proof = Self {
-            bit_distributions: builder.bit_distributions().to_vec(),
+            bit_distributions: final_round_builder.bit_distributions().to_vec(),
             one_evaluation_lengths: one_evaluation_lengths.to_vec(),
-            commitments,
+            first_round_commitments,
+            final_round_commitments,
             sumcheck_proof,
-            pcs_proof_evaluations,
+            first_round_pcs_proof_evaluations,
+            final_round_pcs_proof_evaluations,
             evaluation_proof,
             range_length,
             subpolynomial_constraint_count,
@@ -260,6 +288,7 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
             min_row_num,
             &self.one_evaluation_lengths,
             self.post_result_challenge_count,
+            &self.first_round_commitments,
         );
 
         // These are the challenges that will be consumed by the proof
@@ -272,10 +301,10 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
                 .take(self.post_result_challenge_count)
                 .collect();
 
-        // add the commitments and bit disctibutions to the proof
+        // add the commitments and bit distributions to the proof
         extend_transcript_with_commitments(
             &mut transcript,
-            &self.commitments,
+            &self.final_round_commitments,
             &self.bit_distributions,
         );
 
@@ -298,13 +327,17 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
         )?;
 
         // commit to mle evaluations
-        transcript.extend_canonical_serialize_as_le(&self.pcs_proof_evaluations);
+        transcript.extend_canonical_serialize_as_le(&self.first_round_pcs_proof_evaluations);
+        transcript.extend_canonical_serialize_as_le(&self.final_round_pcs_proof_evaluations);
 
         // draw the random scalars for the evaluation proof
         // (i.e. the folding/random linear combination of the pcs_proof_mles)
         let evaluation_random_scalars: Vec<_> =
             core::iter::repeat_with(|| transcript.scalar_challenge_as_be())
-                .take(self.pcs_proof_evaluations.len())
+                .take(
+                    self.first_round_pcs_proof_evaluations.len()
+                        + self.final_round_pcs_proof_evaluations.len(),
+                )
                 .collect();
 
         // Always prepend input lengths to the one evaluation lengths
@@ -324,7 +357,7 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
             one_evaluation_lengths,
             &subclaim.evaluation_point,
             &sumcheck_random_scalars,
-            &self.pcs_proof_evaluations,
+            &self.final_round_pcs_proof_evaluations,
         );
         let one_eval_map: IndexMap<TableRef, CP::Scalar> = table_length_map
             .iter()
@@ -343,7 +376,8 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
         let pcs_proof_commitments: Vec<_> = column_references
             .iter()
             .map(|col| accessor.get_commitment(col.clone()))
-            .chain(self.commitments.iter().cloned())
+            .chain(self.first_round_commitments.iter().cloned())
+            .chain(self.final_round_commitments.iter().cloned())
             .collect();
         let evaluation_accessor: IndexMap<_, _> = column_references
             .into_iter()
@@ -378,7 +412,7 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
                 &mut transcript,
                 &pcs_proof_commitments,
                 &evaluation_random_scalars,
-                &self.pcs_proof_evaluations,
+                &self.final_round_pcs_proof_evaluations,
                 &subclaim.evaluation_point,
                 min_row_num as u64,
                 self.range_length,
@@ -412,17 +446,20 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
 /// * `range_length` - The length of the range of generators used.
 /// * `min_row_num` - The minimum row number in the index range of the tables referenced by the query.
 /// * `one_evaluation_lengths` - The lengths of the one evaluations.
+/// * `post_result_challenge_count` - The number of post-result challenges.
+/// * `first_round_commitments` - A slice of commitments produced before post-result challenges that are part of the proof.
 ///
 /// # Returns
 ///
 /// A transcript initialized with the provided data.
-fn make_transcript<S: Scalar, T: Transcript>(
+fn make_transcript<C: Commitment, T: Transcript>(
     expr: &(impl ProofPlan + Serialize),
-    result: &OwnedTable<S>,
+    result: &OwnedTable<C::Scalar>,
     range_length: usize,
     min_row_num: usize,
     one_evaluation_lengths: &[usize],
     post_result_challenge_count: usize,
+    first_round_commitments: &[C],
 ) -> T {
     let mut transcript = T::new();
     extend_transcript_with_owned_table(&mut transcript, result);
@@ -431,6 +468,9 @@ fn make_transcript<S: Scalar, T: Transcript>(
     transcript.extend_serialize_as_le(&min_row_num);
     transcript.extend_serialize_as_le(one_evaluation_lengths);
     transcript.extend_serialize_as_le(&post_result_challenge_count);
+    for commitment in first_round_commitments {
+        commitment.append_to_transcript(&mut transcript);
+    }
     transcript
 }
 
@@ -486,10 +526,10 @@ fn extend_transcript_with_owned_table<S: Scalar, T: Transcript>(
 /// * `bit_distributions` - The bit distributions to add to the transcript.
 fn extend_transcript_with_commitments<C: Commitment>(
     transcript: &mut impl Transcript,
-    commitments: &[C],
+    final_round_commitments: &[C],
     bit_distributions: &[BitDistribution],
 ) {
-    for commitment in commitments {
+    for commitment in final_round_commitments {
         commitment.append_to_transcript(transcript);
     }
     transcript.extend_serialize_as_le(bit_distributions);
diff --git a/crates/proof-of-sql/src/sql/proof/query_proof_test.rs b/crates/proof-of-sql/src/sql/proof/query_proof_test.rs
index e88c7ba7b..fb8f0353a 100644
--- a/crates/proof-of-sql/src/sql/proof/query_proof_test.rs
+++ b/crates/proof-of-sql/src/sql/proof/query_proof_test.rs
@@ -48,7 +48,7 @@ impl Default for TrivialTestProofPlan {
 impl ProverEvaluate for TrivialTestProofPlan {
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
@@ -263,7 +263,7 @@ impl Default for SquareTestProofPlan {
 impl ProverEvaluate for SquareTestProofPlan {
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
@@ -442,7 +442,7 @@ impl Default for DoubleSquareTestProofPlan {
 impl ProverEvaluate for DoubleSquareTestProofPlan {
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
@@ -600,7 +600,8 @@ fn verify_fails_if_an_intermediate_commitment_doesnt_match() {
         (),
     );
     let (mut proof, result) = QueryProof::<InnerProductProof>::new(&expr, &accessor, &());
-    proof.commitments[0] = proof.commitments[0] * Curve25519Scalar::from(2u64);
+    proof.final_round_commitments[0] =
+        proof.final_round_commitments[0] * Curve25519Scalar::from(2u64);
     assert!(proof.verify(&expr, &accessor, result, &()).is_err());
 }
 
@@ -652,7 +653,7 @@ struct ChallengeTestProofPlan {}
 impl ProverEvaluate for ChallengeTestProofPlan {
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
@@ -773,3 +774,185 @@ fn we_can_verify_a_proof_with_a_post_result_challenge_and_with_a_zero_offset() {
 fn we_can_verify_a_proof_with_a_post_result_challenge_and_with_a_non_zero_offset() {
     verify_a_proof_with_a_post_result_challenge_and_given_offset(123);
 }
+
+/// prove and verify an artificial query where
+///     `res_i = x_i * x_i`
+/// where the commitment for x is known
+#[derive(Debug, Serialize)]
+struct FirstRoundSquareTestProofPlan {
+    res: [i64; 2],
+    anchored_commit_multiplier: i64,
+}
+impl Default for FirstRoundSquareTestProofPlan {
+    fn default() -> Self {
+        Self {
+            res: [9, 25],
+            anchored_commit_multiplier: 1,
+        }
+    }
+}
+impl ProverEvaluate for FirstRoundSquareTestProofPlan {
+    fn first_round_evaluate<'a, S: Scalar>(
+        &self,
+        builder: &mut FirstRoundBuilder<'a, S>,
+        alloc: &'a Bump,
+        _table_map: &IndexMap<TableRef, Table<'a, S>>,
+    ) -> Table<'a, S> {
+        let res: &[_] = alloc.alloc_slice_copy(&self.res);
+        builder.produce_intermediate_mle(res);
+        builder.produce_one_evaluation_length(2);
+        table([borrowed_bigint("a1", self.res, alloc)])
+    }
+
+    fn final_round_evaluate<'a, S: Scalar>(
+        &self,
+        builder: &mut FinalRoundBuilder<'a, S>,
+        alloc: &'a Bump,
+        table_map: &IndexMap<TableRef, Table<'a, S>>,
+    ) -> Table<'a, S> {
+        let x = *table_map
+            .get(&TableRef::new("sxt.test".parse().unwrap()))
+            .unwrap()
+            .inner_table()
+            .get(&Ident::new("x"))
+            .unwrap();
+        let res: &[_] = alloc.alloc_slice_copy(&self.res);
+        builder.produce_intermediate_mle(res);
+        builder.produce_sumcheck_subpolynomial(
+            SumcheckSubpolynomialType::Identity,
+            vec![
+                (S::ONE, vec![Box::new(res)]),
+                (-S::ONE, vec![Box::new(x), Box::new(x)]),
+            ],
+        );
+        table([borrowed_bigint("a1", self.res, alloc)])
+    }
+}
+impl ProofPlan for FirstRoundSquareTestProofPlan {
+    fn verifier_evaluate<S: Scalar>(
+        &self,
+        builder: &mut VerificationBuilder<S>,
+        accessor: &IndexMap<ColumnRef, S>,
+        _result: Option<&OwnedTable<S>>,
+        _one_eval_map: &IndexMap<TableRef, S>,
+    ) -> Result<TableEvaluation<S>, ProofError> {
+        let x_eval = S::from(self.anchored_commit_multiplier)
+            * *accessor
+                .get(&ColumnRef::new(
+                    "sxt.test".parse().unwrap(),
+                    "x".into(),
+                    ColumnType::BigInt,
+                ))
+                .unwrap();
+        let first_round_res_eval = builder.try_consume_mle_evaluation()?;
+        let final_round_res_eval = builder.try_consume_mle_evaluation()?;
+        assert_eq!(first_round_res_eval, final_round_res_eval);
+        builder.try_produce_sumcheck_subpolynomial_evaluation(
+            SumcheckSubpolynomialType::Identity,
+            final_round_res_eval - x_eval * x_eval,
+            2,
+        )?;
+        Ok(TableEvaluation::new(
+            vec![final_round_res_eval],
+            builder.try_consume_one_evaluation()?,
+        ))
+    }
+    fn get_column_result_fields(&self) -> Vec<ColumnField> {
+        vec![ColumnField::new("a1".into(), ColumnType::BigInt)]
+    }
+    fn get_column_references(&self) -> IndexSet<ColumnRef> {
+        indexset! {ColumnRef::new(
+            "sxt.test".parse().unwrap(),
+            "x".into(),
+            ColumnType::BigInt,
+        )}
+    }
+    fn get_table_references(&self) -> IndexSet<TableRef> {
+        indexset! {TableRef::new("sxt.test".parse().unwrap())}
+    }
+}
+
+fn verify_a_proof_with_a_commitment_and_given_offset(offset_generators: usize) {
+    // prove and verify an artificial query where
+    //     res_i = x_i * x_i
+    // where the commitment for x is known
+    let expr = FirstRoundSquareTestProofPlan {
+        ..Default::default()
+    };
+    let accessor = OwnedTableTestAccessor::<InnerProductProof>::new_from_table(
+        "sxt.test".parse().unwrap(),
+        owned_table([bigint("x", [3, 5])]),
+        offset_generators,
+        (),
+    );
+    let (proof, result) = QueryProof::<InnerProductProof>::new(&expr, &accessor, &());
+    let QueryData {
+        verification_hash,
+        table,
+    } = proof
+        .clone()
+        .verify(&expr, &accessor, result.clone(), &())
+        .unwrap();
+    assert_ne!(verification_hash, [0; 32]);
+    let expected_result = owned_table([bigint("a1", [9, 25])]);
+    assert_eq!(table, expected_result);
+
+    // invalid offset will fail to verify
+    let accessor = OwnedTableTestAccessor::<InnerProductProof>::new_from_table(
+        "sxt.test".parse().unwrap(),
+        owned_table([bigint("x", [3, 5])]),
+        offset_generators + 1,
+        (),
+    );
+    assert!(proof.verify(&expr, &accessor, result, &()).is_err());
+}
+
+#[test]
+fn we_can_verify_a_proof_with_a_commitment_and_with_a_zero_offset() {
+    verify_a_proof_with_a_commitment_and_given_offset(0);
+}
+
+#[test]
+fn we_can_verify_a_proof_with_a_commitment_and_with_a_non_zero_offset() {
+    verify_a_proof_with_a_commitment_and_given_offset(123);
+}
+
+#[test]
+fn verify_fails_if_the_result_doesnt_satisfy_an_equation() {
+    // attempt to prove and verify an artificial query where
+    //     res_i = x_i * x_i
+    // where the commitment for x is known and
+    //     res_i != x_i * x_i
+    // for some i
+    let expr = FirstRoundSquareTestProofPlan {
+        res: [9, 26],
+        ..Default::default()
+    };
+    let accessor = OwnedTableTestAccessor::<InnerProductProof>::new_from_table(
+        "sxt.test".parse().unwrap(),
+        owned_table([bigint("x", [3, 5])]),
+        0,
+        (),
+    );
+    let (proof, result) = QueryProof::<InnerProductProof>::new(&expr, &accessor, &());
+    assert!(proof.verify(&expr, &accessor, result, &()).is_err());
+}
+
+#[test]
+fn verify_fails_if_the_commitment_doesnt_match() {
+    // prove and verify an artificial query where
+    //     res_i = x_i * x_i
+    // where the commitment for x is known
+    let expr = FirstRoundSquareTestProofPlan {
+        anchored_commit_multiplier: 2,
+        ..Default::default()
+    };
+    let accessor = OwnedTableTestAccessor::<InnerProductProof>::new_from_table(
+        "sxt.test".parse().unwrap(),
+        owned_table([bigint("x", [3, 5])]),
+        0,
+        (),
+    );
+    let (proof, result) = QueryProof::<InnerProductProof>::new(&expr, &accessor, &());
+    assert!(proof.verify(&expr, &accessor, result, &()).is_err());
+}
diff --git a/crates/proof-of-sql/src/sql/proof/verifiable_query_result_test.rs b/crates/proof-of-sql/src/sql/proof/verifiable_query_result_test.rs
index f49a75d94..0bfe992ca 100644
--- a/crates/proof-of-sql/src/sql/proof/verifiable_query_result_test.rs
+++ b/crates/proof-of-sql/src/sql/proof/verifiable_query_result_test.rs
@@ -27,7 +27,7 @@ pub(super) struct EmptyTestQueryExpr {
 impl ProverEvaluate for EmptyTestQueryExpr {
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
diff --git a/crates/proof-of-sql/src/sql/proof/verifiable_query_result_test_utility.rs b/crates/proof-of-sql/src/sql/proof/verifiable_query_result_test_utility.rs
index 461ae0c1a..35b27376f 100644
--- a/crates/proof-of-sql/src/sql/proof/verifiable_query_result_test_utility.rs
+++ b/crates/proof-of-sql/src/sql/proof/verifiable_query_result_test_utility.rs
@@ -19,7 +19,7 @@ use serde::Serialize;
 /// Will panic if:
 /// - The verification of `res` does not succeed, causing the assertion `assert!(res.verify(...).is_ok())` to fail.
 /// - `res.proof` is `None`, causing `res.proof.as_ref().unwrap()` to panic.
-/// - Attempting to modify `pcs_proof_evaluations` or `commitments` if `res_p.proof` is `None`, leading to a panic on `unwrap()`.
+/// - Attempting to modify `final_round_pcs_proof_evaluations` or `commitments` if `res_p.proof` is `None`, leading to a panic on `unwrap()`.
 /// - `fake_accessor.update_offset` fails, causing a panic if it is designed to do so in the implementation.
 pub fn exercise_verification(
     res: &VerifiableQueryResult<InnerProductProof>,
@@ -43,9 +43,13 @@ pub fn exercise_verification(
     assert!(res_p.verify(expr, accessor, &()).is_err());
 
     // try changing MLE evaluations
-    for i in 0..proof.pcs_proof_evaluations.len() {
+    for i in 0..proof.final_round_pcs_proof_evaluations.len() {
         let mut res_p = res.clone();
-        res_p.proof.as_mut().unwrap().pcs_proof_evaluations[i] += Curve25519Scalar::one();
+        res_p
+            .proof
+            .as_mut()
+            .unwrap()
+            .final_round_pcs_proof_evaluations[i] += Curve25519Scalar::one();
         assert!(res_p.verify(expr, accessor, &()).is_err());
     }
 
@@ -59,9 +63,9 @@ pub fn exercise_verification(
         &(),
     )[0];
 
-    for i in 0..proof.commitments.len() {
+    for i in 0..proof.final_round_commitments.len() {
         let mut res_p = res.clone();
-        res_p.proof.as_mut().unwrap().commitments[i] = commit_p;
+        res_p.proof.as_mut().unwrap().final_round_commitments[i] = commit_p;
         assert!(res_p.verify(expr, accessor, &()).is_err());
     }
 
@@ -71,7 +75,10 @@ pub fn exercise_verification(
     // the inner product proof isn't dependent on the generators since it simply sends the input
     // vector; hence, changing the offset would have no effect.
     if accessor.get_length(table_ref) > 1
-        || proof.commitments.iter().any(|&c| c != Identity::identity())
+        || proof
+            .final_round_commitments
+            .iter()
+            .any(|&c| c != Identity::identity())
     {
         let offset_generators = accessor.get_offset(table_ref);
         let mut fake_accessor = accessor.clone();
diff --git a/crates/proof-of-sql/src/sql/proof_gadgets/range_check_test.rs b/crates/proof-of-sql/src/sql/proof_gadgets/range_check_test.rs
index 1909b2ed6..e7f7217a4 100644
--- a/crates/proof-of-sql/src/sql/proof_gadgets/range_check_test.rs
+++ b/crates/proof-of-sql/src/sql/proof_gadgets/range_check_test.rs
@@ -22,7 +22,7 @@ impl ProverEvaluate for RangeCheckTestPlan {
     #[doc = " Evaluate the query, modify `FirstRoundBuilder` and return the result."]
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         _alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
diff --git a/crates/proof-of-sql/src/sql/proof_plans/demo_mock_plan.rs b/crates/proof-of-sql/src/sql/proof_plans/demo_mock_plan.rs
index 7a3352caa..1ddf3c796 100644
--- a/crates/proof-of-sql/src/sql/proof_plans/demo_mock_plan.rs
+++ b/crates/proof-of-sql/src/sql/proof_plans/demo_mock_plan.rs
@@ -53,7 +53,7 @@ impl ProofPlan for DemoMockPlan {
 impl ProverEvaluate for DemoMockPlan {
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        _builder: &mut FirstRoundBuilder,
+        _builder: &mut FirstRoundBuilder<'a, S>,
         _alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
diff --git a/crates/proof-of-sql/src/sql/proof_plans/empty_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/empty_exec.rs
index 9f3cf0815..d4b8107ef 100644
--- a/crates/proof-of-sql/src/sql/proof_plans/empty_exec.rs
+++ b/crates/proof-of-sql/src/sql/proof_plans/empty_exec.rs
@@ -67,7 +67,7 @@ impl ProverEvaluate for EmptyExec {
     #[tracing::instrument(name = "EmptyExec::first_round_evaluate", level = "debug", skip_all)]
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        _builder: &mut FirstRoundBuilder,
+        _builder: &mut FirstRoundBuilder<'a, S>,
         _alloc: &'a Bump,
         _table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
diff --git a/crates/proof-of-sql/src/sql/proof_plans/filter_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/filter_exec.rs
index ed8d1a14b..cd94b56c4 100644
--- a/crates/proof-of-sql/src/sql/proof_plans/filter_exec.rs
+++ b/crates/proof-of-sql/src/sql/proof_plans/filter_exec.rs
@@ -145,7 +145,7 @@ impl ProverEvaluate for FilterExec {
     #[tracing::instrument(name = "FilterExec::first_round_evaluate", level = "debug", skip_all)]
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
diff --git a/crates/proof-of-sql/src/sql/proof_plans/filter_exec_test_dishonest_prover.rs b/crates/proof-of-sql/src/sql/proof_plans/filter_exec_test_dishonest_prover.rs
index 9ec4f5701..443b28934 100644
--- a/crates/proof-of-sql/src/sql/proof_plans/filter_exec_test_dishonest_prover.rs
+++ b/crates/proof-of-sql/src/sql/proof_plans/filter_exec_test_dishonest_prover.rs
@@ -37,7 +37,7 @@ impl ProverEvaluate for DishonestFilterExec {
     )]
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
diff --git a/crates/proof-of-sql/src/sql/proof_plans/group_by_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/group_by_exec.rs
index 3ca31d42b..7e040ba89 100644
--- a/crates/proof-of-sql/src/sql/proof_plans/group_by_exec.rs
+++ b/crates/proof-of-sql/src/sql/proof_plans/group_by_exec.rs
@@ -193,7 +193,7 @@ impl ProverEvaluate for GroupByExec {
     #[tracing::instrument(name = "GroupByExec::first_round_evaluate", level = "debug", skip_all)]
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
diff --git a/crates/proof-of-sql/src/sql/proof_plans/projection_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/projection_exec.rs
index 80e8e85b4..c47e7ecd7 100644
--- a/crates/proof-of-sql/src/sql/proof_plans/projection_exec.rs
+++ b/crates/proof-of-sql/src/sql/proof_plans/projection_exec.rs
@@ -94,7 +94,7 @@ impl ProverEvaluate for ProjectionExec {
     )]
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        _builder: &mut FirstRoundBuilder,
+        _builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
diff --git a/crates/proof-of-sql/src/sql/proof_plans/slice_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/slice_exec.rs
index aca38aaf1..36e80c8c7 100644
--- a/crates/proof-of-sql/src/sql/proof_plans/slice_exec.rs
+++ b/crates/proof-of-sql/src/sql/proof_plans/slice_exec.rs
@@ -112,7 +112,7 @@ impl ProverEvaluate for SliceExec {
     #[tracing::instrument(name = "SliceExec::first_round_evaluate", level = "debug", skip_all)]
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
diff --git a/crates/proof-of-sql/src/sql/proof_plans/table_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/table_exec.rs
index 0e4b62fba..f23235a9d 100644
--- a/crates/proof-of-sql/src/sql/proof_plans/table_exec.rs
+++ b/crates/proof-of-sql/src/sql/proof_plans/table_exec.rs
@@ -76,7 +76,7 @@ impl ProverEvaluate for TableExec {
     #[tracing::instrument(name = "TableExec::first_round_evaluate", level = "debug", skip_all)]
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        _builder: &mut FirstRoundBuilder,
+        _builder: &mut FirstRoundBuilder<'a, S>,
         _alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
diff --git a/crates/proof-of-sql/src/sql/proof_plans/union_exec.rs b/crates/proof-of-sql/src/sql/proof_plans/union_exec.rs
index 48a273398..76b948d16 100644
--- a/crates/proof-of-sql/src/sql/proof_plans/union_exec.rs
+++ b/crates/proof-of-sql/src/sql/proof_plans/union_exec.rs
@@ -108,7 +108,7 @@ impl ProverEvaluate for UnionExec {
     #[tracing::instrument(name = "UnionExec::first_round_evaluate", level = "debug", skip_all)]
     fn first_round_evaluate<'a, S: Scalar>(
         &self,
-        builder: &mut FirstRoundBuilder,
+        builder: &mut FirstRoundBuilder<'a, S>,
         alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {