Introduce OuterZkvmHost trait and wire SP1 aggregation into parallel prover

crates/adapters/mock-zkvm/src/host.rs

@@ -1,6 +1,8 @@
 use crate::notifier::NotificationManager;
 use crate::{MockCodeCommitment, MockProof, MockZkGuest};
 use serde::Serialize;
+use sov_rollup_interface::da::DaSpec;
+use sov_rollup_interface::zk::aggregated_proof::{BlockProof, OuterZkvmHost};
 
 /// A mock implementing the zkVM trait.
 #[derive(Clone)]
@@ -41,6 +43,17 @@ impl MockZkvmHost {
         })
         .unwrap()
     }
+
+    fn add_hint_and_run_inner<T: Serialize>(&self, item: &T) -> anyhow::Result<Vec<u8>> {
+        let pub_data = bincode::serialize(item)?;
+        if self.wait_for_proof {
+            self.notification_manager.wait();
+        }
+        Ok(bincode::serialize(&MockProof {
+            is_valid: true,
+            pub_data,
+        })?)
+    }
 }
 
 impl Default for MockZkvmHost {
@@ -59,17 +72,32 @@ impl sov_rollup_interface::zk::ZkvmHost for MockZkvmHost {
     }
 
     fn add_hint_and_run<T: Serialize>(&mut self, item: &T) -> anyhow::Result<Vec<u8>> {
-        let pub_data = bincode::serialize(item)?;
-        if self.wait_for_proof {
-            self.notification_manager.wait();
-        }
-        Ok(bincode::serialize(&MockProof {
-            is_valid: true,
-            pub_data,
-        })?)
+        self.add_hint_and_run_inner(item)
     }
 
     fn from_args(_args: &Self::HostArgs) -> Self {
         Self::default()
     }
 }
+
+impl OuterZkvmHost for MockZkvmHost {
+    fn run_proof_aggregation<Address: Serialize + Clone, Da: DaSpec, Root: Serialize + Clone>(
+        &self,
+        genesis_state_root: Root,
+        headers_with_block_proofs: Vec<(Da::BlockHeader, BlockProof<Address, Da, Root>)>,
+    ) -> anyhow::Result<Vec<u8>> {
+        use sov_rollup_interface::zk::aggregated_proof::AggregatedProofPublicData;
+
+        let block_proofs_data = headers_with_block_proofs
+            .iter()
+            .map(|(_, bp)| bp)
+            .collect::<Vec<_>>();
+
+        let public_data = AggregatedProofPublicData::from_block_proofs(
+            block_proofs_data.as_slice(),
+            genesis_state_root,
+        );
+
+        self.add_hint_and_run_inner(&public_data)
+    }
+}

crates/adapters/mock-zkvm/src/lib.rs

@@ -59,6 +59,9 @@ impl Zkvm for MockZkvm {
     #[cfg(feature = "native")]
     type Host = crate::host::MockZkvmHost;
 
+    #[cfg(feature = "native")]
+    type OuterHost = crate::host::MockZkvmHost;
+
     #[cfg(feature = "native")]
     type Network = crate::network::MockZkvmNetwork;
 }

crates/adapters/risc0/src/host.rs

@@ -1,10 +1,13 @@
 //! This module implements the [`ZkvmHost`] trait for the RISC0 VM.
 
-use risc0_zkvm::{ExecutorEnvBuilder, ExecutorImpl, Session};
-use sov_rollup_interface::zk::ZkvmHost;
-
 use crate::guest::Risc0Guest;
 use crate::Risc0MethodId;
+use risc0_zkvm::{ExecutorEnvBuilder, ExecutorImpl, Session};
+use serde::Serialize;
+use sov_rollup_interface::da::DaSpec;
+use sov_rollup_interface::zk::aggregated_proof::BlockProof;
+use sov_rollup_interface::zk::aggregated_proof::OuterZkvmHost;
+use sov_rollup_interface::zk::ZkvmHost;
 
 /// A [`Risc0Host`] stores a binary to execute in the Risc0 VM, and accumulates hints to be
 /// provided to its execution.
@@ -109,3 +112,13 @@ impl ZkvmHost for Risc0Host<'static> {
         ))
     }
 }
