feat: impl zkvm support for pacaya (#526)

* feat(provers): add zkvm support for batch requests

* chore(provers): update elf

* feat(host): enable zk_any for batch API

* chore: cargo fmt

* fix(host): remove zk_any opts from request

* chore(host): add logs

* fix(reqactor): return early for wip,success task

* feat(provers/sp1): hardcode sp1_proof_mode

* feat(host): fill batch_id into response of batch-request

* fix(host): remove fulfill_sp1_params

* chore(host/config): deprecate sp1_params.recursion

---------

Co-authored-by: smtmfft <[email protected]>

Author: keroro520

host/config/config.devnet.json

@@ -30,11 +30,10 @@
     "execution_po2": 20
   },
   "sp1": {
-    "recursion": "core",
     "prover": "network",
     "verify": false
   },
   "native": {
     "json_guest_input": null
   }
-}
+}

host/config/config.json

@@ -20,7 +20,6 @@
     "execution_po2": 20
   },
   "sp1": {
-    "recursion": "core",
     "prover": "network",
     "verify": false
   },

host/config/config.sgx.json

@@ -30,11 +30,10 @@
         "verify": true
     },
     "sp1": {
-        "recursion": "core",
         "prover": "network",
         "verify": true
     },
     "native": {
         "json_guest_input": null
     }
-}
+}

host/config/config.taiko_hekla.json

@@ -17,7 +17,6 @@
     "execution_po2": 20
   },
   "sp1": {
-    "recursion": "core",
     "prover": "network",
     "verify": true
   },
@@ -33,4 +32,4 @@
   "native": {
     "json_guest_input": null
   }
-}
+}

host/config/config.taiko_mainnet.json

@@ -17,7 +17,6 @@
     "execution_po2": 20
   },
   "sp1": {
-    "recursion": "core",
     "prover": "network",
     "verify": true
   },
@@ -33,4 +32,4 @@
   "native": {
     "json_guest_input": null
   }
-}
+}

host/src/server/api/v2/mod.rs

