feat: add membership check

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

@@ -846,10 +846,10 @@ impl ProofPlan for FirstRoundSquareTestProofPlan {
                 .unwrap();
         let first_round_res_eval = builder.try_consume_mle_evaluation()?;
         let final_round_res_eval = builder.try_consume_mle_evaluation()?;
-        //assert_eq!(first_round_res_eval, final_round_res_eval);
+        assert_eq!(first_round_res_eval, final_round_res_eval);
         builder.try_produce_sumcheck_subpolynomial_evaluation(
             SumcheckSubpolynomialType::Identity,
-            first_round_res_eval - x_eval * x_eval,
+            final_round_res_eval - x_eval * x_eval,
             2,
         )?;
         Ok(TableEvaluation::new(

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

@@ -0,0 +1,140 @@
+use super::{verify_constant_abs_decomposition, verify_constant_sign_decomposition};
+use crate::{
+    base::{
+        bit::{compute_varying_bit_matrix, BitDistribution},
+        proof::ProofError,
+        scalar::Scalar,
+    },
+    sql::proof::{
+        FinalRoundBuilder, FirstRoundBuilder, SumcheckSubpolynomialTerm, SumcheckSubpolynomialType,
+        VerificationBuilder,
+    },
+};
+use alloc::{boxed::Box, vec, vec::Vec};
+use bumpalo::Bump;
+
+/// Perform first round evaluation of the membership check.
+pub fn first_round_evaluate_membership_check<'a, S: Scalar>(
+    builder: &mut FirstRoundBuilder<'a, S>,
+    indexes: &[usize],
+    num_rows: usize,
+    alloc: &'a Bump,
+) {
+    let multiplicity_map = indexes.into_iter().counts();
+    let multiplicities = (0..num_rows - 1)
+        .map(|i| multiplicity_map.get(&i).copied().unwrap_or(0))
+        .collect::<Vec<_>>();
+    let alloc_multiplicities = alloc.alloc_slice_fill_copy(&multiplicities);
+    builder.produce_intermediate_mle(alloc_multiplicities as &[_]);
+    builder.request_post_result_challenges(2);
+}
+
+/// Perform final round evaluation of the membership check.
+pub fn final_round_evaluate_membership_check<'a, S: Scalar>(
+    builder: &mut FinalRoundBuilder<'a, S>,
+    columns: &[Column<'a, S>],
+    candidate_subset: &[Column<'a, S>],
+    indexes: &[usize],
+    num_rows: usize,
+    alloc: &'a Bump,
+) {
+    // 1. Get multiplicity of each index
+    let multiplicity_map = indexes.into_iter().counts();
+    let multiplicities = (0..num_rows - 1)
+        .map(|i| multiplicity_map.get(&i).copied().unwrap_or(0))
+        .collect::<Vec<_>>();
+    let alloc_multiplicities = alloc.alloc_slice_fill_copy(&multiplicities);
+    builder.produce_intermediate_mle(alloc_multiplicities as &[_]);
+    // 2. Fold the columns
+    let alpha = builder.consume_post_result_challenge();
+    let beta = builder.consume_post_result_challenge();
+    let ones = alloc.alloc_slice_fill_copy(num_rows, true);
+    let c_fold = alloc.alloc_slice_fill_copy(n, Zero::zero());
+    fold_columns(c_fold, alpha, beta, c);
+    let d_fold = alloc.alloc_slice_fill_copy(m, Zero::zero());
+    fold_columns(d_fold, alpha, beta, d);
+    let c_star = alloc.alloc_slice_copy(c_fold);
+    slice_ops::add_const::<S, S>(c_star, One::one());
+    slice_ops::batch_inversion(c_star);
+
+    let d_star = alloc.alloc_slice_copy(d_fold);
+    slice_ops::add_const::<S, S>(d_star, One::one());
+    slice_ops::batch_inversion(d_star);
+
+    builder.produce_intermediate_mle(c_star as &[_]);
+    builder.produce_intermediate_mle(d_star as &[_]);
+
+    // sum c_star * multiplicities - d_star = 0
+    builder.produce_sumcheck_subpolynomial(
+        SumcheckSubpolynomialType::ZeroSum,
+        vec![
+            (S::one(), vec![Box::new(c_star as &[_]), Box::new(s)]),
+            (-S::one(), vec![Box::new(d_star as &[_])]),
+        ],
+    );
+
+    // c_star + c_fold * c_star - input_ones = 0
+    builder.produce_sumcheck_subpolynomial(
+        SumcheckSubpolynomialType::Identity,
+        vec![
+            (S::one(), vec![Box::new(c_star as &[_])]),
+            (
+                S::one(),
+                vec![Box::new(c_star as &[_]), Box::new(c_fold as &[_])],
+            ),
+            (-S::one(), vec![Box::new(input_ones as &[_])]),
+        ],
+    );
+
+    // d_star + d_fold * d_star - chi = 0
+    builder.produce_sumcheck_subpolynomial(
+        SumcheckSubpolynomialType::Identity,
+        vec![
+            (S::one(), vec![Box::new(d_star as &[_])]),
+            (
+                S::one(),
+                vec![Box::new(d_star as &[_]), Box::new(d_fold as &[_])],
+            ),
+            (-S::one(), vec![Box::new(chi as &[_])]),
+        ],
+    );
+}
+
+pub(super) fn verify_filter<S: Scalar>(
+    builder: &mut VerificationBuilder<S>,
+    alpha: S,
+    beta: S,
+    one_eval: S,
+    chi_eval: S,
+    c_evals: &[S],
+    s_eval: S,
+    d_evals: &[S],
+) -> Result<(), ProofError> {
+    let c_fold_eval = alpha * fold_vals(beta, c_evals);
+    let d_fold_eval = alpha * fold_vals(beta, d_evals);
+    let c_star_eval = builder.try_consume_mle_evaluation()?;
+    let d_star_eval = builder.try_consume_mle_evaluation()?;
+
+    // sum c_star * s - d_star = 0
+    builder.try_produce_sumcheck_subpolynomial_evaluation(
+        SumcheckSubpolynomialType::ZeroSum,
+        c_star_eval * s_eval - d_star_eval,
+        2,
+    )?;
+
+    // c_star + c_fold * c_star - input_ones = 0
+    builder.try_produce_sumcheck_subpolynomial_evaluation(
+        SumcheckSubpolynomialType::Identity,
+        c_star_eval + c_fold_eval * c_star_eval - one_eval,
+        2,
+    )?;
+
+    // d_star + d_fold * d_star - chi = 0
+    builder.try_produce_sumcheck_subpolynomial_evaluation(
+        SumcheckSubpolynomialType::Identity,
+        d_star_eval + d_fold_eval * d_star_eval - chi_eval,
+        2,
+    )?;
+
+    Ok(())
+}

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

@@ -3,6 +3,7 @@ mod bitwise_verification;
 use bitwise_verification::{verify_constant_abs_decomposition, verify_constant_sign_decomposition};
 #[cfg(test)]
 mod bitwise_verification_test;
+mod membership_check;
 mod sign_expr;
 pub(crate) use sign_expr::{prover_evaluate_sign, result_evaluate_sign, verifier_evaluate_sign};
 pub mod range_check;