refactor!: generalize one eval to arbitrary lengths

crates/proof-of-sql/src/base/database/mod.rs

@@ -77,6 +77,9 @@ pub use table::{Table, TableOptions};
 mod table_test;
 pub mod table_utility;
 
+mod table_evaluation;
+pub use table_evaluation::TableEvaluation;
+
 /// TODO: add docs
 pub(crate) mod expression_evaluation;
 mod expression_evaluation_error;

crates/proof-of-sql/src/base/database/table_evaluation.rs

@@ -0,0 +1,34 @@
+use crate::base::scalar::Scalar;
+use alloc::vec::Vec;
+
+/// The result of evaluating a table
+#[derive(Debug, Eq, PartialEq, Clone)]
+pub struct TableEvaluation<S: Scalar> {
+    /// Evaluation of each column in the table
+    column_evals: Vec<S>,
+    /// Evaluation of an all-one column with the same length as the table
+    one_eval: S,
+}
+
+impl<S: Scalar> TableEvaluation<S> {
+    /// Creates a new [`TableEvaluation`].
+    #[must_use]
+    pub fn new(column_evals: Vec<S>, one_eval: S) -> Self {
+        Self {
+            column_evals,
+            one_eval,
+        }
+    }
+
+    /// Returns the evaluation of each column in the table.
+    #[must_use]
+    pub fn column_evals(&self) -> &[S] {
+        &self.column_evals
+    }
+
+    /// Returns the evaluation of an all-one column with the same length as the table.
+    #[must_use]
+    pub fn one_eval(&self) -> &S {
+        &self.one_eval
+    }
+}

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

@@ -17,6 +17,7 @@ pub struct FinalRoundBuilder<'a, S: Scalar> {
     commitment_descriptor: Vec<CommittableColumn<'a>>,
     pcs_proof_mles: Vec<Box<dyn MultilinearExtension<S> + 'a>>,
     sumcheck_subpolynomials: Vec<SumcheckSubpolynomial<'a, S>>,
+    one_evaluation_lengths: Vec<usize>,
     /// The challenges used in creation of the constraints 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
@@ -36,6 +37,7 @@ impl<'a, S: Scalar> FinalRoundBuilder<'a, S> {
             pcs_proof_mles: Vec::new(),
             sumcheck_subpolynomials: Vec::new(),
             post_result_challenges,
+            one_evaluation_lengths: Vec::new(),
         }
     }
 
@@ -51,6 +53,15 @@ impl<'a, S: Scalar> FinalRoundBuilder<'a, S> {
         &self.pcs_proof_mles
     }
 
