fix: address review comments

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

@@ -7,24 +7,23 @@ use crate::{
         proof_plans::{fold_columns, fold_vals},
     },
 };
-use alloc::{boxed::Box, vec, vec::Vec};
+use alloc::{boxed::Box, vec};
 use bumpalo::Bump;
 use num_traits::{One, Zero};
 
 /// Perform first round evaluation of the membership check.
-#[allow(dead_code)]
 pub(crate) fn first_round_evaluate_membership_check<'a, S: Scalar>(
     builder: &mut FirstRoundBuilder<'a, S>,
-    multiplicities: &[usize],
-    alloc: &'a Bump,
+    multiplicities: &'a [i128],
 ) {
-    let casted_multiplicities: Vec<i128> = multiplicities.iter().map(|&x| x as i128).collect();
-    let alloc_multiplicities = alloc.alloc_slice_copy(casted_multiplicities.as_slice());
-    builder.produce_intermediate_mle(alloc_multiplicities as &[_]);
+    builder.produce_intermediate_mle(multiplicities);
     builder.request_post_result_challenges(2);
 }
 
 /// Perform final round evaluation of the membership check.
+///
+/// # Panics
+/// Panics if the number of source and candidate columns are not equal.
 #[allow(dead_code)]
 #[allow(clippy::too_many_arguments)]
 pub(crate) fn final_round_evaluate_membership_check<'a, S: Scalar>(
