feat: allow FirstRoundBuilder to produce MLEs

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,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);
+}

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