feat(zk-backend): two-stage groth16 dry-run

crates/proof/zk/backend/src/lib.rs

@@ -2,9 +2,10 @@
 
 mod succinct;
 pub use succinct::{
-    ClusterSessionId, ClusterZkProver, ClusterZkProverConfig, DRY_RUN_PREFIX, DryRunZkProver,
-    L1HeadSource, MOCK_PROOF_BYTES, MOCK_SNARK_PREFIX, MockZkProver, NetworkZkProver,
-    NetworkZkProverConfig, OpSuccinctWitnessProvider, SuccinctClusterBackendConfig,
-    SuccinctNetworkBackendConfig, SuccinctRpcConfig, SuccinctZkBackendConfig,
-    SuccinctZkProverBuildError, SuccinctZkProverBuilder, WitnessError, WitnessParams,
+    ClusterSessionId, ClusterZkProver, ClusterZkProverConfig, DRY_RUN_SNARK_PREFIX,
+    DRY_RUN_STARK_PREFIX, DryRunZkProver, L1HeadSource, MOCK_PROOF_BYTES, MOCK_SNARK_PREFIX,
+    MockZkProver, NetworkZkProver, NetworkZkProverConfig, OpSuccinctWitnessProvider,
+    SuccinctClusterBackendConfig, SuccinctNetworkBackendConfig, SuccinctRpcConfig,
+    SuccinctZkBackendConfig, SuccinctZkProverBuildError, SuccinctZkProverBuilder, WitnessError,
+    WitnessParams,
 };

crates/proof/zk/backend/src/succinct/dry_run.rs

@@ -9,7 +9,10 @@ use async_trait::async_trait;
 use base_proof_succinct_client_utils::client::DEFAULT_INTERMEDIATE_ROOT_INTERVAL;
 use base_proof_succinct_proof_utils::get_range_elf_embedded;
 use base_proof_zk_host::{ZkProofRequestKind, ZkProver, ZkProverError, ZkSessionState};
-use base_prover_service_protocol::{ExecutionStats, ProofResult, SessionType, ZkProofResult, ZkVm};
+use base_prover_service_protocol::{
+    ExecutionStats, ProofResult, SessionType, SnarkGroth16ProofRequest, SnarkGroth16ProofResult,
+    ZkProofRequest, ZkProofResult, ZkVm,
+};
 use sp1_sdk::{
     Elf, SP1Stdin,
     blocking::{LightProver, Prover},
@@ -28,8 +31,11 @@ macro_rules! backend_error {
     };
 }
 
-/// Prefix marking a dry-run backend session id.
-pub const DRY_RUN_PREFIX: &str = "dry-run-stark-";
+/// Prefix marking a dry-run STARK (range) backend session id.
+pub const DRY_RUN_STARK_PREFIX: &str = "dry-run-stark-";
+
+/// Prefix marking a dry-run SNARK (aggregation) backend session id.
+pub const DRY_RUN_SNARK_PREFIX: &str = "dry-run-snark-";
 
 /// [`ZkProver`] that generates a witness and executes the range program locally.
 #[derive(Clone, Debug)]
@@ -130,7 +136,7 @@ impl DryRunZkProver {
     /// Generate the witness, execute it locally, and return an empty-proof dry-run result.
     pub async fn prove_range(
         &self,
-        request: &base_prover_service_protocol::ZkProofRequest,
+        request: &ZkProofRequest,
         request_session_id: &str,
     ) -> Result<ProofResult, ZkProverError> {
         if request.number_of_blocks_to_prove == 0 {
@@ -206,6 +212,60 @@ impl DryRunZkProver {
             execution_stats: Some(execution_stats),
         }))
     }
+
+    /// Derive the deterministic backend session id for a request. Every request first proves a
+    /// range (STARK) proof, so the initial stage is always [`SessionType::Stark`].
+    pub fn backend_session_id(request_session_id: &str) -> String {
+        Self::backend_session_id_for_type(SessionType::Stark, request_session_id)
+    }
+
+    /// Derive the deterministic backend session id for a stage.
+    pub fn backend_session_id_for_type(
+        session_type: SessionType,
+        request_session_id: &str,
+    ) -> String {
+        match session_type {
+            SessionType::Stark => format!("{DRY_RUN_STARK_PREFIX}{request_session_id}"),
+            SessionType::Snark => format!("{DRY_RUN_SNARK_PREFIX}{request_session_id}"),
+        }
+    }
+
+    /// Atomically swap a completed range result for its aggregation (SNARK) result under
+    /// `backend_session_id`, returning that id.
+    ///
+    /// The lock is held across the whole check-remove-insert so the swap (drop the range result,
+    /// store the snark result) is atomic and no concurrent caller observes an intermediate state
+    /// where neither key is present. Idempotent: a crash-resume may re-invoke this with the same
+    /// pair, so if the snark result was already produced it is returned without re-processing.
+    pub fn advance_range_to_snark(
+        completed_results: &Mutex<HashMap<String, ProofResult>>,
+        backend_session_id: String,
+        completed_backend_session_id: &str,
+    ) -> Result<String, ZkProverError> {
+        let mut store = completed_results
+            .lock()
+            .map_err(|e| backend_error!("dry-run result store lock poisoned: {e}"))?;
+
+        if store.contains_key(&backend_session_id) {
+            return Ok(backend_session_id);
+        }
+
+        let range_result = store.remove(completed_backend_session_id).ok_or_else(|| {
+            ZkProverError::BackendSessionNotFound {
+                backend_session_id: completed_backend_session_id.to_owned(),
+            }
+        })?;
+        let execution_stats = match range_result {
+            ProofResult::Compressed(proof) => proof.execution_stats,
+            _ => return Err(backend_error!("dry-run range result had unexpected proof type")),
+        };
+        let result = ProofResult::SnarkGroth16(SnarkGroth16ProofResult {
+            proof: ZkProofResult { zk_vm: ZkVm::Sp1, proof: Vec::new().into(), execution_stats },
+        });
+        store.insert(backend_session_id.clone(), result);
+
+        Ok(backend_session_id)
+    }
 }
 
 /// Dry-run submission is synchronous: `submit` generates the witness and runs local SP1 execution
