fix: update first round builder to accept initial range length

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

@@ -8,19 +8,32 @@ pub struct FirstRoundBuilder {
     num_post_result_challenges: usize,
     /// The extra one evaluation lengths used in the proof.
     one_evaluation_lengths: Vec<usize>,
+
+    range_length: usize,
 }
 
 impl Default for FirstRoundBuilder {
     fn default() -> Self {
-        Self::new()
+        Self::new(0)
     }
 }
 
 impl FirstRoundBuilder {
-    pub fn new() -> Self {
+    pub fn new(initial_range_length: usize) -> Self {
         Self {
             num_post_result_challenges: 0,
             one_evaluation_lengths: Vec::new(),
+            range_length: initial_range_length,
+        }
+    }
+
+    pub(crate) fn range_length(&self) -> usize {
+        self.range_length
+    }
+
+    pub(crate) fn update_range_length(&mut self, new_range_length: usize) {
+        if new_range_length > self.range_length {
+            self.range_length = new_range_length;
         }
     }
 

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

@@ -100,18 +100,13 @@ 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::new(initial_range_length);
         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();
         let one_evaluation_lengths = first_round_builder.one_evaluation_lengths();
 
-        let range_length = one_evaluation_lengths
-            .iter()
-            .copied()
-            .chain(core::iter::once(initial_range_length))
-            .max()
-            .expect("Will always have at least one element"); // safe to unwrap because we have at least one element
+        let range_length = first_round_builder.range_length();
 
         let num_sumcheck_variables = cmp::max(log2_up(range_length), 1);
         assert!(num_sumcheck_variables > 0);

crates/proof-of-sql/src/sql/proof_gadgets/range_check.rs

@@ -22,13 +22,20 @@
 //! * Parallelization: Single-threaded execution of these operations is a performance bottleneck
 use crate::{
     base::{polynomial::MultilinearExtension, proof::ProofSizeMismatch, scalar::Scalar, slice_ops},
-    sql::proof::{FinalRoundBuilder, SumcheckSubpolynomialType, VerificationBuilder},
+    sql::proof::{
+        FinalRoundBuilder, FirstRoundBuilder, SumcheckSubpolynomialType, VerificationBuilder,
+    },
 };
 use alloc::{boxed::Box, vec::Vec};
 use bumpalo::Bump;
 use bytemuck::cast_slice;
 use core::{cmp::max, iter::repeat};
 
+/// Update the max range length for the range check.
+pub fn first_round_evaluate_range_check<'a, S: Scalar + 'a>(builder: &mut FirstRoundBuilder) {
+    builder.update_range_length(256);
+}
+
 /// Prove that a word-wise decomposition of a collection of scalars
 /// are all within the range 0 to 2^248.
 pub fn final_round_evaluate_range_check<'a, S: Scalar + 'a>(

crates/proof-of-sql/src/sql/proof_gadgets/range_check_test.rs

@@ -28,6 +28,7 @@ impl ProverEvaluate for RangeCheckTestPlan {
     ) -> Table<'a, S> {
         builder.request_post_result_challenges(1);
         builder.produce_one_evaluation_length(256);
+        builder.update_range_length(256);
         table_map[&self.column.table_ref()].clone()
     }
 

crates/proof-of-sql/src/sql/proof_plans/filter_exec_test.rs

@@ -192,7 +192,7 @@ fn we_can_get_an_empty_result_from_a_basic_filter_on_an_empty_table_using_first_
             ColumnType::Decimal75(Precision::new(75).unwrap(), 0),
         ),
     ];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,
@@ -241,7 +241,7 @@ fn we_can_get_an_empty_result_from_a_basic_filter_using_first_round_evaluate() {
             ColumnType::Decimal75(Precision::new(1).unwrap(), 0),
         ),
     ];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,
@@ -278,7 +278,7 @@ fn we_can_get_no_columns_from_a_basic_filter_with_no_selected_columns_using_firs
     let where_clause: DynProofExpr = equal(column(t, "a", &accessor), const_int128(5));
     let expr = filter(cols_expr_plan(t, &[], &accessor), tab(t), where_clause);
     let fields = &[];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,
@@ -321,7 +321,7 @@ fn we_can_get_the_correct_result_from_a_basic_filter_using_first_round_evaluate(
             ColumnType::Decimal75(Precision::new(1).unwrap(), 0),
         ),
     ];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,

crates/proof-of-sql/src/sql/proof_plans/projection_exec_test.rs

@@ -174,7 +174,7 @@ fn we_can_get_an_empty_result_from_a_basic_projection_on_an_empty_table_using_fi
             ColumnType::Decimal75(Precision::new(75).unwrap(), 0),
         ),
     ];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,
@@ -211,7 +211,7 @@ fn we_can_get_no_columns_from_a_basic_projection_with_no_selected_columns_using_
     accessor.add_table(t, data, 0);
     let expr: DynProofPlan = projection(cols_expr_plan(t, &[], &accessor), tab(t));
     let fields = &[];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,
@@ -260,7 +260,7 @@ fn we_can_get_the_correct_result_from_a_basic_projection_using_first_round_evalu
             ColumnType::Decimal75(Precision::new(1).unwrap(), 0),
         ),
     ];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,

crates/proof-of-sql/src/sql/proof_plans/slice_exec_test.rs

@@ -102,7 +102,7 @@ fn we_can_get_an_empty_result_from_a_slice_on_an_empty_table_using_first_round_e
             ColumnType::Decimal75(Precision::new(75).unwrap(), 0),
         ),
     ];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,
@@ -156,7 +156,7 @@ fn we_can_get_an_empty_result_from_a_slice_using_first_round_evaluate() {
             ColumnType::Decimal75(Precision::new(1).unwrap(), 0),
         ),
     ];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,
@@ -197,7 +197,7 @@ fn we_can_get_no_columns_from_a_slice_with_empty_input_using_first_round_evaluat
         None,
     );
     let fields = &[];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,
@@ -244,7 +244,7 @@ fn we_can_get_the_correct_result_from_a_slice_using_first_round_evaluate() {
             ColumnType::Decimal75(Precision::new(1).unwrap(), 0),
         ),
     ];
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(expr.first_round_evaluate(
         first_round_builder,
         &alloc,

crates/proof-of-sql/src/sql/proof_plans/union_exec.rs

@@ -120,6 +120,7 @@ impl ProverEvaluate for UnionExec {
         let res = table_union(&inputs, alloc, self.schema.clone()).expect("Failed to union tables");
         builder.request_post_result_challenges(2);
         builder.produce_one_evaluation_length(res.num_rows());
+        builder.update_range_length(res.num_rows());
         res
     }
 

crates/proof-of-sql/src/sql/proof_plans/union_exec_test.rs

@@ -272,7 +272,7 @@ fn we_can_get_result_from_union_using_first_round_evaluate() {
         ],
         fields.clone(),
     );
-    let first_round_builder = &mut FirstRoundBuilder::new();
+    let first_round_builder = &mut FirstRoundBuilder::new(0);
     let res: OwnedTable<Curve25519Scalar> = ProvableQueryResult::from(ast.first_round_evaluate(
         first_round_builder,
         &alloc,