+    pub fn one_evaluation_lengths(&self) -> &[usize] {
+        &self.one_evaluation_lengths
+    }
+
+    /// Whenever we need to evaluate a column of 1s with a given length, we push the length
+    pub fn push_one_evaluation_length(&mut self, length: usize) {
+        self.one_evaluation_lengths.push(length);
+    }
+
     /// Produce a bit distribution that describes which bits are constant
     /// and which bits varying in a column of data
     pub fn produce_bit_distribution(&mut self, dist: BitDistribution) {

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

@@ -1,6 +1,6 @@
 use super::{CountBuilder, FinalRoundBuilder, FirstRoundBuilder, VerificationBuilder};
 use crate::base::{
-    database::{ColumnField, ColumnRef, OwnedTable, Table, TableRef},
+    database::{ColumnField, ColumnRef, OwnedTable, Table, TableEvaluation, TableRef},
     map::{IndexMap, IndexSet},
     proof::ProofError,
     scalar::Scalar,
@@ -21,7 +21,8 @@ pub trait ProofPlan: Debug + Send + Sync + ProverEvaluate {
         builder: &mut VerificationBuilder<S>,
         accessor: &IndexMap<ColumnRef, S>,
         result: Option<&OwnedTable<S>>,
-    ) -> Result<Vec<S>, ProofError>;
+        one_eval_map: &IndexMap<TableRef, S>,
+    ) -> Result<TableEvaluation<S>, ProofError>;
 
     /// Return all the result column fields
     fn get_column_result_fields(&self) -> Vec<ColumnField>;

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

@@ -27,15 +27,15 @@ use serde::{Deserialize, Serialize};
 ///
 /// Basically we are looking for the smallest offset and the largest offset + length
 /// so that we have an index range of the table rows that the query is referencing.
-fn get_index_range(
+fn get_index_range<'a>(
     accessor: &dyn MetadataAccessor,
-    table_refs: impl IntoIterator<Item = TableRef>,
+    table_refs: impl IntoIterator<Item = &'a TableRef>,
 ) -> (usize, usize) {
     table_refs
         .into_iter()
         .map(|table_ref| {
-            let length = accessor.get_length(table_ref);
-            let offset = accessor.get_offset(table_ref);
+            let length = accessor.get_length(*table_ref);
+            let offset = accessor.get_offset(*table_ref);
             (offset, offset + length)
         })
         .reduce(|(min_start, max_end), (start, end)| (min_start.min(start), max_end.max(end)))
@@ -52,6 +52,8 @@ fn get_index_range(
 pub struct QueryProof<CP: CommitmentEvaluationProof> {
     /// Bit distributions
     pub bit_distributions: Vec<BitDistribution>,
+    /// One evaluation lengths
+    pub one_evaluation_lengths: Vec<usize>,
     /// Commitments
     pub commitments: Vec<CP::Commitment>,
     /// Sumcheck Proof
@@ -72,7 +74,7 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
         accessor: &impl DataAccessor<CP::Scalar>,
         setup: &CP::ProverPublicSetup<'_>,
     ) -> (Self, ProvableQueryResult) {
-        let (min_row_num, max_row_num) = get_index_range(accessor, expr.get_table_references());
+        let (min_row_num, max_row_num) = get_index_range(accessor, &expr.get_table_references());
         let initial_range_length = max_row_num - min_row_num;
         let alloc = Bump::new();
 
@@ -178,6 +180,7 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
 
         let proof = Self {
             bit_distributions: builder.bit_distributions().to_vec(),
+            one_evaluation_lengths: builder.one_evaluation_lengths().to_vec(),
             commitments,
             sumcheck_proof,
             pcs_proof_evaluations,
@@ -197,7 +200,8 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
         setup: &CP::VerifierPublicSetup<'_>,
     ) -> QueryResult<CP::Scalar> {
         let owned_table_result = result.to_owned_table(&expr.get_column_result_fields())?;
-        let (min_row_num, _) = get_index_range(accessor, expr.get_table_references());
+        let table_refs = expr.get_table_references();
+        let (min_row_num, _) = get_index_range(accessor, &table_refs);
         let num_sumcheck_variables = cmp::max(log2_up(self.range_length), 1);
         assert!(num_sumcheck_variables > 0);
 
@@ -274,21 +278,38 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
                 .take(self.pcs_proof_evaluations.len())
                 .collect();
 
+        // Always prepend input lengths to the one evaluation lengths
+        let table_length_map = table_refs
+            .iter()
+            .map(|table_ref| (table_ref, accessor.get_length(*table_ref)))
+            .collect::<IndexMap<_, _>>();
+
+        let one_evaluation_lengths = table_length_map
+            .values()
+            .chain(self.one_evaluation_lengths.clone().iter())
+            .copied()
+            .collect::<Vec<_>>();
+
         // pass over the provable AST to fill in the verification builder
         let sumcheck_evaluations = SumcheckMleEvaluations::new(
             self.range_length,
-            owned_table_result.num_rows(),
+            one_evaluation_lengths,
             &subclaim.evaluation_point,
             &sumcheck_random_scalars,
             &self.pcs_proof_evaluations,
         );
+        let one_eval_map: IndexMap<TableRef, CP::Scalar> = table_length_map
+            .iter()
+            .map(|(table_ref, length)| (**table_ref, sumcheck_evaluations.one_evaluations[length]))
+            .collect();
         let mut builder = VerificationBuilder::new(
             min_row_num,
             sumcheck_evaluations,
             &self.bit_distributions,
             sumcheck_random_scalars.subpolynomial_multipliers,
             &evaluation_random_scalars,
             post_result_challenges,
+            self.one_evaluation_lengths.clone(),
         );
 
         let pcs_proof_commitments: Vec<_> = column_references
@@ -305,11 +326,12 @@ impl<CP: CommitmentEvaluationProof> QueryProof<CP> {
             &mut builder,
             &evaluation_accessor,
             Some(&owned_table_result),
+            &one_eval_map,
         )?;
         // compute the evaluation of the result MLEs
         let result_evaluations = owned_table_result.mle_evaluations(&subclaim.evaluation_point);
         // check the evaluation of the result MLEs
-        if verifier_evaluations != result_evaluations {
+        if verifier_evaluations.column_evals() != result_evaluations {
             Err(ProofError::VerificationError {
                 error: "result evaluation check failed",
             })?;

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

@@ -8,7 +8,7 @@ use crate::{
             owned_table_utility::{bigint, owned_table},
             table_utility::*,
             ColumnField, ColumnRef, ColumnType, OwnedTable, OwnedTableTestAccessor, Table,
-            TableRef,
+            TableEvaluation, TableRef,
         },
         map::{indexset, IndexMap, IndexSet},
         proof::ProofError,
@@ -65,6 +65,7 @@ impl ProverEvaluate for TrivialTestProofPlan {
             SumcheckSubpolynomialType::Identity,
             vec![(S::ONE, vec![Box::new(col as &[_])])],
         );
+        builder.push_one_evaluation_length(self.length);
         table([borrowed_bigint(
             "a1",
             vec![self.column_fill_value; self.length],
@@ -85,13 +86,17 @@ impl ProofPlan for TrivialTestProofPlan {
         builder: &mut VerificationBuilder<S>,
         _accessor: &IndexMap<ColumnRef, S>,
         _result: Option<&OwnedTable<S>>,
-    ) -> Result<Vec<S>, ProofError> {
+        _one_eval_map: &IndexMap<TableRef, S>,
+    ) -> Result<TableEvaluation<S>, ProofError> {
         assert_eq!(builder.consume_intermediate_mle(), S::ZERO);
         builder.produce_sumcheck_subpolynomial_evaluation(
             &SumcheckSubpolynomialType::ZeroSum,
             S::from(self.evaluation),
         );
-        Ok(vec![S::ZERO])
+        Ok(TableEvaluation::new(
+            vec![S::ZERO],
+            builder.consume_one_evaluation(),
+        ))
     }
     ///
     /// # Panics
@@ -246,6 +251,7 @@ impl ProverEvaluate for SquareTestProofPlan {
                 (-S::ONE, vec![Box::new(x), Box::new(x)]),
             ],
         );
+        builder.push_one_evaluation_length(2);
         table([borrowed_bigint("a1", self.res, alloc)])
     }
 }
@@ -261,7 +267,8 @@ impl ProofPlan for SquareTestProofPlan {
         builder: &mut VerificationBuilder<S>,
         accessor: &IndexMap<ColumnRef, S>,
         _result: Option<&OwnedTable<S>>,
-    ) -> Result<Vec<S>, ProofError> {
+        _one_eval_map: &IndexMap<TableRef, S>,
+    ) -> Result<TableEvaluation<S>, ProofError> {
         let x_eval = S::from(self.anchored_commit_multiplier)
             * *accessor
                 .get(&ColumnRef::new(
@@ -275,7 +282,10 @@ impl ProofPlan for SquareTestProofPlan {
             &SumcheckSubpolynomialType::Identity,
             res_eval - x_eval * x_eval,
         );
-        Ok(vec![res_eval])
+        Ok(TableEvaluation::new(
+            vec![res_eval],
+            builder.consume_one_evaluation(),
+        ))
     }
     fn get_column_result_fields(&self) -> Vec<ColumnField> {
         vec![ColumnField::new("a1".parse().unwrap(), ColumnType::BigInt)]
@@ -436,6 +446,7 @@ impl ProverEvaluate for DoubleSquareTestProofPlan {
             ],
         );
         builder.produce_intermediate_mle(res);
+        builder.push_one_evaluation_length(2);
         table([borrowed_bigint("a1", self.res, alloc)])
     }
 }
@@ -451,7 +462,8 @@ impl ProofPlan for DoubleSquareTestProofPlan {
         builder: &mut VerificationBuilder<S>,
         accessor: &IndexMap<ColumnRef, S>,
         _result: Option<&OwnedTable<S>>,
-    ) -> Result<Vec<S>, ProofError> {
+        _one_eval_map: &IndexMap<TableRef, S>,
+    ) -> Result<TableEvaluation<S>, ProofError> {
         let x_eval = *accessor
             .get(&ColumnRef::new(
                 "sxt.test".parse().unwrap(),
@@ -473,7 +485,10 @@ impl ProofPlan for DoubleSquareTestProofPlan {
             &SumcheckSubpolynomialType::Identity,
             res_eval - z_eval * z_eval,
         );
-        Ok(vec![res_eval])
+        Ok(TableEvaluation::new(
+            vec![res_eval],
+            builder.consume_one_evaluation(),
+        ))
     }
     fn get_column_result_fields(&self) -> Vec<ColumnField> {
         vec![ColumnField::new("a1".parse().unwrap(), ColumnType::BigInt)]
@@ -634,6 +649,7 @@ impl ProverEvaluate for ChallengeTestProofPlan {
                 (-alpha, vec![Box::new(x), Box::new(x)]),
             ],
         );
+        builder.push_one_evaluation_length(2);
         table([borrowed_bigint("a1", [9, 25], alloc)])
     }
 }
@@ -650,7 +666,8 @@ impl ProofPlan for ChallengeTestProofPlan {
         builder: &mut VerificationBuilder<S>,
         accessor: &IndexMap<ColumnRef, S>,
         _result: Option<&OwnedTable<S>>,
-    ) -> Result<Vec<S>, ProofError> {
+        _one_eval_map: &IndexMap<TableRef, S>,
+    ) -> Result<TableEvaluation<S>, ProofError> {
         let alpha = builder.consume_post_result_challenge();
         let _beta = builder.consume_post_result_challenge();
         let x_eval = *accessor
@@ -665,7 +682,10 @@ impl ProofPlan for ChallengeTestProofPlan {
             &SumcheckSubpolynomialType::Identity,
             alpha * res_eval - alpha * x_eval * x_eval,
         );
-        Ok(vec![res_eval])
+        Ok(TableEvaluation::new(
+            vec![res_eval],
+            builder.consume_one_evaluation(),
+        ))
     }
     fn get_column_result_fields(&self) -> Vec<ColumnField> {
         vec![ColumnField::new("a1".parse().unwrap(), ColumnType::BigInt)]