+
+impl OuterZkvmHost for Risc0Host<'static> {
+    fn run_proof_aggregation<Address: Serialize + Clone, Da: DaSpec, Root: Serialize + Clone>(
+        &self,
+        _genesis_state_root: Root,
+        _headers_with_block_proofs: Vec<(Da::BlockHeader, BlockProof<Address, Da, Root>)>,
+    ) -> anyhow::Result<Vec<u8>> {
+        unimplemented!("Proof aggregation not supported for Risc0")
+    }
+}

crates/adapters/risc0/src/lib.rs

@@ -118,6 +118,9 @@ impl sov_rollup_interface::zk::Zkvm for Risc0 {
     #[cfg(feature = "native")]
     type Host = crate::host::Risc0Host<'static>;
 
+    #[cfg(feature = "native")]
+    type OuterHost = crate::host::Risc0Host<'static>;
+
     #[cfg(feature = "native")]
     type Network = sov_rollup_interface::zk::NoopZkvmNetwork<crate::guest::Risc0Guest>;
 }

crates/adapters/sp1/src/host.rs

@@ -1,6 +1,6 @@
 //! Implementation of the SP1 host for the Sovereign ZkvmHost trait.
 
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use crate::guest::SP1Guest;
 use crate::SP1MethodId;
@@ -10,7 +10,9 @@ use sov_rollup_interface::reexports::anyhow;
 use sov_rollup_interface::zk::aggregated_proof::common::{
     AggregatedProofWitness, DeferredProofInput, PreviousOuterProofWitness,
 };
-use sov_rollup_interface::zk::aggregated_proof::{BlockHeaderWithProof, CodeCommitmentHash};
+use sov_rollup_interface::zk::aggregated_proof::{
+    BlockHeaderWithProof, BlockProof, CodeCommitmentHash, OuterZkvmHost,
+};
 use sov_rollup_interface::zk::ZkvmHost;
 use sp1_sdk::blocking::ProveRequest;
 use sp1_sdk::blocking::{CpuProver, MockProver, Prover, ProverClient};
@@ -24,12 +26,17 @@ use sp1_sdk::{HashableKey, SP1Proof, SP1ProvingKey, SP1Stdin};
 /// that covers one batch of inner proofs.  When called multiple times the host
 /// automatically chains proofs: the previous aggregation proof is fed back as
 /// a deferred proof input so the guest can verify continuity.
