perf(recursion): remove extraneous functions from TwoAdicPcs verify (…

…#1152)

* perf(wip): remove extraneous functions

* chore: remove debugging

* chore: fix workflow

.github/workflows/benchmarks.yml

@@ -10,7 +10,7 @@ on:
       - "crates/circuits/**"
       - "crates/vm/**"
       - "crates/toolchain/**"
-      - "crates/extensions/**"
+      - "extensions/**"
       - "benchmarks/**"
       - ".github/workflows/benchmark-call.yml"
       - ".github/workflows/benchmarks.yml"

extensions/native/circuit/examples/fibonacci.rs

@@ -44,7 +44,6 @@ fn main() {
     let expected_value = F::from_canonical_u32(fibonacci(n_val));
     builder.assert_felt_eq(a, expected_value);
 
-    //builder.print_f(a);
     builder.halt();
 
     let program = builder.compile_isa();

extensions/native/recursion/src/fri/mod.rs

@@ -9,12 +9,8 @@ use openvm_native_compiler::{
 use openvm_stark_backend::p3_field::{AbstractField, Field, TwoAdicField};
 pub use two_adic_pcs::*;
 
-use self::types::{
-    DimensionsVariable, FriChallengesVariable, FriConfigVariable, FriProofVariable,
-    FriQueryProofVariable,
-};
+use self::types::{DimensionsVariable, FriConfigVariable, FriQueryProofVariable};
 use crate::{
-    challenger::ChallengerVariable,
     digest::{CanPoseidon2Digest, DigestVariable},
     outer_poseidon2::Poseidon2CircuitBuilder,
     utils::cond_eval,
@@ -27,88 +23,6 @@ pub mod two_adic_pcs;
 pub mod types;
 pub mod witness;
 
-/// Reference: https://github.com/Plonky3/Plonky3/blob/4809fa7bedd9ba8f6f5d3267b1592618e3776c57/fri/src/verifier.rs#L27
-pub fn verify_shape_and_sample_challenges<C: Config>(
-    builder: &mut Builder<C>,
-    config: &FriConfigVariable<C>,
-    proof: &FriProofVariable<C>,
-    challenger: &mut impl ChallengerVariable<C>,
-) -> FriChallengesVariable<C> {
-    let betas: Array<C, Ext<C::F, C::EF>> = builder.array(proof.commit_phase_commits.len());
-
-    builder
-        .range(0, proof.commit_phase_commits.len())
-        .for_each(|i, builder| {
-            let comm = builder.get(&proof.commit_phase_commits, i);
-            challenger.observe_digest(builder, comm);
-            let sample = challenger.sample_ext(builder);
-            builder.set(&betas, i, sample);
-        });
-
-    let final_poly_felts = builder.ext2felt(proof.final_poly);
-    challenger.observe_slice(builder, final_poly_felts);
-
-    let num_query_proofs = proof.query_proofs.len().clone();
-    builder
-        .if_ne(num_query_proofs, RVar::from(config.num_queries))
-        .then(|builder| {
-            builder.error();
-        });
-
-    challenger.check_witness(builder, config.proof_of_work_bits, proof.pow_witness);
-
-    let log_max_height =
-        builder.eval_expr(proof.commit_phase_commits.len() + RVar::from(config.log_blowup));
-    let query_indices = builder.array(config.num_queries);
-    builder.range(0, config.num_queries).for_each(|i, builder| {
-        let index_bits = challenger.sample_bits(builder, log_max_height);
-        builder.set(&query_indices, i, index_bits);
-    });
-
-    FriChallengesVariable {
-        query_indices,
-        betas,
-    }
-}
-
-/// Verifies a set of FRI challenges.
-///
-/// Reference: https://github.com/Plonky3/Plonky3/blob/4809fa7bedd9ba8f6f5d3267b1592618e3776c57/fri/src/verifier.rs#L67
-#[allow(clippy::type_complexity)]
-pub fn verify_challenges<C: Config>(
-    builder: &mut Builder<C>,
-    config: &FriConfigVariable<C>,
-    proof: &FriProofVariable<C>,
-    challenges: &FriChallengesVariable<C>,
-    reduced_openings: &Array<C, Array<C, Ext<C::F, C::EF>>>,
-) where
-    C::F: TwoAdicField,
-    C::EF: TwoAdicField,
-{
-    let log_max_height =
-        builder.eval_expr(proof.commit_phase_commits.len() + RVar::from(config.log_blowup));
-    builder
-        .range(0, challenges.query_indices.len())
-        .for_each(|i, builder| {
-            let index_bits = builder.get(&challenges.query_indices, i);
-            let query_proof = builder.get(&proof.query_proofs, i);
-            let ro = builder.get(reduced_openings, i);
-
-            let folded_eval = verify_query(
-                builder,
-                config,
-                &proof.commit_phase_commits,
-                &index_bits,
-                &query_proof,
-                &challenges.betas,
-                &ro,
-                log_max_height,
-            );
-
-            builder.assert_ext_eq(folded_eval, proof.final_poly);
-        });
-}
-
 /// Verifies a FRI query.
 ///
 /// Currently assumes the index that is accessed is constant.

extensions/native/recursion/src/fri/two_adic_pcs.rs

@@ -9,8 +9,7 @@ use super::{
     types::{
         DimensionsVariable, FriConfigVariable, TwoAdicPcsMatsVariable, TwoAdicPcsRoundVariable,
     },
-    verify_batch, verify_challenges, verify_shape_and_sample_challenges, NestedOpenedValues,
-    TwoAdicMultiplicativeCosetVariable,
+    verify_batch, verify_query, NestedOpenedValues, TwoAdicMultiplicativeCosetVariable,
 };
 use crate::{
     challenger::ChallengerVariable, commit::PcsVariable, digest::DigestVariable,
@@ -29,6 +28,7 @@ use crate::{
 /// Reference:
 /// <https://github.com/Plonky3/Plonky3/blob/27b3127dab047e07145c38143379edec2960b3e1/merkle-tree/src/mmcs.rs#L87>
 /// <https://github.com/Plonky3/Plonky3/blob/27b3127dab047e07145c38143379edec2960b3e1/merkle-tree/src/merkle_tree.rs#L100>
+/// <https://github.com/Plonky3/Plonky3/blob/784b7dd1fa87c1202e63350cc8182d7c5327a7af/fri/src/verifier.rs#L22>
 pub fn verify_two_adic_pcs<C: Config>(
     builder: &mut Builder<C>,
     config: &FriConfigVariable<C>,
@@ -45,21 +45,36 @@ pub fn verify_two_adic_pcs<C: Config>(
     let blowup = config.blowup;
     let alpha = challenger.sample_ext(builder);
 
-    builder.cycle_tracker_start("stage-d-1-verify-shape-and-sample-challenges");
-    let fri_challenges = verify_shape_and_sample_challenges(builder, config, &proof, challenger);
-    builder.cycle_tracker_end("stage-d-1-verify-shape-and-sample-challenges");
+    builder.cycle_tracker_start("stage-d-verifier-verify");
+    let betas: Array<C, Ext<C::F, C::EF>> = builder.array(proof.commit_phase_commits.len());
+    builder
+        .range(0, proof.commit_phase_commits.len())
+        .for_each(|i, builder| {
+            let comm = builder.get(&proof.commit_phase_commits, i);
+            challenger.observe_digest(builder, comm);
+            let sample = challenger.sample_ext(builder);
+            builder.set(&betas, i, sample);
+        });
+    let final_poly_felts = builder.ext2felt(proof.final_poly);
+    challenger.observe_slice(builder, final_poly_felts);
 
-    let log_global_max_height =
-        builder.eval_expr(proof.commit_phase_commits.len() + RVar::from(log_blowup));
+    let num_query_proofs = proof.query_proofs.len().clone();
+    builder
+        .if_ne(num_query_proofs, RVar::from(config.num_queries))
+        .then(|builder| {
+            builder.error();
+        });
+
+    challenger.check_witness(builder, config.proof_of_work_bits, proof.pow_witness);
 
-    let reduced_openings: Array<_, Array<_, Ext<_, _>>> = builder.array(proof.query_proofs.len());
+    let log_max_height =
+        builder.eval_expr(proof.commit_phase_commits.len() + RVar::from(log_blowup));
 
-    builder.cycle_tracker_start("stage-d-2-fri-fold");
     builder
         .range(0, proof.query_proofs.len())
         .for_each(|i, builder| {
             let query_proof = builder.get(&proof.query_proofs, i);
-            let index_bits = builder.get(&fri_challenges.query_indices, i);
+            let index_bits = challenger.sample_bits(builder, log_max_height);
 
             let ro: Array<C, Ext<C::F, C::EF>> = builder.array(32);
             let alpha_pow: Array<C, Ext<C::F, C::EF>> = builder.array(32);
@@ -128,8 +143,7 @@ pub fn verify_two_adic_pcs<C: Config>(
                 });
                 let permed_opened_values = NestedOpenedValues::Felt(permed_opened_values);
 
-                let bits_reduced: Usize<_> =
-                    builder.eval(log_global_max_height - log_batch_max_height);
+                let bits_reduced: Usize<_> = builder.eval(log_max_height - log_batch_max_height);
                 let index_bits_shifted_v1 = index_bits.shift(builder, bits_reduced);
 
                 builder.cycle_tracker_start("verify-batch");
@@ -161,8 +175,7 @@ pub fn verify_two_adic_pcs<C: Config>(
                         let cur_ro = builder.get(&ro, log_height);
                         let cur_alpha_pow = builder.get(&alpha_pow, log_height);
 
-                        let bits_reduced: Usize<_> =
-                            builder.eval(log_global_max_height - log_height);
+                        let bits_reduced: Usize<_> = builder.eval(log_max_height - log_height);
                         let index_bits_shifted = index_bits.shift(builder, bits_reduced);
 
                         let two_adic_generator = config.get_two_adic_generator(builder, log_height);
@@ -209,13 +222,20 @@ pub fn verify_two_adic_pcs<C: Config>(
                 builder.cycle_tracker_end("compute-reduced-opening");
             });
 
-            builder.set_value(&reduced_openings, i, ro);
+            let folded_eval = verify_query(
+                builder,
+                config,
+                &proof.commit_phase_commits,
+                &index_bits,
+                &query_proof,
+                &betas,
+                &ro,
+                log_max_height,
+            );
+
+            builder.assert_ext_eq(folded_eval, proof.final_poly);
         });
-    builder.cycle_tracker_end("stage-d-2-fri-fold");
-
-    builder.cycle_tracker_start("stage-d-3-verify-challenges");
-    verify_challenges(builder, config, &proof, &fri_challenges, &reduced_openings);
-    builder.cycle_tracker_end("stage-d-3-verify-challenges");
+    builder.cycle_tracker_end("stage-d-verifier-verify");
 }
 
 impl<C: Config> FromConstant<C> for TwoAdicPcsRoundVariable<C>