@@ -34,20 +33,19 @@ pub(crate) fn final_round_evaluate_membership_check<'a, S: Scalar>(
     beta: S,
     columns: &[Column<'a, S>],
     candidate_subset: &[Column<'a, S>],
-    multiplicities: &[usize],
-    num_rows: usize,
-    candidate_num_rows: usize,
+    multiplicities: &'a [i128],
+    input_ones: &'a [bool],
+    candidate_ones: &'a [bool],
 ) {
-    let casted_multiplicities: Vec<i128> = multiplicities.iter().map(|&x| x as i128).collect();
-    let alloc_multiplicities = alloc.alloc_slice_copy(casted_multiplicities.as_slice());
-
+    assert_eq!(
+        columns.len(),
+        candidate_subset.len(),
+        "The number of source and candidate columns should be equal"
+    );
     // Fold the columns
-    let input_ones = alloc.alloc_slice_fill_copy(num_rows, true);
-    let candidate_ones = alloc.alloc_slice_fill_copy(candidate_num_rows, true);
-
-    let c_fold = alloc.alloc_slice_fill_copy(num_rows, Zero::zero());
+    let c_fold = alloc.alloc_slice_fill_copy(input_ones.len(), Zero::zero());
     fold_columns(c_fold, alpha, beta, columns);
-    let d_fold = alloc.alloc_slice_fill_copy(candidate_num_rows, Zero::zero());
+    let d_fold = alloc.alloc_slice_fill_copy(candidate_ones.len(), Zero::zero());
     fold_columns(d_fold, alpha, beta, candidate_subset);
 
     let c_star = alloc.alloc_slice_copy(c_fold);
@@ -67,10 +65,7 @@ pub(crate) fn final_round_evaluate_membership_check<'a, S: Scalar>(
         vec![
             (
                 S::one(),
-                vec![
-                    Box::new(c_star as &[_]),
-                    Box::new(alloc_multiplicities as &[_]),
-                ],
+                vec![Box::new(c_star as &[_]), Box::new(multiplicities as &[_])],
             ),
             (-S::one(), vec![Box::new(d_star as &[_])]),
         ],
@@ -113,6 +108,12 @@ pub(crate) fn verify_membership_check<S: Scalar>(
     column_evals: &[S],
     candidate_evals: &[S],
 ) -> Result<(), ProofError> {
+    // Check that the source and candidate columns have the same amount of columns
+    if column_evals.len() != candidate_evals.len() {
+        return Err(ProofError::VerificationError {
+            error: "The number of source and candidate columns should be equal",
+        });
+    }
     let multiplicity_eval = builder.try_consume_first_round_mle_evaluation()?;
     let c_fold_eval = fold_vals(beta, column_evals);
     let d_fold_eval = fold_vals(beta, candidate_evals);

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

@@ -29,7 +29,7 @@ pub struct MembershipCheckTestPlan {
     pub candidate_table: TableRef,
     pub source_columns: Vec<ColumnRef>,
     pub candidate_columns: Vec<ColumnRef>,
-    pub proposed_multiplicities: Vec<usize>,
+    pub proposed_multiplicities: Vec<i128>,
 }
 
 impl ProverEvaluate for MembershipCheckTestPlan {
@@ -40,12 +40,6 @@ impl ProverEvaluate for MembershipCheckTestPlan {
         alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
-        // Check that the source and candidate columns have the same amount of columns
-        assert_eq!(
-            self.source_columns.len(),
-            self.candidate_columns.len(),
-            "Source and candidate columns must have the same length"
-        );
         // Check that the source columns belong to the source table
         for col_ref in &self.source_columns {
             assert_eq!(self.source_table, col_ref.table_ref(), "Table not found");
@@ -64,7 +58,8 @@ impl ProverEvaluate for MembershipCheckTestPlan {
         builder.produce_one_evaluation_length(source_table.num_rows());
         builder.produce_one_evaluation_length(candidate_table.num_rows());
         // Evaluate the first round
-        first_round_evaluate_membership_check(builder, &self.proposed_multiplicities, alloc);
+        let alloc_proposed_multiplicities = alloc.alloc_slice_copy(&self.proposed_multiplicities);
+        first_round_evaluate_membership_check(builder, alloc_proposed_multiplicities);
         // This is just a dummy table, the actual data is not used
         Table::try_new_with_options(IndexMap::default(), TableOptions { row_count: Some(0) })
             .unwrap()
@@ -76,12 +71,6 @@ impl ProverEvaluate for MembershipCheckTestPlan {
         alloc: &'a Bump,
         table_map: &IndexMap<TableRef, Table<'a, S>>,
     ) -> Table<'a, S> {
-        // Check that the source and candidate columns have the same amount of columns
-        assert_eq!(
-            self.source_columns.len(),
-            self.candidate_columns.len(),
-            "Source and candidate columns must have the same length"
-        );
         // Check that the source columns belong to the source table
         for col_ref in &self.source_columns {
             assert_eq!(self.source_table, col_ref.table_ref(), "Table not found");
@@ -122,6 +111,9 @@ impl ProverEvaluate for MembershipCheckTestPlan {
             .collect_in::<BumpVec<_>>(alloc);
         let alpha = builder.consume_post_result_challenge();
         let beta = builder.consume_post_result_challenge();
+        let source_ones = alloc.alloc_slice_fill_copy(source_table.num_rows(), true);
+        let candidate_ones = alloc.alloc_slice_fill_copy(candidate_table.num_rows(), true);
+        let alloc_proposed_multiplicities = alloc.alloc_slice_copy(&self.proposed_multiplicities);
         // Perform final membership check
         final_round_evaluate_membership_check(
             builder,
@@ -130,9 +122,9 @@ impl ProverEvaluate for MembershipCheckTestPlan {
             beta,
             &source_columns,
             &candidate_columns,
-            &self.proposed_multiplicities,
-            source_table.num_rows(),
-            candidate_table.num_rows(),
+            alloc_proposed_multiplicities,
+            source_ones,
+            candidate_ones,
         );
         // Return a dummy table
         Table::try_new_with_options(IndexMap::default(), TableOptions { row_count: Some(0) })
@@ -170,12 +162,6 @@ impl ProofPlan for MembershipCheckTestPlan {
         let alpha = builder.try_consume_post_result_challenge()?;
         let beta = builder.try_consume_post_result_challenge()?;
         let num_columns = self.source_columns.len();
-        // Check that the source and candidate columns have the same amount of columns
-        if num_columns != self.candidate_columns.len() {
-            return Err(ProofError::VerificationError {
-                error: "Source and candidate columns must have the same length",
-            });
-        }
         // Get the columns
         let column_evals = builder.try_consume_final_round_mle_evaluations(num_columns)?;
         // Get the target columns

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

@@ -4,7 +4,7 @@ use bitwise_verification::{verify_constant_abs_decomposition, verify_constant_si
 #[cfg(test)]
 mod bitwise_verification_test;
 mod membership_check;
-#[allow(unused_imports)]
+#[allow(unused_imports, dead_code)]
 use membership_check::{
     final_round_evaluate_membership_check, first_round_evaluate_membership_check,
     verify_membership_check,