feat: allow FirstRoundBuilder to produce MLEs

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

@@ -12,6 +12,8 @@ pub struct FinalRoundBuilder<'a, S: Scalar> {
     num_sumcheck_variables: usize,
     bit_distributions: Vec<BitDistribution>,
     commitment_descriptor: Vec<CommittableColumn<'a>>,
+    first_round_commitment_descriptor: Vec<CommittableColumn<'a>>,
+    consumed_first_round_commitment_descriptors: usize,
     pcs_proof_mles: Vec<Box<dyn MultilinearExtension<S> + 'a>>,
     sumcheck_subpolynomials: Vec<SumcheckSubpolynomial<'a, S>>,
     /// The challenges used in creation of the constraints in the proof.
@@ -25,11 +27,17 @@ pub struct FinalRoundBuilder<'a, S: Scalar> {
 }
 
 impl<'a, S: Scalar> FinalRoundBuilder<'a, S> {
-    pub fn new(num_sumcheck_variables: usize, post_result_challenges: Vec<S>) -> Self {
+    pub fn new(
+        num_sumcheck_variables: usize,
+        post_result_challenges: Vec<S>,
+        first_round_commitment_descriptor: Vec<CommittableColumn<'a>>,
+    ) -> Self {
         Self {
             num_sumcheck_variables,
             bit_distributions: Vec::new(),
             commitment_descriptor: Vec::new(),
+            first_round_commitment_descriptor,
+            consumed_first_round_commitment_descriptors: 0,
             pcs_proof_mles: Vec::new(),
             sumcheck_subpolynomials: Vec::new(),
             post_result_challenges,
@@ -44,6 +52,12 @@ impl<'a, S: Scalar> FinalRoundBuilder<'a, S> {
         self.sumcheck_subpolynomials.len()
     }
 
+    pub fn consume_first_round_mle(&mut self) -> CommittableColumn<'a> {
+        let index = self.consumed_first_round_commitment_descriptors;
+        self.consumed_first_round_commitment_descriptors += 1;
+        self.first_round_commitment_descriptor[index].clone()
+    }
+
     pub fn pcs_proof_mles(&self) -> &[Box<dyn MultilinearExtension<S> + 'a>] {
         &self.pcs_proof_mles
     }

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

@@ -17,7 +17,7 @@ use curve25519_dalek::RistrettoPoint;
 fn we_can_compute_commitments_for_intermediate_mles_using_a_zero_offset() {
     let mle1 = [1, 2];
     let mle2 = [10i64, 20];
-    let mut builder = FinalRoundBuilder::<Curve25519Scalar>::new(1, Vec::new());
+    let mut builder = FinalRoundBuilder::<Curve25519Scalar>::new(1, Vec::new(), Vec::new());
     builder.produce_anchored_mle(&mle1);
     builder.produce_intermediate_mle(&mle2[..]);
     let offset_generators = 0_usize;
@@ -36,7 +36,7 @@ fn we_can_compute_commitments_for_intermediate_mles_using_a_zero_offset() {
 fn we_can_compute_commitments_for_intermediate_mles_using_a_non_zero_offset() {
     let mle1 = [1, 2];
     let mle2 = [10i64, 20];
-    let mut builder = FinalRoundBuilder::<Curve25519Scalar>::new(1, Vec::new());
+    let mut builder = FinalRoundBuilder::<Curve25519Scalar>::new(1, Vec::new(), Vec::new());
     builder.produce_anchored_mle(&mle1);
     builder.produce_intermediate_mle(&mle2[..]);
     let offset_generators = 123_usize;
@@ -55,7 +55,7 @@ fn we_can_compute_commitments_for_intermediate_mles_using_a_non_zero_offset() {
 fn we_can_evaluate_pcs_proof_mles() {
     let mle1 = [1, 2];
     let mle2 = [10i64, 20];
-    let mut builder = FinalRoundBuilder::new(1, Vec::new());
+    let mut builder = FinalRoundBuilder::new(1, Vec::new(), Vec::new());
     builder.produce_anchored_mle(&mle1);
     builder.produce_intermediate_mle(&mle2[..]);
     let evaluation_vec = [
@@ -112,6 +112,7 @@ fn we_can_consume_post_result_challenges_in_proof_builder() {
             Curve25519Scalar::from(456),
             Curve25519Scalar::from(789),
         ],
+        Vec::new(),
     );
     assert_eq!(
         Curve25519Scalar::from(789),
@@ -126,3 +127,23 @@ fn we_can_consume_post_result_challenges_in_proof_builder() {
         builder.consume_post_result_challenge()
     );
 }
+
+#[test]
+fn we_can_consume_first_round_committable_columns() {
+    let mut builder: FinalRoundBuilder<Curve25519Scalar> = FinalRoundBuilder::new(
+        2,
+        Vec::new(),
+        vec![
+            CommittableColumn::TinyInt(&[2_i8, 3, 4]),
+            CommittableColumn::BigInt(&[5_i64, 6, 7]),
+        ],
+    );
+    assert_eq!(
+        CommittableColumn::TinyInt(&[2_i8, 3, 4]),
+        builder.consume_first_round_mle()
+    );
+    assert_eq!(
+        CommittableColumn::BigInt(&[5_i64, 6, 7]),
+        builder.consume_first_round_mle()
+    );
+}

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

@@ -1,6 +1,13 @@
-use alloc::vec::Vec;
+use crate::base::{
+    commitment::{Commitment, CommittableColumn, VecCommitmentExt},
+    polynomial::MultilinearExtension,
+    scalar::Scalar,
+};
+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 +17,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 commitment_descriptor(&self) -> &[CommittableColumn<'a>] {
+        &self.commitment_descriptor
+    }
+
     /// Get the one evaluation lengths used in the proof.
     pub(crate) fn one_evaluation_lengths(&self) -> &[usize] {
         &self.one_evaluation_lengths
@@ -34,6 +47,58 @@ impl FirstRoundBuilder {
         self.one_evaluation_lengths.push(length);
     }
 
+    /// Produce an anchored MLE that we can reference in sumcheck.
+    ///
+    /// An anchored MLE is an MLE where the verifier has access to the commitment.
+    pub fn produce_anchored_mle(&mut self, data: impl MultilinearExtension<S> + 'a) {
+        self.pcs_proof_mles.push(Box::new(data));
+    }
+
+    /// 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.produce_anchored_mle(data);
+    }
+
+    /// Compute commitments of all the interemdiate MLEs used in sumcheck
+    #[tracing::instrument(
+        name = "FinalRoundBuilder::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 = "FinalRoundBuilder::evaluate_pcs_proof_mles",
+        level = "debug",
+        skip_all
+    )]
+    pub fn evaluate_pcs_proof_mles(&self, evaluation_vec: &[S]) -> Vec<S> {
+        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));
+        }
+        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_anchored_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_anchored_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_anchored_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

@@ -68,6 +68,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

@@ -39,7 +39,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>;