feat: allow FirstRoundBuilder to produce MLEs (#417)

Please be sure to look over the pull request guidelines here:
https://github.com/spaceandtimelabs/sxt-proof-of-sql/blob/main/CONTRIBUTING.md#submit-pr.

# Please go through the following checklist
- [x] The PR title and commit messages adhere to guidelines here:
https://github.com/spaceandtimelabs/sxt-proof-of-sql/blob/main/CONTRIBUTING.md.
In particular `!` is used if and only if at least one breaking change
has been introduced.
- [x] I have run the ci check script with `source
scripts/run_ci_checks.sh`.

# Rationale for this change
For range check and joins we need to have pre-challenge committing of
columns. As a result we need to add such functionalities to
`FirstRoundBuilder`.
<!--
Why are you proposing this change? If this is already explained clearly
in the linked issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.

 Example:
 Add `NestedLoopJoinExec`.
 Closes #345.

Since we added `HashJoinExec` in #323 it has been possible to do
provable inner joins. However performance is not satisfactory in some
cases. Hence we need to fix the problem by implement
`NestedLoopJoinExec` and speed up the code
 for `HashJoinExec`.
-->

# What changes are included in this PR?
See above.
<!--
There is no need to duplicate the description in the ticket here but it
is sometimes worth providing a summary of the individual changes in this
PR.

Example:
- Add `NestedLoopJoinExec`.
- Speed up `HashJoinExec`.
- Route joins to `NestedLoopJoinExec` if the outer input is sufficiently
small.
-->

# Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?

Example:
Yes.
-->
Will be.

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

crates/proof-of-sql/src/sql/proof/first_round_builder_test.rs

@@ -0,0 +1,75 @@
+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, &());
+    let expected_commitments: Vec<RistrettoPoint> = RistrettoPoint::compute_commitments(
+        &[
+            CommittableColumn::from(&mle1[..]),
+            CommittableColumn::from(&mle2[..]),
+        ],
+        offset_generators,
+        &(),
+    );
+    assert_eq!(commitments, expected_commitments,);
+}
+
+#[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, &());
+    let expected_commitments: Vec<RistrettoPoint> = RistrettoPoint::compute_commitments(
+        &[
+            CommittableColumn::from(&mle1[..]),
+            CommittableColumn::from(&mle2[..]),
+        ],
+        offset_generators,
+        &(),
+    );
+    assert_eq!(commitments, expected_commitments,);
+}
+
+#[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);
+}

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;

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>;