Repository: Plonky3/Plonky3 (Rust)

Issue:
Multi batch opening with different max heights in FRI
When attempting to make a `open_multi_batches` call with batches that have different maximum heights, the code panics for accessing invalid rows. 

Here is an example of a test failure:
https://github.com/Plonky3/Plonky3/blob/multi-batch-fail/fri/tests/pcs.rs#L108

In the context of FRI, the problem seems to be that the shifts of degrees on the tree assume the `log_max_degree` is the same for all batches in https://github.com/Plonky3/Plonky3/blob/main/fri/src/two_adic_pcs.rs#L299:
```rust
        let query_openings = query_indices
            .into_iter()
            .map(|index| {
                prover_data_and_points
                    .iter()
                    .map(|(data, _)| {
                        let (opened_values, opening_proof) = self.mmcs.open_batch(index, data);
                        BatchOpening {
                            opened_values,
                            opening_proof,
                        }
                    })
                    .collect()
            })
            .collect();
```
That is because `query_indices` are computed with the maximum height of all batches as basis, whereas the shifts in `open_batch`, which are computed by the following code:
```rust
        let openings = prover_data
            .leaves
            .iter()
            .map(|matrix| {
                let log2_height = log2_ceil_usize(matrix.height());
                let bits_reduced = log_max_height - log2_height;
                let reduced_index = index >> bits_reduced;
                matrix.row(reduced_index).collect()
            })
            .collect_vec();
 ```
are computing the shift with respect to the maximum height of that batch. This ends up with overflow in the rows if the `max_height_for batch < max_height_on_all_batches`.


I have also checked an earlier commit of FRI, and had the same failure, so this behavior did not originate from recent fixes. 

Is this the intended behavior of `open_multi_batches`? If so, we should probably add some comments in the trait as the errors were a little confusing to debug. If it's not the intended behavior, happy to explore a fix (some naive fixes I tried didn't work). 


Discussion:
I see a few options for addressing this:
1. Replace `Mmcs::open_batch` with a method called `open_multi_batches`, which takes multiple `ProverData`s and handles openings for all of them in a single call, similar to `UnivariatePcs::open_multi_batches`.
2. Change `Mmcs::open_batch` to take a `log_max_height` param, representing the max height of any matrix.
3. Leave `Mmcs` unchanged, and instead change `TwoAdicFriPcs` to determine the "global max height" and shift off certain bits when it calls `Mmcs::open_batch`.

I think I'm leaning toward (1) - it would make `Mmcs` more complicated, but the extra layer of batching already exists in PCS APIs, and it seems natural for `Mmcs` to have the same batching approach. That said, I'll need to try coding and see how clean/simple (or not) the changes turn out before I can confidently recommend anything.
(1) does sound really nice, I can try to implement it too and see if it works
@tamirhemo that would be great, I'll tentatively reassign to you, but let me know if I (and/or @rjeli) can help
WIP pull request: https://github.com/Plonky3/Plonky3/pull/277
I made progress on (1) but it was starting to feel like too much extra logic without a great reason for it. (3) is a very easy change, see #286, and seems like a reasonable solution I think?
That's great, thank you! 

I think I ran into some issues with (1) as well since it seems like the logic of open_multi_batches makes less sense for the MMCS given how it's used in Pcs. 



Patch:
diff --git a/fri/src/two_adic_pcs.rs b/fri/src/two_adic_pcs.rs
index 9f99ca213..6ec584069 100644
--- a/fri/src/two_adic_pcs.rs
+++ b/fri/src/two_adic_pcs.rs
@@ -197,15 +197,16 @@ where
             .iter()
             .map(|(data, points)| (self.mmcs.get_matrices(data), points))
             .collect_vec();
-
-        let max_width = mats_and_points
+        let mats = mats_and_points
             .iter()
             .flat_map(|(mats, _)| mats)
-            .map(|m| m.width())
-            .max()
-            .unwrap();
+            .collect_vec();
 
