fix: handle batch_size > 1 in prove_batch/verify_batch RLC combination (#205)

Author: shreyas-londhe

src/whir/mod.rs

@@ -516,4 +516,173 @@ mod tests {
             "Verifier should reject mismatched polynomial"
         );
     }
+
+    /// Test batch proving with batch_size > 1 (multiple polynomials per commitment).
+    ///
+    /// This tests the case where each commitment contains multiple stacked polynomials
+    /// (e.g., masked witness + random blinding for ZK), and we batch-prove multiple
+    /// such commitments together.
+    ///
+    /// This was a regression test for a bug where the RLC combination of stacked
+    /// leaf answers was incorrect when batch_size > 1.
+    fn make_whir_batch_with_batch_size(
+        num_variables: usize,
+        folding_factor: FoldingFactor,
+        num_points_per_poly: usize,
+        num_witnesses: usize,
+        batch_size: usize,
+        soundness_type: SoundnessType,
+        pow_bits: usize,
+    ) {
+        let num_coeffs = 1 << num_variables;
+        let mut rng = ark_std::test_rng();
+
+        let (leaf_hash_params, two_to_one_params) =
+            default_config::<EF, Blake3LeafHash<EF>, Blake3Compress>(&mut rng);
+
+        let mv_params = MultivariateParameters::new(num_variables);
+        let whir_params = ProtocolParameters::<MerkleConfig, PowStrategy> {
+            initial_statement: true,
+            security_level: 32,
+            pow_bits,
+            folding_factor,
+            leaf_hash_params,
+            two_to_one_params,
+            soundness_type,
+            _pow_parameters: Default::default(),
+            starting_log_inv_rate: 1,
+            batch_size, // KEY: batch_size > 1
+            deduplication_strategy: DeduplicationStrategy::Enabled,
+            merkle_proof_strategy: MerkleProofStrategy::Compressed,
+        };
+
+        let reed_solomon = Arc::new(RSDefault);
+        let basefield_reed_solomon = reed_solomon.clone();
+        let params = WhirConfig::new(reed_solomon, basefield_reed_solomon, mv_params, whir_params);
+
+        // Create polynomials for each witness
+        // Each witness will contain batch_size polynomials committed together
+        let mut all_polynomials: Vec<Vec<CoefficientList<F>>> = Vec::new();
+        for w in 0..num_witnesses {
+            let witness_polys: Vec<_> = (0..batch_size)
+                .map(|b| {
+                    // Different polynomials: witness 0 batch 0 = all 1s, witness 0 batch 1 = all 2s, etc.
+                    CoefficientList::new(vec![
+                        F::from(((w * batch_size + b) + 1) as u64);
+                        num_coeffs
+                    ])
+                })
+                .collect();
+            all_polynomials.push(witness_polys);
+        }
+
+        // Create statements - one per witness, based on the first polynomial in each batch
+        // (the internally-batched polynomial is what the statement is about)
+        let mut statements = Vec::new();
+        for witness_polys in &all_polynomials {
+            let poly = &witness_polys[0]; // Use first poly for statement (will be batched internally)
+            let mut statement = Statement::new(num_variables);
+
+            for _ in 0..num_points_per_poly {
+                let point = MultilinearPoint::rand(&mut rng, num_variables);
+                let eval = poly.evaluate_at_extension(&point);
+                statement.add_constraint(Weights::evaluation(point), eval);
+            }
+
+            // Add a linear constraint
+            let input = CoefficientList::new(
+                (0..1 << num_variables)
+                    .map(<EF as Field>::BasePrimeField::from)
+                    .collect(),
+            );
+            let linear_claim_weight = Weights::linear(input.into());
+            let poly_evals = EvaluationsList::from(poly.clone().to_extension());
+            let sum = linear_claim_weight.weighted_sum(&poly_evals);
+            statement.add_constraint(linear_claim_weight, sum);
+
+            statements.push(statement);
+        }
+
+        // Set up domain separator
+        let mut domainsep = DomainSeparator::new("🌪️");
+        for _ in 0..num_witnesses {
+            domainsep = domainsep.commit_statement(&params);
+        }
+
+        let total_constraints = num_witnesses * params.committment_ood_samples
+            + statements
+                .iter()
+                .map(|s| s.constraints.len())
+                .sum::<usize>();
+        domainsep = domainsep.add_whir_batch_proof(&params, num_witnesses, total_constraints);
+        let mut prover_state = domainsep.to_prover_state();
+
+        // Commit using commit_batch (stacks batch_size polynomials per witness)
+        let committer = CommitmentWriter::new(params.clone());
+        let mut witnesses = Vec::new();
+
+        for witness_polys in &all_polynomials {
+            let poly_refs: Vec<_> = witness_polys.iter().collect();
+            let witness = committer
+                .commit_batch(&mut prover_state, &poly_refs)
+                .unwrap();
+            witnesses.push(witness);
+        }
+
+        // Batch prove all witnesses together
+        let prover = Prover::new(params.clone());
+        let result = prover.prove_batch(&mut prover_state, &statements, &witnesses);
+        assert!(
+            result.is_ok(),
+            "Batch proving with batch_size={} failed: {:?}",
+            batch_size,
+            result.err()
+        );
+
+        // Verify
+        let commitment_reader = CommitmentReader::new(&params);
+        let verifier = Verifier::new(&params);
+        let mut verifier_state = domainsep.to_verifier_state(prover_state.narg_string());
+
+        let mut parsed_commitments = Vec::new();
+        for _ in 0..num_witnesses {
+            let parsed_commitment = commitment_reader
+                .parse_commitment(&mut verifier_state)
+                .unwrap();
+            parsed_commitments.push(parsed_commitment);
+        }
+
+        let verify_result =
+            verifier.verify_batch(&mut verifier_state, &parsed_commitments, &statements);
+        assert!(
+            verify_result.is_ok(),
+            "Batch verification with batch_size={} failed: {:?}",
+            batch_size,
+            verify_result.err()
+        );
+    }
+
+    #[test]
+    fn test_whir_batch_with_batch_size_2() {
+        // This is the key regression test for the batch_size > 1 bug
+        let batch_sizes = [2, 3];
+        let num_witnesses = [2, 3];
+        let folding_factors = [2, 3];
+
+        for batch_size in batch_sizes {
+            for num_witness in num_witnesses {
+                for folding_factor in folding_factors {
+                    make_whir_batch_with_batch_size(
+                        folding_factor * 2, // num_variables
+                        FoldingFactor::Constant(folding_factor),
+                        1, // num_points_per_poly
+                        num_witness,
+                        batch_size,
+                        SoundnessType::ConjectureList,
+                        0, // pow_bits
+                    );
+                }
+            }
+        }
+    }
 }

