feat: couple of traits to genericize verify circuit (#1305)

recursion/circuit-v2/src/challenger.rs

@@ -34,6 +34,8 @@ pub trait FeltChallenger<C: Config>:
     CanObserveVariable<C, Felt<C::F>> + CanSampleVariable<C, Felt<C::F>> + CanSampleBitsVariable<C>
 {
     fn sample_ext(&mut self, builder: &mut Builder<C>) -> Ext<C::F, C::EF>;
+
+    fn check_witness(&mut self, builder: &mut Builder<C>, nb_bits: usize, witness: Felt<C::F>);
 }
 
 pub trait CanSampleBitsVariable<C: Config> {
@@ -128,29 +130,13 @@ impl<C: Config> DuplexChallengerVariable<C> {
             .expect("output buffer should be non-empty")
     }
 
-    pub fn sample_ext(&mut self, builder: &mut Builder<C>) -> Ext<C::F, C::EF> {
-        let a = self.sample(builder);
-        let b = self.sample(builder);
-        let c = self.sample(builder);
-        let d = self.sample(builder);
-        builder.ext_from_base_slice(&[a, b, c, d])
-    }
-
     pub fn sample_bits(&mut self, builder: &mut Builder<C>, nb_bits: usize) -> Vec<Felt<C::F>> {
         assert!(nb_bits <= NUM_BITS);
         let rand_f = self.sample(builder);
         let mut rand_f_bits = builder.num2bits_v2_f(rand_f, NUM_BITS);
         rand_f_bits.truncate(nb_bits);
         rand_f_bits
     }
-
-    pub fn check_witness(&mut self, builder: &mut Builder<C>, nb_bits: usize, witness: Felt<C::F>) {
-        self.observe(builder, witness);
-        let element_bits = self.sample_bits(builder, nb_bits);
-        for bit in element_bits {
-            builder.assert_felt_eq(bit, C::F::zero());
-        }
-    }
 }
 
 impl<C: Config> CanObserveVariable<C, Felt<C::F>> for DuplexChallengerVariable<C> {
@@ -212,7 +198,24 @@ impl<C: Config> CanObserveVariable<C, VerifyingKeyVariable<C>> for DuplexChallen
 
 impl<C: Config> FeltChallenger<C> for DuplexChallengerVariable<C> {
     fn sample_ext(&mut self, builder: &mut Builder<C>) -> Ext<C::F, C::EF> {
-        DuplexChallengerVariable::sample_ext(self, builder)
+        let a = self.sample(builder);
+        let b = self.sample(builder);
+        let c = self.sample(builder);
+        let d = self.sample(builder);
+        builder.ext_from_base_slice(&[a, b, c, d])
+    }
+
+    fn check_witness(
+        &mut self,
+        builder: &mut Builder<C>,
+        nb_bits: usize,
+        witness: Felt<<C as Config>::F>,
+    ) {
+        self.observe(builder, witness);
+        let element_bits = self.sample_bits(builder, nb_bits);
+        for bit in element_bits {
+            builder.assert_felt_eq(bit, C::F::zero());
+        }
     }
 }
 
@@ -238,6 +241,7 @@ pub(crate) mod tests {
     use sp1_recursion_core_v2::Runtime;
 
     use crate::challenger::DuplexChallengerVariable;
+    use crate::challenger::FeltChallenger;
     use crate::witness::Witness;
 
     use sp1_core::utils::run_test_machine;

recursion/circuit-v2/src/fri.rs

@@ -14,20 +14,21 @@ use std::{
     ops::{Add, Mul},
 };
 
-use crate::challenger::DuplexChallengerVariable;
+use crate::challenger::CanSampleBitsVariable;
+use crate::challenger::FeltChallenger;
 use crate::*;
 
-pub fn verify_shape_and_sample_challenges<C: Config, Mmcs>(
+pub fn verify_shape_and_sample_challenges<C: Config, SC: BabyBearFriConfigVariable<C = C>>(
     builder: &mut Builder<C>,
-    config: &FriConfig<Mmcs>,
+    config: &FriConfig<FriMmcs<SC>>,
     proof: &FriProofVariable<C>,
-    challenger: &mut DuplexChallengerVariable<C>,
+    challenger: &mut SC::FriChallengerVariable,
 ) -> FriChallenges<C> {
     let betas = proof
         .commit_phase_commits
         .iter()
         .map(|&commitment| {
-            challenger.observe_commitment(builder, commitment);
+            challenger.observe_slice(builder, commitment);
             challenger.sample_ext(builder)
         })
         .collect();
@@ -53,17 +54,17 @@ pub fn verify_shape_and_sample_challenges<C: Config, Mmcs>(
     }
 }
 
-pub fn verify_two_adic_pcs<C: Config, Mmcs>(
+pub fn verify_two_adic_pcs<C: Config, SC: BabyBearFriConfigVariable<C = C>>(
     builder: &mut Builder<C>,
-    config: &FriConfig<Mmcs>,
+    config: &FriConfig<FriMmcs<SC>>,
     proof: &TwoAdicPcsProofVariable<C>,
-    challenger: &mut DuplexChallengerVariable<C>,
+    challenger: &mut SC::FriChallengerVariable,
     rounds: Vec<TwoAdicPcsRoundVariable<C>>,
 ) {
     let alpha = challenger.sample_ext(builder);
 
     let fri_challenges =
-        verify_shape_and_sample_challenges(builder, config, &proof.fri_proof, challenger);
+        verify_shape_and_sample_challenges::<C, SC>(builder, config, &proof.fri_proof, challenger);
 
     let log_global_max_height = proof.fri_proof.commit_phase_commits.len() + config.log_blowup;
 
@@ -73,12 +74,11 @@ pub fn verify_two_adic_pcs<C: Config, Mmcs>(
         .zip(&fri_challenges.query_indices)
         .map(|(query_opening, index_bits)| {
             // The powers of alpha, where the ith element is alpha^i.
-            let mut alpha_pows: Vec<Ext<C::F, C::EF>> =
-                [builder.constant(C::EF::one()); 32].to_vec();
+            let mut alpha_pows: [Ext<C::F, C::EF>; 32] = [builder.constant(C::EF::one()); 32];
             let mut ro: [Ext<C::F, C::EF>; 32] =
                 [builder.eval(SymbolicExt::from_f(C::EF::zero())); 32];
 
-            for (batch_opening, round) in izip!(query_opening, rounds.iter().cloned()) {
+            for (batch_opening, round) in zip(query_opening, rounds.iter().cloned()) {
                 let batch_commit = round.batch_commit;
                 let mats = round.mats;
                 let batch_heights = mats
@@ -96,14 +96,15 @@ pub fn verify_two_adic_pcs<C: Config, Mmcs>(
 
                 let reduced_index_bits = index_bits[bits_reduced..].to_vec();
 
-                verify_batch::<C, 1>(
+                verify_batch::<C, SC, 1>(
                     builder,
                     batch_commit,
                     batch_dims,
                     reduced_index_bits,
                     batch_opening.opened_values.clone(),
                     batch_opening.opening_proof.clone(),
                 );
+
                 for (mat_opening, mat) in izip!(&batch_opening.opened_values, mats) {
                     let mat_domain = mat.domain;
                     let mat_points = mat.points;
@@ -137,7 +138,7 @@ pub fn verify_two_adic_pcs<C: Config, Mmcs>(
         })
         .collect::<Vec<_>>();
 
-    verify_challenges(
+    verify_challenges::<C, SC>(
         builder,
         config,
         &proof.fri_proof,
@@ -146,9 +147,9 @@ pub fn verify_two_adic_pcs<C: Config, Mmcs>(
     );
 }
 
-pub fn verify_challenges<C: Config, Mmcs>(
+pub fn verify_challenges<C: Config, SC: BabyBearFriConfigVariable<C = C>>(
     builder: &mut Builder<C>,
-    config: &FriConfig<Mmcs>,
+    config: &FriConfig<FriMmcs<SC>>,
     proof: &FriProofVariable<C>,
     challenges: &FriChallenges<C>,
     reduced_openings: Vec<[Ext<C::F, C::EF>; 32]>,
@@ -160,7 +161,7 @@ pub fn verify_challenges<C: Config, Mmcs>(
         .zip(&proof.query_proofs)
         .zip(reduced_openings)
     {
-        let folded_eval = verify_query(
+        let folded_eval = verify_query::<C, SC>(
             builder,
             proof.commit_phase_commits.clone(),
             index_bits,
@@ -174,7 +175,7 @@ pub fn verify_challenges<C: Config, Mmcs>(
     }
 }
 
-pub fn verify_query<C: Config>(
+pub fn verify_query<C: Config, SC: BabyBearFriConfigVariable<C = C>>(
     builder: &mut Builder<C>,
     commit_phase_commits: Vec<DigestVariable<C>>,
     index_bits: &[Felt<C::F>],
@@ -219,7 +220,7 @@ pub fn verify_query<C: Config>(
             width: 2,
             height: (1 << log_folded_height),
         }];
-        verify_batch::<C, 4>(
+        verify_batch::<C, SC, 4>(
             builder,
             commit,
             dims.to_vec(),
@@ -238,7 +239,7 @@ pub fn verify_query<C: Config>(
     folded_eval
 }
 
-pub fn verify_batch<C: Config, const D: usize>(
+pub fn verify_batch<C: Config, SC: BabyBearFriConfigVariable<C = C>, const D: usize>(
     builder: &mut Builder<C>,
     commit: DigestVariable<C>,
     dimensions: Vec<Dimensions>,
@@ -312,10 +313,9 @@ where
     C: Config,
     R: Variable<C> + Into<<R as Variable<C>>::Expression> + 'a,
     S: Variable<C> + Into<<S as Variable<C>>::Expression> + 'a,
-    <R as Variable<C>>::Expression: AbstractField + 'a,
+    <R as Variable<C>>::Expression: AbstractField,
     <S as Variable<C>>::Expression: Add<Output = <S as Variable<C>>::Expression>
-        + Mul<<R as Variable<C>>::Expression, Output = <S as Variable<C>>::Expression>
-        + 'a,
+        + Mul<<R as Variable<C>>::Expression, Output = <S as Variable<C>>::Expression>,
 {
     let should_swap: <R as Variable<C>>::Expression = should_swap.into();
     let one = <R as Variable<C>>::Expression::one();
@@ -560,7 +560,7 @@ mod tests {
             .into_iter()
             .map(|p| p.map(|x| builder.eval(x)))
             .collect();
-        verify_batch::<_, 1>(
+        verify_batch::<_, SC, 1>(
             &mut builder,
             commit,
             large_mat_dims
@@ -633,11 +633,15 @@ mod tests {
 
         let mut challenger = DuplexChallengerVariable::new(&mut builder);
         let commit: [_; DIGEST_SIZE] = commit.into();
-        let commit = commit.map(|x| builder.eval(x));
-        challenger.observe_commitment(&mut builder, commit);
+        let commit: [Felt<InnerVal>; DIGEST_SIZE] = commit.map(|x| builder.eval(x));
+        challenger.observe_slice(&mut builder, commit);
         let _ = challenger.sample_ext(&mut builder);
-        let fri_challenges =
-            verify_shape_and_sample_challenges(&mut builder, &config, &fri_proof, &mut challenger);
+        let fri_challenges = verify_shape_and_sample_challenges::<InnerConfig, BabyBearPoseidon2>(
+            &mut builder,
+            &config,
+            &fri_proof,
+            &mut challenger,
+        );
 
         for i in 0..fri_challenges_gt.betas.len() {
             builder.assert_ext_eq(
@@ -719,11 +723,17 @@ mod tests {
         let proof = const_two_adic_pcs_proof(&mut builder, proof);
         let (commit, rounds) = const_two_adic_pcs_rounds(&mut builder, commit.into(), os);
         let mut challenger = DuplexChallengerVariable::new(&mut builder);
-        challenger.observe_commitment(&mut builder, commit);
+        challenger.observe_slice(&mut builder, commit);
         let x2 = challenger.sample_ext(&mut builder);
         let x1: Ext<_, _> = builder.constant(x1);
         builder.assert_ext_eq(x1, x2);
-        verify_two_adic_pcs(&mut builder, &config, &proof, &mut challenger, rounds);
+        verify_two_adic_pcs::<_, BabyBearPoseidon2>(
+            &mut builder,
+            &config,
+            &proof,
+            &mut challenger,
+            rounds,
+        );
 
         run_test_recursion(builder.operations, std::iter::empty());
         // let mut backend = ConstraintCompiler::<InnerConfig>::default();

recursion/circuit-v2/src/lib.rs

@@ -1,7 +1,11 @@
 //! Copied from [`sp1_recursion_program`].
 
+use challenger::{CanObserveVariable, DuplexChallengerVariable, FeltChallenger};
 use p3_commit::TwoAdicMultiplicativeCoset;
-use sp1_recursion_compiler::ir::{Config, Ext, Felt};
+use sp1_recursion_compiler::{
+    config::InnerConfig,
+    ir::{Config, Ext, Felt},
+};
 use sp1_recursion_core_v2::DIGEST_SIZE;
 
 pub mod build_wrap_v2;
@@ -66,3 +70,78 @@ pub struct TwoAdicPcsMatsVariable<C: Config> {
     pub points: Vec<Ext<C::F, C::EF>>,
     pub values: Vec<Vec<Ext<C::F, C::EF>>>,
 }
+
+use p3_challenger::{CanObserve, CanSample, FieldChallenger, GrindingChallenger};
+use p3_commit::{ExtensionMmcs, Mmcs};
+use p3_dft::Radix2DitParallel;
+use p3_fri::{FriConfig, TwoAdicFriPcs};
+use p3_matrix::dense::RowMajorMatrix;
+use sp1_recursion_core::stark::config::{BabyBearPoseidon2Outer, OuterValMmcs};
+
+use p3_baby_bear::BabyBear;
+use sp1_core::{stark::StarkGenericConfig, utils::BabyBearPoseidon2};
+
+type EF = <BabyBearPoseidon2 as StarkGenericConfig>::Challenge;
+
+pub type PcsConfig<SC> = FriConfig<
+    ExtensionMmcs<
+        <SC as StarkGenericConfig>::Val,
+        <SC as StarkGenericConfig>::Challenge,
+        <SC as BabyBearFriConfig>::ValMmcs,
+    >,
+>;
+
+pub type FriMmcs<SC> = ExtensionMmcs<BabyBear, EF, <SC as BabyBearFriConfig>::ValMmcs>;
+
+pub trait BabyBearFriConfig:
+    StarkGenericConfig<
+    Val = BabyBear,
+    Challenge = EF,
+    Challenger = Self::FriChallenger,
+    Pcs = TwoAdicFriPcs<
+        BabyBear,
+        Radix2DitParallel,
+        Self::ValMmcs,
+        ExtensionMmcs<BabyBear, EF, Self::ValMmcs>,
+    >,
+>
+{
+    type ValMmcs: Mmcs<BabyBear>;
+    // type RowMajorProverData: Clone;
+    type FriChallenger: CanObserve<<Self::ValMmcs as Mmcs<BabyBear>>::Commitment>
+        + CanSample<EF>
+        + GrindingChallenger<Witness = BabyBear>
+        + FieldChallenger<BabyBear>;
+
+    fn fri_config(&self) -> &FriConfig<FriMmcs<Self>>;
+}
+
+pub trait BabyBearFriConfigVariable: BabyBearFriConfig {
+    // Is this is the best place to put this?
+    type C: Config<F = Self::Val, EF = Self::Challenge>;
+    type FriChallengerVariable: FeltChallenger<Self::C>;
+}
+
+impl BabyBearFriConfig for BabyBearPoseidon2 {
+    type ValMmcs = sp1_core::utils::baby_bear_poseidon2::ValMmcs;
+    type FriChallenger = <Self as StarkGenericConfig>::Challenger;
+
+    fn fri_config(&self) -> &FriConfig<FriMmcs<Self>> {
+        self.pcs().fri_config()
+    }
+}
+
+impl BabyBearFriConfigVariable for BabyBearPoseidon2 {
+    type C = InnerConfig;
+
+    type FriChallengerVariable = DuplexChallengerVariable<Self::C>;
+}
+
+impl BabyBearFriConfig for BabyBearPoseidon2Outer {
+    type ValMmcs = OuterValMmcs;
+    type FriChallenger = <Self as StarkGenericConfig>::Challenger;
+
+    fn fri_config(&self) -> &FriConfig<FriMmcs<Self>> {
+        self.pcs().fri_config()
+    }
+}