fix: two adic pcs issue when verifying tables of different heights (#442

)

Cargo.toml

@@ -46,3 +46,26 @@ p3-symmetric = { git = "https://github.com/Plonky3/Plonky3.git", branch = "sp1"
 p3-uni-stark = { git = "https://github.com/Plonky3/Plonky3.git", branch = "sp1" }
 p3-maybe-rayon = { git = "https://github.com/Plonky3/Plonky3.git", branch = "sp1" }
 p3-bn254-fr = { git = "https://github.com/Plonky3/Plonky3.git", branch = "sp1" }
+
+# For local development.
+#
+# p3-air = { path = "../Plonky3/air" }
+# p3-field = { path = "../Plonky3/field" }
+# p3-commit = { path = "../Plonky3/commit" }
+# p3-matrix = { path = "../Plonky3/matrix" }
+# p3-baby-bear = { path = "../Plonky3/baby-bear" }
+# p3-util = { path = "../Plonky3/util" }
+# p3-challenger = { path = "../Plonky3/challenger" }
+# p3-dft = { path = "../Plonky3/dft" }
+# p3-fri = { path = "../Plonky3/fri" }
+# p3-goldilocks = { path = "../Plonky3/goldilocks" }
+# p3-keccak = { path = "../Plonky3/keccak" }
+# p3-keccak-air = { path = "../Plonky3/keccak-air" }
+# p3-blake3 = { path = "../Plonky3/blake3" }
+# p3-mds = { path = "../Plonky3/mds" }
+# p3-merkle-tree = { path = "../Plonky3/merkle-tree" }
+# p3-poseidon2 = { path = "../Plonky3/poseidon2" }
+# p3-symmetric = { path = "../Plonky3/symmetric" }
+# p3-uni-stark = { path = "../Plonky3/uni-stark" }
+# p3-maybe-rayon = { path = "../Plonky3/maybe-rayon" }
+# p3-bn254-fr = { path = "../Plonky3/bn254-fr" }

recursion/compiler/src/verifier/fri/mod.rs

@@ -239,16 +239,17 @@ pub fn verify_batch<C: Config, const D: usize>(
         );
 
         let new_root = builder.poseidon2_compress(&left, &right);
-
         builder.assign(root.clone(), new_root);
         builder.assign(current_height, current_height * (C::N::two().inverse()));
 
         let next_height = builder.get(&dimensions, index).height;
-        builder.if_eq(next_height, current_height).then(|builder| {
-            let next_height_openings_digest =
-                reduce::<C, D>(builder, index, &dimensions, current_height, &opened_values);
-            let new_root = builder.poseidon2_compress(&root, &next_height_openings_digest);
-            builder.assign(root.clone(), new_root);
+        builder.if_ne(index, dimensions.len()).then(|builder| {
+            builder.if_eq(next_height, current_height).then(|builder| {
+                let next_height_openings_digest =
+                    reduce::<C, D>(builder, index, &dimensions, current_height, &opened_values);
+                let new_root = builder.poseidon2_compress(&root, &next_height_openings_digest);
+                builder.assign(root.clone(), new_root);
+            });
         })
     });
 
@@ -273,23 +274,27 @@ pub fn reduce<C: Config, const D: usize>(
 ) -> Array<C, Felt<C::F>> {
     let nb_opened_values = builder.eval(C::N::zero());
     let mut flattened_opened_values = builder.dyn_array(8192);
-    builder.range(dim_idx, dims.len()).for_each(|i, builder| {
-        let height = builder.get(dims, i).height;
-        builder.if_eq(height, curr_height_padded).then(|builder| {
-            let opened_values = builder.get(opened_values, i);
-            builder
-                .range(0, opened_values.len())
-                .for_each(|j, builder| {
-                    let opened_value = builder.get(&opened_values, j);
-                    let opened_value_flat = builder.ext2felt(opened_value);
-                    for k in 0..D {
-                        let base = builder.get(&opened_value_flat, k);
-                        builder.set(&mut flattened_opened_values, nb_opened_values, base);
-                        builder.assign(nb_opened_values, nb_opened_values + C::N::one());
-                    }
-                });
+    let start_dim_idx: Var<_> = builder.eval(dim_idx);
+    builder
+        .range(start_dim_idx, dims.len())
+        .for_each(|i, builder| {
+            let height = builder.get(dims, i).height;
+            builder.if_eq(height, curr_height_padded).then(|builder| {
+                let opened_values = builder.get(opened_values, i);
+                builder
+                    .range(0, opened_values.len())
+                    .for_each(|j, builder| {
+                        let opened_value = builder.get(&opened_values, j);
+                        let opened_value_flat = builder.ext2felt(opened_value);
+                        for k in 0..D {
+                            let base = builder.get(&opened_value_flat, k);
+                            builder.set(&mut flattened_opened_values, nb_opened_values, base);
+                            builder.assign(nb_opened_values, nb_opened_values + C::N::one());
+                        }
+                    });
+                builder.assign(dim_idx, dim_idx + C::N::one());
+            });
         });
-    });
     flattened_opened_values.truncate(builder, Usize::Var(nb_opened_values));
     builder.poseidon2_hash(&flattened_opened_values)
 }

recursion/compiler/src/verifier/fri/pcs.rs

@@ -55,8 +55,12 @@ pub fn verify_two_adic_pcs<C: Config>(
     let fri_challenges =
         verify_shape_and_sample_challenges(builder, config, &proof.fri_proof, challenger);
 
-    // let commit_phase_commits_len = builder.materialize(proof.fri_proof.commit_phase_commits.len());
-    // let log_max_height: Var<_> = builder.eval(commit_phase_commits_len + config.log_blowup);
+    let commit_phase_commits_len = proof
+        .fri_proof
+        .commit_phase_commits
+        .len()
+        .materialize(builder);
+    let log_global_max_height: Var<_> = builder.eval(commit_phase_commits_len + config.log_blowup);
 
     let mut reduced_openings: Array<C, Array<C, Ext<C::F, C::EF>>> =
         builder.array(proof.query_openings.len());
@@ -82,22 +86,31 @@ pub fn verify_two_adic_pcs<C: Config>(
                 let batch_commit = round.batch_commit;
                 let mats = round.mats;
 
+                let mut batch_heights_log2: Array<C, Var<C::N>> = builder.array(mats.len());
+                builder.range(0, mats.len()).for_each(|k, builder| {
+                    let mat = builder.get(&mats, k);
+                    let height_log2: Var<_> = builder.eval(mat.domain.log_n + config.log_blowup);
+                    builder.set(&mut batch_heights_log2, k, height_log2);
+                });
                 let mut batch_dims: Array<C, Dimensions<C>> = builder.array(mats.len());
                 builder.range(0, mats.len()).for_each(|k, builder| {
                     let mat = builder.get(&mats, k);
                     let dim = Dimensions::<C> {
-                        height: mat.domain.size(),
+                        height: builder.eval(mat.domain.size() * C::N::two()), // TODO: fix this to use blowup
                     };
                     builder.set(&mut batch_dims, k, dim);
                 });
 
+                let log_batch_max_height = builder.get(&batch_heights_log2, 0);
+                let bits_reduced: Var<_> =
+                    builder.eval(log_global_max_height - log_batch_max_height);
                 let index_bits = builder.num2bits_v(index);
-
+                let index_bits_shifted_v1 = index_bits.shift(builder, bits_reduced);
                 fri::verify_batch::<C, 1>(
                     builder,
                     &batch_commit,
                     batch_dims,
-                    index_bits,
+                    index_bits_shifted_v1,
                     batch_opening.opened_values.clone(),
                     &batch_opening.opening_proof,
                 );
@@ -114,10 +127,12 @@ pub fn verify_two_adic_pcs<C: Config>(
                         let log_height: Var<C::N> =
                             builder.eval(log2_domain_size + config.log_blowup);
 
-                        // let bits_reduced: Var<C::N> = builder.eval(log_max_height - log_height);
-                        // TODO: divide by bits_reduced
+                        let bits_reduced: Var<C::N> =
+                            builder.eval(log_global_max_height - log_height);
+                        let index_bits_shifted_v2 = index_bits.shift(builder, bits_reduced);
+                        let index_shifted_v2 = builder.bits_to_num_var(&index_bits_shifted_v2);
                         let rev_reduced_index =
-                            builder.reverse_bits_len(index, Usize::Var(log_height));
+                            builder.reverse_bits_len(index_shifted_v2, Usize::Var(log_height));
                         let rev_reduced_index = rev_reduced_index.materialize(builder);
 
                         let g = builder.generator();

recursion/compiler/src/verifier/fri/utils.rs

@@ -119,11 +119,12 @@ impl<C: Config> Builder<C> {
 
     pub fn bits_to_num_var(&mut self, bits: &Array<C, Var<C::N>>) -> Var<C::N> {
         let num: Var<_> = self.eval(C::N::zero());
-        for i in 0..NUM_BITS {
-            let bit = self.get(bits, i);
-            // Add `bit * 2^i` to the sum.
-            self.assign(num, num + bit * C::N::from_canonical_u32(1 << i));
-        }
+        let power: Var<_> = self.eval(C::N::one());
+        self.range(0, bits.len()).for_each(|i, builder| {
+            let bit = builder.get(bits, i);
+            builder.assign(num, num + bit * power);
+            builder.assign(power, power * C::N::from_canonical_u32(2));
+        });
         num
     }
 

recursion/compiler/tests/pcs.rs

@@ -237,7 +237,7 @@ fn default_fri_config() -> FriConfig<ChallengeMmcs> {
 #[test]
 fn test_pcs_verify() {
     let mut rng = &mut OsRng;
-    let log_degrees = &[16];
+    let log_degrees = &[16, 9, 7, 4, 2];
     let perm = Perm::new(8, 22, RC_16_30.to_vec(), DiffusionMatrixBabybear);
     let fri_config = default_fri_config();
     let hash = Hash::new(perm.clone());
@@ -257,7 +257,7 @@ fn test_pcs_verify() {
         .map(|&d| {
             (
                 <CustomPcs as Pcs<Challenge, Challenger>>::natural_domain_for_degree(&pcs, 1 << d),
-                RowMajorMatrix::<Val>::rand(&mut rng, 1 << d, 100),
+                RowMajorMatrix::<Val>::rand(&mut rng, 1 << d, 10),
             )
         })
         .collect::<Vec<_>>();