@@ -220,14 +280,13 @@ impl ZkProver for DryRunZkProver {
         request: &ZkProofRequestKind,
         request_session_id: &str,
     ) -> Result<String, ZkProverError> {
-        let ZkProofRequestKind::Compressed(request) = request else {
-            return Err(backend_error!(
-                "dry-run backend only supports compressed proof types; SNARK_GROTH16 requires a proof-producing backend"
-            ));
+        let range_request = match request {
+            ZkProofRequestKind::Compressed(request) => request,
+            ZkProofRequestKind::SnarkGroth16(request) => &request.proof,
         };
 
-        let backend_session_id = format!("{DRY_RUN_PREFIX}{request_session_id}");
-        let result = self.prove_range(request, request_session_id).await?;
+        let backend_session_id = Self::backend_session_id(request_session_id);
+        let result = self.prove_range(range_request, request_session_id).await?;
         self.completed_results
             .lock()
             .map_err(|e| backend_error!("dry-run result store lock poisoned: {e}"))?
@@ -236,6 +295,21 @@ impl ZkProver for DryRunZkProver {
         Ok(backend_session_id)
     }
 
+    async fn submit_next(
+        &self,
+        _request: &SnarkGroth16ProofRequest,
+        request_session_id: &str,
+        completed_backend_session_id: &str,
+    ) -> Result<String, ZkProverError> {
+        let backend_session_id =
+            Self::backend_session_id_for_type(SessionType::Snark, request_session_id);
+        Self::advance_range_to_snark(
+            &self.completed_results,
+            backend_session_id,
+            completed_backend_session_id,
+        )
+    }
+
     async fn poll(&self, backend_session_id: &str) -> Result<ZkSessionState, ZkProverError> {
         let contains_result = self
             .completed_results
@@ -259,3 +333,72 @@ impl ZkProver for DryRunZkProver {
             })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::{collections::HashMap, sync::Mutex};
+
+    use base_proof_zk_host::ZkProverError;
+    use base_prover_service_protocol::{ProofResult, SnarkGroth16ProofResult, ZkProofResult, ZkVm};
+
+    use super::{DRY_RUN_SNARK_PREFIX, DRY_RUN_STARK_PREFIX, DryRunZkProver};
+
+    fn stark_key(session: &str) -> String {
+        format!("{DRY_RUN_STARK_PREFIX}{session}")
+    }
+
+    fn snark_key(session: &str) -> String {
+        format!("{DRY_RUN_SNARK_PREFIX}{session}")
+    }
+
+    fn empty_proof() -> ZkProofResult {
+        ZkProofResult { zk_vm: ZkVm::Sp1, proof: Vec::new().into(), execution_stats: None }
+    }
+
+    #[test]
+    fn advance_range_to_snark_swaps_range_result_for_snark() {
+        let session = "session-1";
+        let store = Mutex::new(HashMap::from([(
+            stark_key(session),
+            ProofResult::Compressed(empty_proof()),
+        )]));
+
+        let snark_id =
+            DryRunZkProver::advance_range_to_snark(&store, snark_key(session), &stark_key(session))
+                .unwrap();
+
+        assert_eq!(snark_id, snark_key(session));
+        let guard = store.lock().unwrap();
+        assert!(!guard.contains_key(&stark_key(session)), "range result should be consumed");
+        assert!(matches!(guard.get(&snark_key(session)), Some(ProofResult::SnarkGroth16(_))));
+    }
+
+    #[test]
+    fn advance_range_to_snark_is_idempotent_on_resume() {
+        let session = "session-1";
+        // A previous call already produced the snark result; the range entry is gone.
+        let store = Mutex::new(HashMap::from([(
+            snark_key(session),
+            ProofResult::SnarkGroth16(SnarkGroth16ProofResult { proof: empty_proof() }),
+        )]));
+
+        let snark_id =
+            DryRunZkProver::advance_range_to_snark(&store, snark_key(session), &stark_key(session))
+                .unwrap();
+
+        assert_eq!(snark_id, snark_key(session));
+        assert_eq!(store.lock().unwrap().len(), 1, "store should be unchanged on resume");
+    }
+
+    #[test]
+    fn advance_range_to_snark_errors_when_range_result_missing() {
+        let session = "session-1";
+        let store = Mutex::new(HashMap::<String, ProofResult>::new());
+
+        let err =
+            DryRunZkProver::advance_range_to_snark(&store, snark_key(session), &stark_key(session))
+                .unwrap_err();
+
+        assert!(matches!(err, ZkProverError::BackendSessionNotFound { .. }));
+    }
+}

crates/proof/zk/backend/src/succinct/mod.rs

@@ -20,7 +20,7 @@ mod network;
 pub use network::{NetworkZkProver, NetworkZkProverConfig, SuccinctNetworkBackendConfig};
 
 mod dry_run;
-pub use dry_run::{DRY_RUN_PREFIX, DryRunZkProver};
+pub use dry_run::{DRY_RUN_SNARK_PREFIX, DRY_RUN_STARK_PREFIX, DryRunZkProver};
 
 mod mock;
 pub use mock::{MOCK_PROOF_BYTES, MOCK_SNARK_PREFIX, MockZkProver};