feat(provers): add zkvm support for batch requests

Author: keroro520

provers/risc0/builder/src/main.rs

@@ -6,7 +6,7 @@ use std::path::PathBuf;
 fn main() {
     let pipeline = Risc0Pipeline::new("provers/risc0/guest", "release");
     pipeline.bins(
-        &["risc0-guest", "risc0-aggregation"],
+        &["risc0-guest", "risc0-aggregation", "risc0-batch"],
         "provers/risc0/driver/src/methods",
     );
     #[cfg(feature = "test")]

provers/risc0/driver/src/bonsai.rs

@@ -1,5 +1,4 @@
 use crate::{
-    methods::risc0_guest::RISC0_GUEST_ID,
     snarks::{stark2snark, verify_groth16_from_snark_receipt},
     Risc0Response,
 };
@@ -313,8 +312,10 @@ pub async fn bonsai_stark_to_snark(
     stark_uuid: String,
     stark_receipt: Receipt,
     input: B256,
+    elf: &[u8],
 ) -> ProverResult<Risc0Response> {
-    let image_id = Digest::from(RISC0_GUEST_ID);
+    let image_id = risc0_zkvm::compute_image_id(elf)
+        .map_err(|e| ProverError::GuestError(format!("Failed to compute image id: {e:?}")))?;
     let (snark_uuid, snark_receipt) = stark2snark(
         image_id,
         stark_uuid.clone(),

provers/risc0/driver/src/lib.rs

@@ -3,8 +3,8 @@
 #[cfg(feature = "bonsai-auto-scaling")]
 use crate::bonsai::auto_scaling::shutdown_bonsai;
 use crate::{
-    methods::risc0_aggregation::RISC0_AGGREGATION_ELF,
-    methods::risc0_guest::{RISC0_GUEST_ELF, RISC0_GUEST_ID},
+    methods::risc0_aggregation::RISC0_AGGREGATION_ELF, methods::risc0_batch::RISC0_BATCH_ELF,
+    methods::risc0_guest::RISC0_GUEST_ELF,
 };
 use alloy_primitives::{hex::ToHexExt, B256};
 use bonsai::{cancel_proof, maybe_prove};
@@ -18,8 +18,10 @@ use raiko_lib::{
     prover::{IdStore, IdWrite, Proof, ProofKey, Prover, ProverConfig, ProverError, ProverResult},
 };
 use risc0_zkvm::{
-    compute_image_id, default_prover, serde::to_vec, sha::Digestible, ExecutorEnv, ProverOpts,
-    Receipt,
+    compute_image_id, default_prover,
+    serde::to_vec,
+    sha::{Digest, Digestible},
+    ExecutorEnv, ProverOpts, Receipt,
 };
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
@@ -93,7 +95,7 @@ impl Prover for Risc0Prover {
         .await?;
 
         let proof_gen_result = if config.snark && config.bonsai {
-            bonsai::bonsai_stark_to_snark(uuid, receipt, output.hash)
+            bonsai::bonsai_stark_to_snark(uuid, receipt, output.hash, RISC0_GUEST_ELF)
                 .await
                 .map(|r0_response| r0_response.into())
                 .map_err(|e| ProverError::GuestError(e.to_string()))
@@ -149,11 +151,16 @@ impl Prover for Risc0Prover {
             .iter()
             .map(|proof| proof.input.unwrap())
             .collect::<Vec<_>>();
+
+        let input_proof_hex_str = input.proofs[0].proof.as_ref().unwrap();
+        let input_proof_bytes = hex::decode(&input_proof_hex_str[2..]).unwrap();
+        let input_image_id_bytes: [u8; 32] = input_proof_bytes[32..64].try_into().unwrap();
+        let input_proof_image_id = Digest::from(input_image_id_bytes);
         let input = ZkAggregationGuestInput {
-            image_id: RISC0_GUEST_ID,
+            image_id: input_proof_image_id.as_words().try_into().unwrap(),
             block_inputs,
         };
-        info!("Start aggregate proofs");
+
         // add_assumption makes the receipt to be verified available to the prover.
         let env = {
             let mut env = ExecutorEnv::builder();
@@ -173,10 +180,9 @@ impl Prover for Risc0Prover {
             "Generate aggregation receipt journal: {:?}",
             alloy_primitives::hex::encode_prefixed(receipt.journal.bytes.clone())
         );
-        let block_proof_image_id = compute_image_id(RISC0_GUEST_ELF).unwrap();
         let aggregation_image_id = compute_image_id(RISC0_AGGREGATION_ELF).unwrap();
         let proof_data = snarks::verify_aggregation_groth16_proof(
-            block_proof_image_id,
+            input_proof_image_id,
             aggregation_image_id,
             receipt.clone(),
         )
@@ -223,12 +229,60 @@ impl Prover for Risc0Prover {
 
     async fn batch_run(
         &self,
-        _input: GuestBatchInput,
-        _output: &GuestBatchOutput,
-        _config: &ProverConfig,
-        _store: Option<&mut dyn IdWrite>,
+        input: GuestBatchInput,
+        output: &GuestBatchOutput,
+        config: &ProverConfig,
+        id_store: Option<&mut dyn IdWrite>,
     ) -> ProverResult<Proof> {
-        unimplemented!();
+        let mut id_store = id_store;
+        let config = Risc0Param::deserialize(config.get("risc0").unwrap()).unwrap();
+        let proof_key = (
+            input.taiko.chain_spec.chain_id,
+            input.taiko.batch_id,
+            output.hash,
+            ProofType::Risc0 as u8,
+        );
+
+        let encoded_input = to_vec(&input).expect("Could not serialize proving input!");
+
+        let (uuid, receipt) = maybe_prove::<GuestBatchInput, B256>(
+            &config,
+            encoded_input,
+            RISC0_BATCH_ELF,
+            &output.hash,
+            (Vec::<Receipt>::new(), Vec::new()),
+            proof_key,
+            &mut id_store,
+        )
+        .await?;
+
+        let proof_gen_result = if config.snark && config.bonsai {
+            bonsai::bonsai_stark_to_snark(uuid, receipt, output.hash, RISC0_BATCH_ELF)
+                .await
+                .map(|r0_response| r0_response.into())
+                .map_err(|e| ProverError::GuestError(e.to_string()))
+        } else {
+            if !config.snark {
+                warn!("proof is not in snark mode, please check.");
+            }
+            Ok(Risc0Response {
+                proof: receipt.journal.encode_hex_with_prefix(),
+                receipt: serde_json::to_string(&receipt).unwrap(),
+                uuid,
+                input: output.hash,
+            }
+            .into())
+        };
+
+        #[cfg(feature = "bonsai-auto-scaling")]
+        if config.bonsai {
+            // shutdown bonsai
+            shutdown_bonsai()
+                .await
+                .map_err(|e| ProverError::GuestError(e.to_string()))?;
+        }
+
+        proof_gen_result
     }
 }
 

provers/risc0/guest/Cargo.toml

@@ -13,6 +13,10 @@ path = "src/zk_op.rs"
 name = "risc0-aggregation"
 path = "src/aggregation.rs"
 
+[[bin]]
+name = "risc0-batch"
+path = "src/batch.rs"
+
 [[bin]]
 name = "sha256"
 path = "src/benchmark/sha256.rs"

provers/risc0/guest/src/batch.rs

@@ -0,0 +1,39 @@
+#![no_main]
+harness::entrypoint!(main, tests, zk_op::tests);
+use raiko_lib::{
+    builder::calculate_batch_blocks_final_header, input::GuestBatchInput, proof_type::ProofType,
+    protocol_instance::ProtocolInstance,
+};
+use revm_precompile::zk_op::ZkOperation;
+use risc0_zkvm::guest::env;
+use zk_op::Risc0Operator;
+
+pub mod mem;
+
+pub use mem::*;
+
+fn main() {
+    let batch_input: GuestBatchInput = env::read();
+
+    revm_precompile::zk_op::ZKVM_OPERATOR.get_or_init(|| Box::new(Risc0Operator {}));
+    revm_precompile::zk_op::ZKVM_OPERATIONS
+        .set(Box::new(vec![ZkOperation::Sha256, ZkOperation::Secp256k1]))
+        .expect("Failed to set ZkvmOperations");
+
+    let final_blocks = calculate_batch_blocks_final_header(&batch_input);
+    let pi = ProtocolInstance::new_batch(&batch_input, final_blocks, ProofType::Risc0)
+        .unwrap()
+        .instance_hash();
+
+    env::commit(&pi);
+}
+
+harness::zk_suits!(
+    pub mod tests {
+        #[test]
+        pub fn test_build_from_mock_input() {
+            // Todo: impl mock input for static unit test
+            assert_eq!(1, 1);
+        }
+    }
+);

provers/sp1/builder/src/main.rs

@@ -5,7 +5,10 @@ use std::path::PathBuf;
 
 fn main() {
     let pipeline = Sp1Pipeline::new("provers/sp1/guest", "release");
-    pipeline.bins(&["sp1-guest", "sp1-aggregation"], "provers/sp1/guest/elf");
+    pipeline.bins(
+        &["sp1-guest", "sp1-aggregation", "sp1-batch"],
+        "provers/sp1/guest/elf",
+    );
     #[cfg(feature = "test")]
     pipeline.tests(&["sp1-guest"], "provers/sp1/guest/elf");
     #[cfg(feature = "bench")]