src/whir/prover.rs

@@ -415,13 +415,14 @@ where
         )?;
 
         let fold_size = 1 << self.config.folding_factor.at_round(0);
+        let leaf_size = fold_size * self.config.batch_size;
         let mut all_witness_answers = Vec::with_capacity(witnesses.len());
 
         // For each witness, provide Merkle proof opening at the challenge indices
         for witness in witnesses {
             let answers: Vec<_> = stir_indexes
                 .iter()
-                .map(|&i| witness.merkle_leaves[i * fold_size..(i + 1) * fold_size].to_vec())
+                .map(|&i| witness.merkle_leaves[i * leaf_size..(i + 1) * leaf_size].to_vec())
                 .collect();
             prover_state.hint::<Vec<Vec<F>>>(&answers)?;
             self.merkle_state.write_proof_hint(
@@ -438,10 +439,19 @@ where
             .map(|query_idx| {
                 let mut combined = vec![F::ZERO; fold_size];
                 let mut pow = F::ONE;
-                for witness_answers in &all_witness_answers {
-                    for (dst, src) in combined.iter_mut().zip(&witness_answers[query_idx]) {
-                        *dst += pow * src;
+                for (witness_idx, witness) in witnesses.iter().enumerate() {
+                    let answer = &all_witness_answers[witness_idx][query_idx];
+
+                    // First, internally reduce stacked leaf using witness's batching_randomness
+                    let mut internal_pow = F::ONE;
+                    for poly_idx in 0..self.config.batch_size {
+                        let start = poly_idx * fold_size;
+                        for j in 0..fold_size {
+                            combined[j] += pow * internal_pow * answer[start + j];
+                        }
+                        internal_pow *= witness.batching_randomness;
                     }
+
                     pow *= batching_randomness;
                 }
                 combined

src/whir/verifier.rs

@@ -236,7 +236,7 @@ where
         // Step 1: Read the N×M constraint evaluation matrix from the transcript
 
         // Collect all constraint weights to determine matrix dimensions
-        let mut all_constraint_weights = Vec::new();
+        let mut all_constraints_info: Vec<(Weights<F>, bool)> = Vec::new();
 
         // OOD constraints from each commitment
         for commitment in parsed_commitments {
@@ -245,20 +245,21 @@ where
                     *point,
                     self.params.mv_parameters.num_variables,
                 );
-                all_constraint_weights.push(Weights::evaluation(ml_point));
+                all_constraints_info.push((Weights::evaluation(ml_point), false));
             }
         }
 
         // Statement constraints
         for statement in statements {
             for constraint in &statement.constraints {
-                all_constraint_weights.push(constraint.weights.clone());
+                all_constraints_info
+                    .push((constraint.weights.clone(), constraint.defer_evaluation));
             }
         }
 
         // Read the N×M evaluation matrix from transcript
         let num_polynomials = parsed_commitments.len();
-        let num_constraints = all_constraint_weights.len();
+        let num_constraints = all_constraints_info.len();
         let mut constraint_evals_matrix = Vec::with_capacity(num_polynomials);
 
         for _ in 0..num_polynomials {
@@ -279,7 +280,9 @@ where
         if self.params.initial_statement {
             let mut all_constraints = Vec::new();
 
-            for (constraint_idx, weights) in all_constraint_weights.into_iter().enumerate() {
+            for (constraint_idx, (weights, defer_evaluation)) in
+                all_constraints_info.into_iter().enumerate()
+            {
                 let mut combined_eval = F::ZERO;
                 let mut pow = F::ONE;
                 for poly_evals in &constraint_evals_matrix {
@@ -290,7 +293,7 @@ where
                 all_constraints.push(Constraint {
                     weights,
                     sum: combined_eval,
-                    defer_evaluation: false,
+                    defer_evaluation,
                 });
             }
 
@@ -368,10 +371,19 @@ where
             .map(|query_idx| {
                 let mut combined = vec![F::ZERO; fold_size];
                 let mut pow = F::ONE;
-                for witness_answers in &all_answers {
-                    for (dst, src) in combined.iter_mut().zip(&witness_answers[query_idx]) {
-                        *dst += pow * src;
+                for (commitment, witness_answers) in parsed_commitments.iter().zip(&all_answers) {
+                    let stacked_answer = &witness_answers[query_idx];
+
+                    // First, internally reduce stacked leaf using commitment's batching_randomness
+                    let mut internal_pow = F::ONE;
+                    for poly_idx in 0..self.params.batch_size {
+                        let start = poly_idx * fold_size;
+                        for j in 0..fold_size {
+                            combined[j] += pow * internal_pow * stacked_answer[start + j];
+                        }
+                        internal_pow *= commitment.batching_randomness;
                     }
+
                     pow *= batching_randomness;
                 }
                 combined