-        let alpha_reducer = PowersReducer::<Val, Challenge>::new(alpha, max_width);
+        let global_max_width = mats.iter().map(|m| m.width()).max().unwrap();
+        let global_max_height = mats.iter().map(|m| m.height()).max().unwrap();
+        let log_global_max_height = log2_strict_usize(global_max_height);
+
+        let alpha_reducer = PowersReducer::<Val, Challenge>::new(alpha, global_max_width);
 
         // For each unique opening point z, we will find the largest degree bound
         // for that point, and precompute 1/(X - z) for the largest subgroup (in bitrev order).
@@ -272,8 +273,11 @@ where
                 rounds
                     .iter()
                     .map(|(data, _)| {
-                        // needs to recombine decomposed openings.. or something...
-                        let (opened_values, opening_proof) = self.mmcs.open_batch(index, data);
+                        let log_max_height = log2_strict_usize(self.mmcs.get_max_height(data));
+                        let bits_reduced = log_global_max_height - log_max_height;
+                        let reduced_index = index >> bits_reduced;
+                        let (opened_values, opening_proof) =
+                            self.mmcs.open_batch(reduced_index, data);
                         BatchOpening {
                             opened_values,
                             opening_proof,
@@ -320,7 +324,8 @@ where
             verifier::verify_shape_and_sample_challenges(&self.fri, &proof.fri_proof, challenger)
                 .map_err(VerificationError::FriError)?;
 
-        let log_max_height = proof.fri_proof.commit_phase_commits.len() + self.fri.log_blowup;
+        let log_global_max_height =
+            proof.fri_proof.commit_phase_commits.len() + self.fri.log_blowup;
 
         let reduced_openings: Vec<[Challenge; 32]> = proof
             .query_openings
@@ -329,19 +334,27 @@ where
             .map(|(query_opening, &index)| {
                 let mut ro = [Challenge::zero(); 32];
                 let mut alpha_pow = [Challenge::one(); 32];
+
                 for (batch_opening, (batch_commit, mats)) in izip!(query_opening, &rounds) {
-                    let batch_dims: Vec<Dimensions> = mats
+                    let batch_heights = mats
                         .iter()
-                        .map(|(domain, _)| Dimensions {
-                            // todo: mmcs doesn't really need width
-                            width: 0,
-                            height: domain.size(),
-                        })
+                        .map(|(domain, _)| domain.size() << self.fri.log_blowup)
                         .collect_vec();
+                    let batch_dims = batch_heights
+                        .iter()
+                        // TODO: MMCS doesn't really need width; we put 0 for now.
+                        .map(|&height| Dimensions { width: 0, height })
+                        .collect_vec();
+
+                    let batch_max_height = batch_heights.iter().max().expect("Empty batch?");
+                    let log_batch_max_height = log2_strict_usize(*batch_max_height);
+                    let bits_reduced = log_global_max_height - log_batch_max_height;
+                    let reduced_index = index >> bits_reduced;
+
                     self.mmcs.verify_batch(
                         batch_commit,
                         &batch_dims,
-                        index,
+                        reduced_index,
                         &batch_opening.opened_values,
                         &batch_opening.opening_proof,
                     )?;
@@ -350,7 +363,7 @@ where
                     {
                         let log_height = log2_strict_usize(mat_domain.size()) + self.fri.log_blowup;
 
-                        let bits_reduced = log_max_height - log_height;
+                        let bits_reduced = log_global_max_height - log_height;
                         let rev_reduced_index = reverse_bits_len(index >> bits_reduced, log_height);
 
                         let x = Val::generator()


Test patch:
diff --git a/fri/tests/pcs.rs b/fri/tests/pcs.rs
index d5dc7f180..a47e76a43 100644
--- a/fri/tests/pcs.rs
+++ b/fri/tests/pcs.rs
@@ -1,3 +1,4 @@
+use itertools::{izip, Itertools};
 use p3_baby_bear::{BabyBear, DiffusionMatrixBabybear};
 use p3_challenger::{CanObserve, DuplexChallenger, FieldChallenger};
 use p3_commit::{ExtensionMmcs, Pcs};
@@ -11,96 +12,123 @@ use p3_poseidon2::Poseidon2;
 use p3_symmetric::{PaddingFreeSponge, TruncatedPermutation};
 use rand::thread_rng;
 
-fn make_test_fri_pcs(log_degrees: &[usize]) {
+type Val = BabyBear;
+type Challenge = BinomialExtensionField<Val, 4>;
+
+type Perm = Poseidon2<Val, DiffusionMatrixBabybear, 16, 7>;
+type MyHash = PaddingFreeSponge<Perm, 16, 8, 8>;
+type MyCompress = TruncatedPermutation<Perm, 2, 8, 16>;
+
+type ValMmcs =
+    FieldMerkleTreeMmcs<<Val as Field>::Packing, <Val as Field>::Packing, MyHash, MyCompress, 8>;
+
+type ChallengeMmcs = ExtensionMmcs<Val, Challenge, ValMmcs>;
+
+type Dft = Radix2DitParallel;
+
+type Challenger = DuplexChallenger<Val, Perm, 16>;
+
+type MyPcs = TwoAdicFriPcs<Val, Dft, ValMmcs, ChallengeMmcs>;
+
+fn make_test_fri_pcs(log_degrees_by_round: &[&[usize]]) {
+    let num_rounds = log_degrees_by_round.len();
     let mut rng = thread_rng();
-    type Val = BabyBear;
-    type Challenge = BinomialExtensionField<Val, 4>;
 
-    type Perm = Poseidon2<Val, DiffusionMatrixBabybear, 16, 7>;
     let perm = Perm::new_from_rng(8, 22, DiffusionMatrixBabybear, &mut rng);
-
-    type MyHash = PaddingFreeSponge<Perm, 16, 8, 8>;
     let hash = MyHash::new(perm.clone());
-
-    type MyCompress = TruncatedPermutation<Perm, 2, 8, 16>;
     let compress = MyCompress::new(perm.clone());
 
-    type ValMmcs = FieldMerkleTreeMmcs<
-        <Val as Field>::Packing,
-        <Val as Field>::Packing,
-        MyHash,
-        MyCompress,
-        8,
-    >;
     let val_mmcs = ValMmcs::new(hash, compress);
-
-    type ChallengeMmcs = ExtensionMmcs<Val, Challenge, ValMmcs>;
     let challenge_mmcs = ChallengeMmcs::new(val_mmcs.clone());
 
-    type Dft = Radix2DitParallel;
-    let dft = Dft {};
-
-    type Challenger = DuplexChallenger<Val, Perm, 16>;
-
     let fri_config = FriConfig {
         log_blowup: 1,
         num_queries: 10,
         proof_of_work_bits: 8,
         mmcs: challenge_mmcs,
     };
-    type MyPcs = TwoAdicFriPcs<Val, Dft, ValMmcs, ChallengeMmcs>;
-    let max_log_n = log_degrees.iter().copied().max().unwrap();
-    let pcs: MyPcs = MyPcs::new(max_log_n, dft, val_mmcs, fri_config);
+    let max_log_n = log_degrees_by_round
+        .iter()
+        .copied()
+        .flatten()
+        .copied()
+        .max()
+        .unwrap();
+    let pcs = MyPcs::new(max_log_n, Dft {}, val_mmcs, fri_config);
 
     let mut challenger = Challenger::new(perm.clone());
 
-    let domains_and_polys = log_degrees
+    let domains_and_polys_by_round = log_degrees_by_round
         .iter()
-        .map(|&d| {
-            (
-                <MyPcs as Pcs<Challenge, Challenger>>::natural_domain_for_degree(&pcs, 1 << d),
-                RowMajorMatrix::<Val>::rand(&mut rng, 1 << d, 10),
-            )
+        .map(|log_degrees| {
+            log_degrees
+                .iter()
+                .map(|&log_degree| {
+                    let d = 1 << log_degree;
+                    (
+                        <MyPcs as Pcs<Challenge, Challenger>>::natural_domain_for_degree(&pcs, d),
+                        RowMajorMatrix::<Val>::rand(&mut rng, d, 10),
+                    )
+                })
+                .collect_vec()
         })
-        .collect::<Vec<_>>();
-
-    let (commit, data) =
-        <MyPcs as Pcs<Challenge, Challenger>>::commit(&pcs, domains_and_polys.clone());
+        .collect_vec();
 
-    challenger.observe(commit);
+    let (commits_by_round, data_by_round): (Vec<_>, Vec<_>) = domains_and_polys_by_round
+        .iter()
+        .map(|domains_and_polys| {
+            <MyPcs as Pcs<Challenge, Challenger>>::commit(&pcs, domains_and_polys.clone())
+        })
+        .unzip();
+    assert_eq!(commits_by_round.len(), num_rounds);
+    assert_eq!(data_by_round.len(), num_rounds);
+    challenger.observe_slice(&commits_by_round);
 
-    let zeta = challenger.sample_ext_element::<Challenge>();
+    let zeta: Challenge = challenger.sample_ext_element();
 
-    let points = domains_and_polys
+    let points_by_round = log_degrees_by_round
         .iter()
-        .map(|_| vec![zeta])
-        .collect::<Vec<_>>();
-
-    let (opening, proof) = pcs.open(vec![(&data, points)], &mut challenger);
+        .map(|log_degrees| vec![vec![zeta]; log_degrees.len()])
+        .collect_vec();
+    let data_and_points = data_by_round.iter().zip(points_by_round).collect();
+    let (opening_by_round, proof) = pcs.open(data_and_points, &mut challenger);
+    assert_eq!(opening_by_round.len(), num_rounds);
 
-    // verify the proof.
+    // Verify the proof.
     let mut challenger = Challenger::new(perm);
-    challenger.observe(commit);
-    let _ = challenger.sample_ext_element::<Challenge>();
-
-    let os = domains_and_polys
-        .iter()
-        .zip(&opening[0])
-        .map(|((domain, _), mat_openings)| (*domain, vec![(zeta, mat_openings[0].clone())]))
-        .collect();
-    pcs.verify(vec![(commit, os)], &proof, &mut challenger)
+    challenger.observe_slice(&commits_by_round);
+    let verifier_zeta: Challenge = challenger.sample_ext_element();
+    assert_eq!(verifier_zeta, zeta);
+
+    let commits_and_claims_by_round = izip!(
+        commits_by_round,
+        domains_and_polys_by_round,
+        opening_by_round
+    )
+    .map(|(commit, domains_and_polys, openings)| {
+        let claims = domains_and_polys
+            .iter()
+            .zip(openings)
+            .map(|((domain, _), mat_openings)| (*domain, vec![(zeta, mat_openings[0].clone())]))
+            .collect_vec();
+        (commit, claims)
+    })
+    .collect_vec();
+    assert_eq!(commits_and_claims_by_round.len(), num_rounds);
+
+    pcs.verify(commits_and_claims_by_round, &proof, &mut challenger)
         .unwrap()
 }
 
 #[test]
 fn test_fri_pcs_single() {
-    make_test_fri_pcs(&[3]);
+    make_test_fri_pcs(&[&[3]]);
 }
 
 #[test]
 fn test_fri_pcs_many_equal() {
     for i in 1..4 {
-        make_test_fri_pcs(&[i; 5]);
+        make_test_fri_pcs(&[&[i; 5]]);
     }
 }
 
@@ -108,7 +136,7 @@ fn test_fri_pcs_many_equal() {
 fn test_fri_pcs_many_different() {
     for i in 2..4 {
         let degrees = (3..3 + i).collect::<Vec<_>>();
-        make_test_fri_pcs(&degrees);
+        make_test_fri_pcs(&[&degrees]);
     }
 }
 
@@ -116,6 +144,19 @@ fn test_fri_pcs_many_different() {
 fn test_fri_pcs_many_different_rev() {
     for i in 2..4 {
         let degrees = (3..3 + i).rev().collect::<Vec<_>>();
-        make_test_fri_pcs(&degrees);
+        make_test_fri_pcs(&[&degrees]);
     }
 }
+
+#[test]
+fn test_fri_pcs_multiple_rounds() {
+    make_test_fri_pcs(&[&[3]]);
+    make_test_fri_pcs(&[&[3], &[3]]);
+    make_test_fri_pcs(&[&[3], &[2]]);
+    make_test_fri_pcs(&[&[2], &[3]]);
+    make_test_fri_pcs(&[&[3, 4], &[3, 4]]);
+    make_test_fri_pcs(&[&[4, 2], &[4, 2]]);
+    make_test_fri_pcs(&[&[2, 2], &[3, 3]]);
+    make_test_fri_pcs(&[&[3, 3], &[2, 2]]);
+    make_test_fri_pcs(&[&[2], &[3, 3]]);
+}