+#[derive(Clone)]
 pub struct SP1AggregationHost {
+    inner: Arc<Inner>,
+}
+
+struct Inner {
     host: SP1Host,
     aggregation_vk: sp1_sdk::SP1VerifyingKey,
     code_commitment: SP1MethodId,
     inner_method_id: SP1MethodId,
-    prev_agg_proof: Option<sp1_sdk::SP1ProofWithPublicValues>,
+    prev_agg_proof: Mutex<Option<sp1_sdk::SP1ProofWithPublicValues>>,
 }
 
 impl SP1AggregationHost {
@@ -40,24 +47,26 @@ impl SP1AggregationHost {
         let aggregation_vk = host.pk.verifying_key().clone();
         let code_commitment = SP1MethodId(bincode::serialize(&aggregation_vk)?);
         Ok(Self {
-            host,
-            aggregation_vk,
-            code_commitment,
-            inner_method_id,
-            prev_agg_proof: None,
+            inner: Arc::new(Inner {
+                host,
+                aggregation_vk,
+                code_commitment,
+                inner_method_id,
+                prev_agg_proof: Mutex::new(None),
+            }),
         })
     }
 
     /// Returns the code commitment (verifying key) of the aggregation program.
     pub fn code_commitment(&self) -> SP1MethodId {
-        self.code_commitment.clone()
+        self.inner.code_commitment.clone()
     }
 
     /// Generates a compressed aggregation proof over the supplied inner
     /// `proofs_and_headers`.  If a previous aggregation proof exists it is
     /// included as a deferred proof input for recursive verification.
     pub fn run<Da: DaSpec>(
-        &mut self,
+        &self,
         proofs_and_headers: Vec<BlockHeaderWithProof<Da>>,
     ) -> anyhow::Result<Vec<u8>> {
         anyhow::ensure!(
@@ -67,24 +76,31 @@ impl SP1AggregationHost {
 
         let mut stdin = SP1Stdin::new();
 
-        let prev_outer_proof_witness =
-            if let Some(previous_outer_proof) = self.prev_agg_proof.as_ref() {
-                let serialized = bincode::serialize(previous_outer_proof)?;
-                let public_values =
-                    self.host
-                        .add_proof_helper(&mut stdin, &serialized, &self.code_commitment)?;
+        let mut prev_agg_proof = self
+            .inner
+            .prev_agg_proof
+            .lock()
+            .map_err(|e| anyhow::anyhow!("prev_agg_proof mutex poisoned: {e}"))?;
 
-                Some(PreviousOuterProofWitness { public_values })
-            } else {
-                None
-            };
+        let prev_outer_proof_witness = if let Some(previous_outer_proof) = prev_agg_proof.as_ref() {
+            let serialized = bincode::serialize(previous_outer_proof)?;
+            let public_values = self.inner.host.add_proof_helper(
+                &mut stdin,
+                &serialized,
+                &self.inner.code_commitment,
+            )?;
+
+            Some(PreviousOuterProofWitness { public_values })
+        } else {
+            None
+        };
 
         let mut proof_inputs = Vec::with_capacity(proofs_and_headers.len());
         for proof_and_header in proofs_and_headers {
-            let public_values = self.host.add_proof_helper(
+            let public_values = self.inner.host.add_proof_helper(
                 &mut stdin,
                 &proof_and_header.proof.raw_inner_proof,
-                &self.inner_method_id,
+                &self.inner.inner_method_id,
             )?;
             let proof_input = DeferredProofInput::<Da> {
                 public_values,
@@ -94,9 +110,10 @@ impl SP1AggregationHost {
             proof_inputs.push(proof_input);
         }
 
-        let aggregation_vk_hash = self.aggregation_vk.hash_u32();
-        let inner_vk: sp1_sdk::SP1VerifyingKey = bincode::deserialize(&self.inner_method_id.0)
-            .map_err(|e| anyhow::anyhow!("Failed to deserialize inner SP1VerifyingKey: {e}"))?;
+        let aggregation_vk_hash = self.inner.aggregation_vk.hash_u32();
+        let inner_vk: sp1_sdk::SP1VerifyingKey =
+            bincode::deserialize(&self.inner.inner_method_id.0)
+                .map_err(|e| anyhow::anyhow!("Failed to deserialize inner SP1VerifyingKey: {e}"))?;
         let inner_vkey_hash = CodeCommitmentHash::from_u32_array(inner_vk.hash_u32());
         let outer_vkey_hash = CodeCommitmentHash::from_u32_array(aggregation_vk_hash);
 
@@ -109,9 +126,9 @@ impl SP1AggregationHost {
 
         stdin.write(&witness);
 
-        let agg_proof = self.host.run_helper(stdin)?;
+        let agg_proof = self.inner.host.run_helper(stdin)?;
         let serialized = bincode::serialize(&agg_proof)?;
-        self.prev_agg_proof = Some(agg_proof);
+        *prev_agg_proof = Some(agg_proof);
 
         Ok(serialized)
     }
@@ -257,3 +274,21 @@ impl MockSp1Prover {
             .map_err(|e| anyhow::anyhow!("SP1 verification failed. Error: {:?}", e))
     }
 }
+
+impl OuterZkvmHost for SP1AggregationHost {
+    fn run_proof_aggregation<Address: Serialize + Clone, Da: DaSpec, Root: Serialize + Clone>(
+        &self,
+        _genesis_state_root: Root,
+        headers_with_block_proofs: Vec<(Da::BlockHeader, BlockProof<Address, Da, Root>)>,
+    ) -> anyhow::Result<Vec<u8>> {
+        let proofs_and_headers: Vec<BlockHeaderWithProof<Da>> = headers_with_block_proofs
+            .into_iter()
+            .map(|(header, proof)| BlockHeaderWithProof {
+                da_block_header: header,
+                proof: proof.proof,
+            })
+            .collect();
+
+        self.run(proofs_and_headers)
+    }
+}

crates/adapters/sp1/src/lib.rs

@@ -113,6 +113,9 @@ impl sov_rollup_interface::zk::Zkvm for SP1 {
     #[cfg(feature = "native")]
     type Host = crate::host::SP1Host;
 
+    #[cfg(feature = "native")]
+    type OuterHost = crate::host::SP1AggregationHost;
+
     #[cfg(feature = "native")]
     type Network = crate::network::SP1Network;
 }

crates/full-node/sov-stf-runner/src/processes/prover_service/mod.rs

@@ -187,7 +187,7 @@ pub trait ProverService: Send + Sync + 'static {
     /// This method is not yet fully implemented: see #1185
     async fn create_aggregated_proof(
         &self,
-        block_header_hashes: &[<<Self::DaService as DaService>::Spec as DaSpec>::SlotHash],
+        block_headers: &[<<Self::DaService as DaService>::Spec as DaSpec>::BlockHeader],
         genesis_state_root: &Self::StateRoot,
     ) -> anyhow::Result<ProofAggregationStatus>;
 }

crates/full-node/sov-stf-runner/src/processes/prover_service/network/mod.rs

@@ -96,11 +96,11 @@ where
 
     async fn create_aggregated_proof(
         &self,
-        block_header_hashes: &[<<Self::DaService as DaService>::Spec as DaSpec>::SlotHash],
+        block_headers: &[<<Self::DaService as DaService>::Spec as DaSpec>::BlockHeader],
         genesis_state_root: &Self::StateRoot,
     ) -> anyhow::Result<ProofAggregationStatus> {
         self.prover
-            .create_aggregated_proof(block_header_hashes, genesis_state_root)
+            .create_aggregated_proof(block_headers, genesis_state_root)
             .await
     }
 }

crates/full-node/sov-stf-runner/src/processes/prover_service/network/prover.rs

@@ -8,8 +8,8 @@ use sov_rollup_interface::zk::aggregated_proof::{
     AggregatedProofPublicData, BlockProof, SerializedAggregatedProof,
 };
 use sov_rollup_interface::zk::{
-    StateTransitionPublicData, StateTransitionWitness, StateTransitionWitnessWithAddress, Zkvm,
-    ZkvmNetwork,
+    SerializedInnerProof, StateTransitionPublicData, StateTransitionWitness,
+    StateTransitionWitnessWithAddress, Zkvm, ZkvmNetwork,
 };
 use std::collections::HashMap;
 use std::marker::PhantomData;
@@ -178,15 +178,17 @@ where
 
     pub(crate) async fn create_aggregated_proof(
         &self,
-        block_header_hashes: &[<Da::Spec as DaSpec>::SlotHash],
+        block_headers: &[<Da::Spec as DaSpec>::BlockHeader],
         genesis_state_root: &StateRoot,
     ) -> anyhow::Result<ProofAggregationStatus> {
-        assert!(!block_header_hashes.is_empty());
+        assert!(!block_headers.is_empty());
+
+        let block_header_hashes: Vec<_> = block_headers.iter().map(|h| h.hash()).collect();
 
         let mut proof_statuses = self.tracker.write().await;
         // Phase 1: Poll all Submitted entries and transition them to Proved.
         // We only remove entries confirmed to be Submitted, so Proved/Err entries are untouched.
-        for slot_hash in block_header_hashes {
+        for slot_hash in &block_header_hashes {
             // We want to only remove if we know the entry exists
             if !matches!(
                 proof_statuses.get(slot_hash),
@@ -204,7 +206,9 @@ where
             match self.inner_vm.poll(&handle).await {
                 Ok(Some(proof_bytes)) => {
                     let block_proof = BlockProof {
-                        proof: proof_bytes,
+                        proof: SerializedInnerProof {
+                            raw_inner_proof: proof_bytes,
+                        },
                         slot_number: metadata.slot_number,
                         st: metadata.st,
                     };
@@ -240,7 +244,7 @@ where
         let proof_statuses = proof_statuses.downgrade();
         // Phase 2: Collect all proved block proofs.
         let mut block_proofs_data = Vec::new();
-        for slot_hash in block_header_hashes {
+        for slot_hash in &block_header_hashes {
             match proof_statuses.get(slot_hash) {
                 Some(NetworkProverStatus::Proved(block_proof)) => {
                     assert_eq!(slot_hash, &block_proof.st.slot_hash);
@@ -299,7 +303,7 @@ where
         })??;
 
         let mut tracker = self.tracker.write().await;
-        for slot_hash in block_header_hashes {
+        for slot_hash in &block_header_hashes {
             tracker.remove(slot_hash);
         }