feat: infrastructure for proof shape merkle proofs

crates/prover/src/lib.rs

@@ -30,6 +30,7 @@ use std::{
     thread,
 };
 
+use itertools::Itertools;
 use lru::LruCache;
 
 use tracing::instrument;
@@ -69,11 +70,14 @@ use sp1_recursion_core_v2::{
 };
 
 use sp1_recursion_circuit_v2::{
+    hash::FieldHasher,
     machine::{
-        SP1CompressShape, SP1CompressVerifier, SP1CompressWitnessValues, SP1DeferredVerifier,
-        SP1DeferredWitnessValues, SP1RecursionShape, SP1RecursionWitnessValues,
-        SP1RecursiveVerifier,
+        SP1CompressShape, SP1CompressVerifier, SP1CompressWithVKeyVerifier,
+        SP1CompressWithVKeyWitnessValues, SP1CompressWitnessValues, SP1DeferredVerifier,
+        SP1DeferredWitnessValues, SP1MerkleProofWitnessValues, SP1RecursionShape,
+        SP1RecursionWitnessValues, SP1RecursiveVerifier,
     },
+    merkle_tree::MerkleTree,
     witness::Witnessable,
 };
 
@@ -125,6 +129,12 @@ pub struct SP1Prover<C: SP1ProverComponents = DefaultProverComponents> {
     pub compress_programs: Mutex<LruCache<SP1CompressShape, Arc<RecursionProgram<BabyBear>>>>,
 
     pub compress_cache_misses: AtomicUsize,
+
+    pub root: <InnerSC as FieldHasher<BabyBear>>::Digest,
+
+    pub allowed_vkeys: Vec<<InnerSC as FieldHasher<BabyBear>>::Digest>,
+
+    pub merkle_tree: MerkleTree<BabyBear, InnerSC>,
 }
 
 impl<C: SP1ProverComponents> SP1Prover<C> {
@@ -166,6 +176,10 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
         )
         .expect("PROVER_COMPRESS_CACHE_SIZE must be a non-zero usize");
 
+        let allowed_vkeys = vec![<InnerSC as FieldHasher<BabyBear>>::Digest::default(); 1 << 16];
+
+        let (root, merkle_tree) = MerkleTree::commit(allowed_vkeys.clone());
+
         Self {
             core_prover,
             compress_prover,
@@ -175,6 +189,9 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
             recursion_cache_misses: AtomicUsize::new(0),
             compress_programs: Mutex::new(LruCache::new(compress_cache_size)),
             compress_cache_misses: AtomicUsize::new(0),
+            root,
+            merkle_tree,
+            allowed_vkeys,
         }
     }
 
@@ -276,7 +293,7 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
 
     pub fn compress_program(
         &self,
-        input: &SP1CompressWitnessValues<InnerSC>,
+        input: &SP1CompressWithVKeyWitnessValues<InnerSC>,
     ) -> Arc<RecursionProgram<BabyBear>> {
         let mut cache = self.compress_programs.lock().unwrap();
         cache
@@ -287,7 +304,11 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
                 let builder_span = tracing::debug_span!("build compress program").entered();
                 let mut builder = Builder::<InnerConfig>::default();
                 let input = input.read(&mut builder);
-                SP1CompressVerifier::verify(&mut builder, self.compress_prover.machine(), input);
+                SP1CompressWithVKeyVerifier::verify(
+                    &mut builder,
+                    self.compress_prover.machine(),
+                    input,
+                );
                 let operations = builder.operations;
                 builder_span.exit();
 
@@ -303,7 +324,7 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
 
     pub fn shrink_program(
         &self,
-        input: &SP1CompressWitnessValues<InnerSC>,
+        input: &SP1CompressWithVKeyWitnessValues<InnerSC>,
     ) -> Arc<RecursionProgram<BabyBear>> {
         self.compress_program(input)
     }
@@ -475,6 +496,7 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
     ) -> Result<SP1ReduceProof<InnerSC>, SP1RecursionProverError> {
         // Set the batch size for the reduction tree.
         let batch_size = 2;
+        let first_layer_batch_size = 2;
         let shard_proofs = &proof.proof.0;
 
         // Get the leaf challenger.
@@ -491,7 +513,7 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
             &leaf_challenger,
             shard_proofs,
             &deferred_proofs,
-            batch_size,
+            first_layer_batch_size,
         );
 
         // Calculate the expected height of the tree.