@@ -69,6 +69,8 @@ pub enum Status {
     Ok {
         #[serde(with = "raiko_lib::proof_type::lowercase")]
         proof_type: ProofType,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        batch_id: Option<u64>,
         data: ProofResponse,
     },
     Error {

host/src/server/api/v2/proof/mod.rs

@@ -9,7 +9,7 @@ use raiko_tasks::TaskStatus;
 use serde_json::Value;
 use utoipa::OpenApi;
 
-use crate::server::utils::{draw_for_zk_any_request, fulfill_sp1_params, is_zk_any_request};
+use crate::server::utils::{draw_for_zk_any_request, is_zk_any_request};
 use crate::{
     interfaces::HostResult,
     metrics::{inc_current_req, inc_guest_req_count, inc_host_req_count},
@@ -39,16 +39,9 @@ pub mod report;
 /// - sgx - uses the sgx environment to construct a block and produce proof of execution
 /// - sp1 - uses the sp1 prover
 /// - risc0 - uses the risc0 prover
-async fn proof_handler(
-    State(actor): State<Actor>,
-    Json(mut req): Json<Value>,
-) -> HostResult<Status> {
+async fn proof_handler(State(actor): State<Actor>, Json(req): Json<Value>) -> HostResult<Status> {
     inc_current_req();
 
-    if is_zk_any_request(&req) {
-        fulfill_sp1_params(&mut req);
-    }
-
     // Override the existing proof request config from the config file and command line
     // options with the request from the client.
     let mut config = actor.default_request_config().clone();
@@ -61,6 +54,7 @@ async fn proof_handler(
             None => {
                 return Ok(Status::Ok {
                     proof_type: ProofType::Native,
+                    batch_id: None,
                     data: ProofResponse::Status {
                         status: TaskStatus::ZKAnyNotDrawn,
                     },
@@ -126,10 +120,10 @@ async fn proof_handler(
             .into();
 
             let result = crate::server::prove(&actor, request_key, request_entity).await;
-            Ok(to_v2_status(proof_type, result))
+            Ok(to_v2_status(proof_type, None, result))
         }
-        Ok(_) => Ok(to_v2_status(proof_type, result)),
-        Err(_) => Ok(to_v2_status(proof_type, result)),
+        Ok(_) => Ok(to_v2_status(proof_type, None, result)),
+        Err(_) => Ok(to_v2_status(proof_type, None, result)),
     }
 }
 

host/src/server/api/v3/proof/aggregate/mod.rs

@@ -67,7 +67,7 @@ async fn aggregation_handler(
     .into();
 
     let result = crate::server::prove(&actor, agg_request_key, agg_request_entity).await;
-    Ok(to_v3_status(proof_type, result))
+    Ok(to_v3_status(proof_type, None, result))
 }
 
 #[derive(OpenApi)]

host/src/server/api/v3/proof/batch.rs

@@ -4,12 +4,12 @@ use crate::{
         api::v3::{ProofResponse, Status},
         handler::prove_many,
         prove_aggregation,
-        utils::{is_zk_any_request, to_v3_status},
+        utils::{draw_for_zk_any_batch_request, is_zk_any_request, to_v3_status},
     },
 };
 use axum::{extract::State, routing::post, Json, Router};
 use raiko_core::{
-    interfaces::{BatchMetadata, BatchProofRequest, BatchProofRequestOpt},
+    interfaces::{BatchMetadata, BatchProofRequest, BatchProofRequestOpt, RaikoError},
     merge,
 };
 use raiko_lib::{proof_type::ProofType, prover::Proof};
@@ -41,20 +41,46 @@ async fn batch_handler(
     State(actor): State<Actor>,
     Json(batch_request_opt): Json<Value>,
 ) -> HostResult<Status> {
-    if is_zk_any_request(&batch_request_opt) {
-        return Ok(Status::Ok {
-            proof_type: ProofType::Native,
-            data: ProofResponse::Status {
-                status: TaskStatus::ZKAnyNotDrawn,
-            },
-        });
-    }
+    tracing::debug!(
+        "Received batch request: {}",
+        serde_json::to_string(&batch_request_opt)?
+    );
 
     let batch_request = {
         // Override the existing proof request config from the config file and command line
         // options with the request from the client, and convert to a BatchProofRequest.
         let mut opts = serde_json::to_value(actor.default_request_config())?;
         merge(&mut opts, &batch_request_opt);
+
+        let first_batch_id = {
+            let batches = opts["batches"]
+                .as_array()
+                .ok_or(RaikoError::InvalidRequestConfig(
+                    "Missing batches".to_string(),
+                ))?;
+            let first_batch = batches.first().ok_or(RaikoError::InvalidRequestConfig(
+                "batches is empty".to_string(),
+            ))?;
+            let first_batch_id = first_batch["batch_id"].as_u64().expect("checked above");
+            first_batch_id
+        };
+
+        // For zk_any request, draw zk proof type based on the block hash.
+        if is_zk_any_request(&opts) {
+            match draw_for_zk_any_batch_request(&actor, &opts).await? {
+                Some(proof_type) => opts["proof_type"] = serde_json::to_value(proof_type).unwrap(),
+                None => {
+                    return Ok(Status::Ok {
+                        proof_type: ProofType::Native,
+                        batch_id: Some(first_batch_id),
+                        data: ProofResponse::Status {
+                            status: TaskStatus::ZKAnyNotDrawn,
+                        },
+                    });
+                }
+            }
+        }
+
         let batch_request_opt: BatchProofRequestOpt = serde_json::from_value(opts)?;
         let batch_request: BatchProofRequest = batch_request_opt.try_into()?;
 
@@ -188,7 +214,12 @@ async fn batch_handler(
                 })
         }
     };
-    Ok(to_v3_status(batch_request.proof_type, result))
+    tracing::debug!("Batch proof result: {}", serde_json::to_string(&result)?);
+    Ok(to_v3_status(
+        batch_request.proof_type,
+        Some(batch_request.batches.first().unwrap().batch_id),
+        result,
+    ))
 }
 
 #[derive(OpenApi)]

host/src/server/api/v3/proof/mod.rs

@@ -115,7 +115,7 @@ async fn proof_handler(
         sub_request_entities,
     )
     .await;
-    Ok(to_v3_status(proof_type, result))
+    Ok(to_v3_status(proof_type, None, result))
 }
 
 #[derive(OpenApi)]

host/src/server/utils.rs

@@ -9,10 +9,15 @@ use raiko_reqpool::Status;
 use raiko_tasks::TaskStatus;
 use serde_json::Value;
 
-pub fn to_v2_status(proof_type: ProofType, result: Result<Status, String>) -> v2::Status {
+pub fn to_v2_status(
+    proof_type: ProofType,
+    batch_id: Option<u64>,
+    result: Result<Status, String>,
+) -> v2::Status {
     match result {
         Ok(status) => v2::Status::Ok {
             proof_type,
+            batch_id,
             data: {
                 match status {
                     Status::Registered => v2::ProofResponse::Status {
@@ -57,8 +62,12 @@ pub fn to_v2_cancel_status(result: Result<Status, String>) -> v2::CancelStatus {
 }
 
 // TODO: remove the staled interface
-pub fn to_v3_status(proof_type: ProofType, result: Result<Status, String>) -> v3::Status {
-    to_v2_status(proof_type, result)
+pub fn to_v3_status(
+    proof_type: ProofType,
+    batch_id: Option<u64>,
+    result: Result<Status, String>,
+) -> v3::Status {
+    to_v2_status(proof_type, batch_id, result)
 }
 
 pub fn to_v3_cancel_status(result: Result<Status, String>) -> v3::CancelStatus {
@@ -90,29 +99,29 @@ pub async fn draw_for_zk_any_request(
     Ok(actor.draw(&blockhash))
 }
 
-pub fn fulfill_sp1_params(req: &mut Value) {
-    let zk_any_opts = req["zk_any"].as_object().clone();
-    let sp1_recursion = match zk_any_opts {
-        None => serde_json::Value::String("plonk".to_string()),
-        Some(zk_any) => {
-            let aggregation = zk_any["aggregation"].as_bool().unwrap_or(false);
-            if aggregation {
-                serde_json::Value::String("compressed".to_string())
-            } else {
-                serde_json::Value::String("plonk".to_string())
-            }
-        }
-    };
-
-    let sp1_opts = req["sp1"].as_object_mut();
-    match sp1_opts {
-        None => {
-            let mut sp1_opts = serde_json::Map::new();
-            sp1_opts.insert("recursion".to_string(), sp1_recursion);
-            req["sp1"] = serde_json::Value::Object(sp1_opts);
-        }
-        Some(sp1_opts) => {
-            sp1_opts.insert("recursion".to_string(), sp1_recursion);
-        }
-    }
+pub async fn draw_for_zk_any_batch_request(
+    actor: &Actor,
+    batch_proof_request_opt: &Value,
+) -> HostResult<Option<ProofType>> {
+    let l1_network =
+        batch_proof_request_opt["l1_network"]
+            .as_str()
+            .ok_or(RaikoError::InvalidRequestConfig(
+                "Missing network".to_string(),
+            ))?;
+    let batches =
+        batch_proof_request_opt["batches"]
+            .as_array()
+            .ok_or(RaikoError::InvalidRequestConfig(
+                "Missing batches".to_string(),
+            ))?;
+    let first_batch = batches.first().ok_or(RaikoError::InvalidRequestConfig(
+        "batches is empty".to_string(),
+    ))?;
+    let l1_inclusion_block_number = first_batch["l1_inclusion_block_number"].as_u64().ok_or(
+        RaikoError::InvalidRequestConfig("Missing l1_inclusion_block_number".to_string()),
+    )?;
+    let (_, blockhash) =
+        get_task_data(&l1_network, l1_inclusion_block_number, actor.chain_specs()).await?;
+    Ok(actor.draw(&blockhash))
 }

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(),