Kunming/simple batch fix (#14)

Fixed the bugs in the simple batching case and add tests. Use `cargo
test --package whir --lib -- whir::tests::test_whir --exact
--show-output` to run the new batching test.

src/whir/batch/committer.rs

@@ -160,7 +160,7 @@ where
                         );
                         ood_answers[j] = eval;
                     }
-                });
+                });   
             merlin.add_scalars(&ood_answers)?;
         }
 

src/whir/batch/prover.rs

@@ -229,8 +229,9 @@ where
             let fold_size = 1 << self.0.folding_factor;
             let answers = final_challenge_indexes
                 .into_par_iter()
-                .map(|i| prev_merkle_answers[i * fold_size..(i + 1) * fold_size].to_vec())
+                .map(|i| prev_merkle_answers[i * (fold_size * num_polys)..(i + 1) * (fold_size * num_polys)].to_vec())
                 .collect();
+
             round_state.merkle_proofs.push((merkle_proof, answers));
 
             // PoW

src/whir/batch/verifier.rs

@@ -53,7 +53,6 @@ where
             |evals: &[F], coeff: &[F]| -> F { zip_eq(evals, coeff).map(|(a, b)| *a * *b).sum() };
 
         let random_coeff = super::utils::generate_random_vector_batch_verify(arthur, num_polys)?;
-
         let initial_claims: Vec<_> = parsed_commitment
             .ood_points
             .clone()
@@ -89,7 +88,7 @@ where
             &parsed_commitment,
             &statement,
             whir_proof,
-            random_coeff,
+            random_coeff.clone(),
             num_polys,
         )?;
 
@@ -265,6 +264,7 @@ where
         let mut sumcheck_rounds = Vec::new();
         let mut folding_randomness: MultilinearPoint<F>;
         let initial_combination_randomness;
+
         if self.params.initial_statement {
             // Derive combination randomness and first sumcheck polynomial
             let [combination_randomness_gen]: [F; 1] = arthur.challenge_scalars()?;
@@ -447,6 +447,28 @@ where
             return Err(ProofError::InvalidProof);
         }
 
+        let final_randomness_answers: Vec<_> = if self.params.n_rounds() == 0 {
+            final_randomness_answers
+                .into_iter()
+                .map(|raw_answer| {
+                    if batched_randomness.len() > 0 {
+                        let chunk_size = 1 << self.params.folding_factor;
+                        let mut res = vec![F::ZERO; chunk_size];
+                        for i in 0..chunk_size {
+                            for j in 0..num_polys {
+                                res[i] += raw_answer[i + j * chunk_size] * batched_randomness[j];
+                            }
+                        }
+                        res
+                    } else {
+                        raw_answer.clone()
+                    }
+                })
+                .collect()
+        } else {
+            final_randomness_answers.to_vec()
+        };
+
         if self.params.final_pow_bits > 0. {
             arthur.challenge_pow::<PowStrategy>(self.params.final_pow_bits)?;
         }

src/whir/mod.rs

@@ -45,6 +45,7 @@ mod tests {
     use crate::parameters::{FoldType, MultivariateParameters, SoundnessType, WhirParameters};
     use crate::poly_utils::coeffs::CoefficientList;
     use crate::poly_utils::MultilinearPoint;
+    use crate::whir::batch::WhirBatchIOPattern;
     use crate::whir::Statement;
     use crate::whir::{
         committer::Committer, iopattern::WhirIOPattern, parameters::WhirConfig, prover::Prover,
@@ -120,6 +121,74 @@ mod tests {
         assert!(verifier.verify(&mut arthur, &statement, &proof).is_ok());
     }
 
+    fn make_whir_batch_things(
+        num_polynomials: usize,
+        num_variables: usize,
+        folding_factor: usize,
+        soundness_type: SoundnessType,
+        pow_bits: usize,
+        fold_type: FoldType,
+    ) {
+        println!(
+            "NP = {num_polynomials}, NV = {num_variables}, FOLD_TYPE = {:?}",
+            fold_type
+        );
+        let num_coeffs = 1 << num_variables;
+
+        let mut rng = ark_std::test_rng();
+        let (leaf_hash_params, two_to_one_params) = merkle_tree::default_config::<F>(&mut rng);
+
+        let mv_params = MultivariateParameters::<F>::new(num_variables);
+
+        let whir_params = WhirParameters::<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,
+            fold_optimisation: fold_type,
+        };
+
+        let params = WhirConfig::<F, MerkleConfig, PowStrategy>::new(mv_params, whir_params);
+
+        let polynomials: Vec<CoefficientList<F>> = (0..num_polynomials)
+            .map(|i| CoefficientList::new(vec![F::from((i + 1) as i32); num_coeffs]))
+            .collect();
+
+        let point = MultilinearPoint::rand(&mut rng, num_variables);
+        let evals: Vec<F> = polynomials
+            .iter()
+            .map(|poly| poly.evaluate(&point))
+            .collect();
+        let point = point.0;
+
+        let io = IOPattern::<DefaultHash>::new("🌪️")
+            .commit_batch_statement(&params, num_polynomials)
+            .add_whir_batch_proof(&params, num_polynomials)
+            .clone();
+        let mut merlin = io.to_merlin();
+
+        let committer = Committer::new(params.clone());
+        let witnesses = committer.batch_commit(&mut merlin, &polynomials).unwrap();
+
+        let prover = Prover(params.clone());
+
+        let proof = prover
+            .simple_batch_prove(&mut merlin, &point, &evals, &witnesses)
+            .unwrap();
+
+        let verifier = Verifier::new(params);
+        let mut arthur = io.to_arthur(merlin.transcript());
+        assert!(verifier
+            .simple_batch_verify(&mut arthur, &point, &evals, &proof)
+            .is_ok());
+        println!("PASSED!");
+    }
+
     #[test]
     fn test_whir() {
         let folding_factors = [2, 3, 4, 5];
@@ -130,10 +199,11 @@ mod tests {
         ];
         let fold_types = [FoldType::Naive, FoldType::ProverHelps];
         let num_points = [0, 1, 2];
+        let num_polys = [1, 2, 3];
         let pow_bits = [0, 5, 10];
 
         for folding_factor in folding_factors {
-            let num_variables = folding_factor - 1..= 2 * folding_factor;
+            let num_variables = folding_factor - 1..=2 * folding_factor;
             for num_variables in num_variables {
                 for fold_type in fold_types {
                     for num_points in num_points {
@@ -153,5 +223,27 @@ mod tests {
                 }
             }
         }
+
+        for folding_factor in folding_factors {
+            let num_variables = folding_factor..=3 * folding_factor;
+            for num_variables in num_variables {
+                for fold_type in fold_types {
+                    for num_polys in num_polys {
+                        for soundness_type in soundness_type {
+                            for pow_bits in pow_bits {
+                                make_whir_batch_things(
+                                    num_polys,
+                                    num_variables,
+                                    folding_factor,
+                                    soundness_type,
+                                    pow_bits,
+                                    fold_type,
+                                );
+                            }
+                        }
+                    }
+                }
+            }
+        }
     }
 }

src/whir/verifier.rs

@@ -24,7 +24,7 @@ where
     MerkleConfig: Config,
 {
     pub(crate) params: WhirConfig<F, MerkleConfig, PowStrategy>,
-    two_inv: F,
+    pub(crate) two_inv: F,
 }
 
 #[derive(Clone)]