@@ -566,8 +588,15 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
                                 }
                                 SP1CircuitWitness::Compress(input) => {
                                     let mut witness_stream = Vec::new();
-                                    Witnessable::<InnerConfig>::write(&input, &mut witness_stream);
-                                    (self.compress_program(&input), witness_stream)
+
+                                    let input_with_merkle = self.make_merkle_proofs(input);
+
+                                    Witnessable::<InnerConfig>::write(
+                                        &input_with_merkle,
+                                        &mut witness_stream,
+                                    );
+
+                                    (self.compress_program(&input_with_merkle), witness_stream)
                                 }
                             });
 
@@ -811,7 +840,9 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
             is_complete: true,
         };
 
-        let program = self.shrink_program(&input);
+        let input_with_merkle = self.make_merkle_proofs(input);
+
+        let program = self.shrink_program(&input_with_merkle);
 
         // Run the compress program.
         let mut runtime = RecursionRuntime::<Val<InnerSC>, Challenge<InnerSC>, _>::new(
@@ -820,7 +851,7 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
         );
 
         let mut witness_stream = Vec::new();
-        Witnessable::<InnerConfig>::write(&input, &mut witness_stream);
+        Witnessable::<InnerConfig>::write(&input_with_merkle, &mut witness_stream);
 
         runtime.witness_stream = witness_stream.into();
 
@@ -985,6 +1016,27 @@ impl<C: SP1ProverComponents> SP1Prover<C> {
         digest
     }
 
+    fn make_merkle_proofs(
+        &self,
+        input: SP1CompressWitnessValues<CoreSC>,
+    ) -> SP1CompressWithVKeyWitnessValues<CoreSC>
+where {
+        // TODO: make an index based on the key itself.
+        let vk_indices = input.vks_and_proofs.iter().map(|(_, _)| 0).collect_vec();
+
+        let proofs = vk_indices
+            .iter()
+            .map(|index| {
+                let (_, proof) = MerkleTree::open(&self.merkle_tree, *index);
+                proof
+            })
+            .collect();
+
+        let merkle_val = SP1MerkleProofWitnessValues { root: self.root, vk_merkle_proofs: proofs };
+
+        SP1CompressWithVKeyWitnessValues { compress_val: input, merkle_val }
+    }
+
     fn check_for_high_cycles(cycles: u64) {
         if cycles > 100_000_000 {
             tracing::warn!(

crates/recursion/circuit-v2/src/hash.rs

@@ -115,6 +115,7 @@ impl FieldHasher<BabyBear> for BabyBearPoseidon2Outer {
         [state[0]; BN254_DIGEST_SIZE]
     }
 }
+
 impl<C: CircuitConfig<F = BabyBear, N = Bn254Fr, Bit = Var<Bn254Fr>>> FieldHasherVariable<C>
     for BabyBearPoseidon2Outer
 {

crates/recursion/circuit-v2/src/machine/mod.rs

@@ -1,12 +1,14 @@
 mod compress;
 mod core;
 mod deferred;
+mod vkey_proof;
 mod witness;
 
 #[allow(unused_imports)]
 pub use compress::*;
 pub use core::*;
 pub use deferred::*;
+pub use vkey_proof::*;
 
 #[allow(unused_imports)]
 pub use witness::*;

crates/recursion/circuit-v2/src/machine/vkey_proof.rs

@@ -0,0 +1,163 @@
+use std::marker::PhantomData;
+
+use itertools::Itertools;
+use p3_air::Air;
+use p3_baby_bear::BabyBear;
+use p3_commit::Mmcs;
+use p3_field::AbstractField;
+use p3_matrix::dense::RowMajorMatrix;
+use sp1_recursion_compiler::ir::{Builder, Felt};
+use sp1_recursion_core_v2::DIGEST_SIZE;
+use sp1_stark::{
+    air::MachineAir, Com, InnerChallenge, OpeningProof, StarkGenericConfig, StarkMachine,
+};
+
+use crate::{
+    challenger::DuplexChallengerVariable,
+    constraints::RecursiveVerifierConstraintFolder,
+    hash::{FieldHasher, FieldHasherVariable},
+    merkle_tree::{verify, MerkleProof},
+    stark::MerkleProofVariable,
+    witness::{WitnessWriter, Witnessable},
+    BabyBearFriConfig, BabyBearFriConfigVariable, CircuitConfig, TwoAdicPcsProofVariable,
+};
+
+use super::{
+    SP1CompressShape, SP1CompressVerifier, SP1CompressWitnessValues, SP1CompressWitnessVariable,
+};
+
+/// A program to verify a batch of recursive proofs and aggregate their public values.
+#[derive(Debug, Clone, Copy)]
+pub struct SP1MerkleProofVerifier<C, SC> {
+    _phantom: PhantomData<(C, SC)>,
+}
+
+/// Witness layout for the compress stage verifier.
+pub struct SP1MerkleProofWitnessVariable<
+    C: CircuitConfig<F = BabyBear>,
+    SC: FieldHasherVariable<C> + BabyBearFriConfigVariable<C>,
+> {
+    /// The shard proofs to verify.
+    pub vk_merkle_proofs: Vec<MerkleProofVariable<C, SC>>,
+    pub root: SC::DigestVariable,
+}
+
+/// An input layout for the reduce verifier.
+pub struct SP1MerkleProofWitnessValues<SC: FieldHasher<BabyBear>> {
+    pub vk_merkle_proofs: Vec<MerkleProof<BabyBear, SC>>,
+    pub root: SC::Digest,
+}
+
+impl<C, SC> SP1MerkleProofVerifier<C, SC>
+where
+    SC: BabyBearFriConfigVariable<C>,
+    C: CircuitConfig<F = SC::Val, EF = SC::Challenge>,
+{
+    /// Verify a batch of recursive proofs and aggregate their public values.
+    ///
+    /// The compression verifier can aggregate proofs of different kinds:
+    /// - Core proofs: proofs which are recursive proof of a batch of SP1 shard proofs. The
+    ///   implementation in this function assumes a fixed recursive verifier speicified by
+    ///   `recursive_vk`.
+    /// - Deferred proofs: proofs which are recursive proof of a batch of deferred proofs. The
+    ///   implementation in this function assumes a fixed deferred verification program specified by
+    ///   `deferred_vk`.
+    /// - Compress proofs: these are proofs which refer to a prove of this program. The key for it
+    ///   is part of public values will be propagated accross all levels of recursion and will be
+    ///   checked against itself as in [sp1_prover::Prover] or as in [super::SP1RootVerifier].
+    pub fn verify(
+        builder: &mut Builder<C>,
+        digests: Vec<SC::DigestVariable>,
+        input: SP1MerkleProofWitnessVariable<C, SC>,
+        // TODO: add vk correctness check.
+        // vk_root: SC::Digest,
+        // Inclusion proof for the compressed vk.
+        // vk_inclusion_proof: SP1MerkleProofWitnessVariable<C, SC>,
+    ) {
+        for (proof, value) in input.vk_merkle_proofs.into_iter().zip(digests) {
+            // SC::assert_digest_eq(builder, *root, SC::hash(builder, &proof.root));
+            verify(builder, proof, value, input.root);
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct SP1CompressWithVKeyVerifier<C, SC, A> {
+    _phantom: PhantomData<(C, SC, A)>,
+}
+
+/// Witness layout for the verifier of the proof shape phase of the compress stage.
+pub struct SP1CompressWithVKeyWitnessVariable<
+    C: CircuitConfig<F = BabyBear>,
+    SC: BabyBearFriConfigVariable<C>,
+> {
+    pub compress_var: SP1CompressWitnessVariable<C, SC>,
+    pub merkle_var: SP1MerkleProofWitnessVariable<C, SC>,
+}
+
+/// An input layout for the verifier of the proof shape phase of the compress stage.
+pub struct SP1CompressWithVKeyWitnessValues<SC: StarkGenericConfig + FieldHasher<BabyBear>> {
+    pub compress_val: SP1CompressWitnessValues<SC>,
+    pub merkle_val: SP1MerkleProofWitnessValues<SC>,
+}
+
+impl<C, SC, A> SP1CompressWithVKeyVerifier<C, SC, A>
+where
+    SC: BabyBearFriConfigVariable<
+        C,
+        FriChallengerVariable = DuplexChallengerVariable<C>,
+        DigestVariable = [Felt<BabyBear>; DIGEST_SIZE],
+    >,
+    C: CircuitConfig<F = SC::Val, EF = SC::Challenge, Bit = Felt<BabyBear>>,
+    <SC::ValMmcs as Mmcs<BabyBear>>::ProverData<RowMajorMatrix<BabyBear>>: Clone,
+    A: MachineAir<SC::Val> + for<'a> Air<RecursiveVerifierConstraintFolder<'a, C>>,
+{
+    /// Verify the proof shape phase of the compress stage.
+    pub fn verify(
+        builder: &mut Builder<C>,
+        machine: &StarkMachine<SC, A>,
+        input: SP1CompressWithVKeyWitnessVariable<C, SC>,
+    ) {
+        // These are dummy values.
+        let values = input
+            .compress_var
+            .vks_and_proofs
+            .iter()
+            .map(|(_, _)| [builder.constant(BabyBear::zero()); DIGEST_SIZE])
+            .collect_vec();
+        // TODO: comment back in.
+        //     input.compress_var.vks_and_proofs.iter().map(|(vk, _)| vk.hash(builder)).collect_vec();
+        SP1MerkleProofVerifier::verify(builder, values, input.merkle_var);
+        SP1CompressVerifier::verify(builder, machine, input.compress_var);
+    }
+}
+
+impl<SC: BabyBearFriConfig + FieldHasher<BabyBear>> SP1CompressWithVKeyWitnessValues<SC> {
+    pub fn shape(&self) -> SP1CompressShape {
+        self.compress_val.shape()
+    }
+}
+
+impl<C: CircuitConfig<F = BabyBear, EF = InnerChallenge>, SC: BabyBearFriConfigVariable<C>>
+    Witnessable<C> for SP1CompressWithVKeyWitnessValues<SC>
+where
+    Com<SC>: Witnessable<C, WitnessVariable = <SC as FieldHasherVariable<C>>::DigestVariable>,
+    // As for `SP1MerkleProofWitnessValues`, this is a redundant trait bound.
+    SC: FieldHasher<BabyBear>,
+    <SC as FieldHasher<BabyBear>>::Digest: Witnessable<C, WitnessVariable = SC::DigestVariable>,
+    OpeningProof<SC>: Witnessable<C, WitnessVariable = TwoAdicPcsProofVariable<C, SC>>,
+{
+    type WitnessVariable = SP1CompressWithVKeyWitnessVariable<C, SC>;
+
+    fn read(&self, builder: &mut Builder<C>) -> Self::WitnessVariable {
+        SP1CompressWithVKeyWitnessVariable {
+            compress_var: self.compress_val.read(builder),
+            merkle_var: self.merkle_val.read(builder),
+        }
+    }
+
+    fn write(&self, witness: &mut impl WitnessWriter<C>) {
+        self.compress_val.write(witness);
+        self.merkle_val.write(witness);
+    }
+}