Support multiple blocks through zeth CLI

.vscode/settings.json

@@ -6,6 +6,7 @@
         "./guests/op-block/Cargo.toml",
         "./guests/op-compose/Cargo.toml",
         "./guests/op-derive/Cargo.toml",
+        "./host/Cargo.toml",
         "./testing/ef-tests/testguest/Cargo.toml"
     ]
 }

host/src/main.rs

@@ -14,6 +14,8 @@
 
 extern crate core;
 
+use std::sync::Arc;
+
 use anyhow::Result;
 use clap::Parser;
 use log::info;
@@ -31,7 +33,7 @@ use zeth_lib::{
 #[tokio::main]
 async fn main() -> Result<()> {
     env_logger::init();
-    let cli = Cli::parse();
+    let cli = Arc::new(Cli::parse());
 
     info!("Using the following image ids:");
     info!("  eth-block: {}", Digest::from(ETH_BLOCK_ID));
@@ -41,15 +43,15 @@ async fn main() -> Result<()> {
 
     // execute the command
     let build_args = cli.build_args();
-    let (image_id, stark) = match build_args.network {
+    let (image_id, starks) = match build_args.network {
         Network::Ethereum => {
             let rpc_url = build_args.eth_rpc_url.clone();
             (
                 ETH_BLOCK_ID,
                 build::build_block::<EthereumStrategy>(
-                    &cli,
+                    cli.clone(),
                     rpc_url,
-                    &ETH_MAINNET_CHAIN_SPEC,
+                    Arc::new(ETH_MAINNET_CHAIN_SPEC.clone()),
                     ETH_BLOCK_ELF,
                 )
                 .await?,
@@ -60,9 +62,9 @@ async fn main() -> Result<()> {
             (
                 OP_BLOCK_ID,
                 build::build_block::<OptimismStrategy>(
-                    &cli,
+                    cli.clone(),
                     rpc_url,
-                    &OP_MAINNET_CHAIN_SPEC,
+                    Arc::new(OP_MAINNET_CHAIN_SPEC.clone()),
                     OP_BLOCK_ELF,
                 )
                 .await?,
@@ -72,30 +74,36 @@ async fn main() -> Result<()> {
             if let Some(composition_size) = build_args.composition {
                 (
                     OP_COMPOSE_ID,
-                    rollups::compose_derived_rollup_blocks(&cli, composition_size).await?,
+                    vec![rollups::compose_derived_rollup_blocks(&cli, composition_size).await?],
                 )
             } else {
-                (OP_DERIVE_ID, rollups::derive_rollup_blocks(&cli).await?)
+                (
+                    OP_DERIVE_ID,
+                    vec![rollups::derive_rollup_blocks(&cli).await?],
+                )
             }
         }
     };
 
     // Create/verify Groth16 SNARK
-    if cli.snark() {
-        let Some((stark_uuid, stark_receipt)) = stark else {
-            panic!("No STARK data to snarkify!");
-        };
+    for stark in starks {
+        if cli.snark() {
+            let Some((stark_uuid, stark_receipt)) = stark else {
+                panic!("No STARK data to snarkify!");
+            };
 
-        if !cli.submit_to_bonsai() {
-            panic!("Bonsai submission flag required to create a SNARK!");
-        }
+            if !cli.submit_to_bonsai() {
+                panic!("Bonsai submission flag required to create a SNARK!");
+            }
 
-        let image_id = Digest::from(image_id);
-        let (snark_uuid, snark_receipt) = stark2snark(image_id, stark_uuid, stark_receipt).await?;
+            let image_id = Digest::from(image_id);
+            let (snark_uuid, snark_receipt) =
+                stark2snark(image_id, stark_uuid, stark_receipt).await?;
 
-        info!("Validating SNARK uuid: {}", snark_uuid);
+            info!("Validating SNARK uuid: {}", snark_uuid);
 
-        verify_groth16_snark(&cli, image_id, snark_receipt).await?;
+            verify_groth16_snark(&cli, image_id, snark_receipt).await?;
+        }
     }
 
     Ok(())

host/src/operations/build.rs

@@ -19,6 +19,8 @@ use ethers_core::types::Transaction as EthersTransaction;
 use log::{info, warn};
 use risc0_zkvm::{compute_image_id, Receipt};
 use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use tokio::sync::Semaphore;
 use zeth_lib::{
     builder::BlockBuilderStrategy,
     consts::ChainSpec,
@@ -27,27 +29,24 @@ use zeth_lib::{
     output::BlockBuildOutput,
 };
 
+const MAX_CONCURRENT_REQUESTS: usize = 5;
+
 use crate::{
-    cli::Cli,
+    cli::{BuildArgs, Cli},
     operations::{execute, maybe_prove, verify_bonsai_receipt},
 };
 
 /// Build a single block using the specified strategy.
-pub async fn build_block<N: BlockBuilderStrategy>(
-    cli: &Cli,
+async fn preflight_block<N: BlockBuilderStrategy>(
+    build_args: BuildArgs,
+    current_block: u64,
     rpc_url: Option<String>,
-    chain_spec: &ChainSpec,
-    guest_elf: &[u8],
-) -> anyhow::Result<Option<(String, Receipt)>>
+    chain_spec: Arc<ChainSpec>,
+) -> anyhow::Result<(BlockBuildInput<N::TxEssence>, BlockBuildOutput)>
 where
     N::TxEssence: 'static + Send + TryFrom<EthersTransaction> + Serialize + Deserialize<'static>,
     <N::TxEssence as TryFrom<EthersTransaction>>::Error: Debug,
 {
-    let build_args = cli.build_args().clone();
-    if build_args.block_count > 1 {
-        warn!("Building multiple blocks is not supported. Only the first block will be built.");
-    }
-
     // Fetch all of the initial data
     let rpc_cache = build_args.cache.as_ref().map(|dir| {
         cache_file_path(
@@ -60,7 +59,7 @@ where
 
     let init_spec = chain_spec.clone();
     let preflight_result = tokio::task::spawn_blocking(move || {
-        N::preflight_with_external_data(&init_spec, rpc_cache, rpc_url, build_args.block_number)
+        N::preflight_with_external_data(&init_spec, rpc_cache, rpc_url, current_block)
     })
     .await?;
     let preflight_data = preflight_result.context("preflight failed")?;
@@ -73,7 +72,7 @@ where
 
     // Verify that the transactions run correctly
     info!("Running from memory ...");
-    let output = N::build_from(chain_spec, input.clone()).context("Error while building block")?;
+    let output = N::build_from(&chain_spec, input.clone()).context("Error while building block")?;
 
     match &output {
         BlockBuildOutput::SUCCESS {
@@ -89,40 +88,110 @@ where
         }
     }
 
+    Ok((input, output))
+}
+
+/// Build a single block using the specified strategy.
+async fn execute_block<N: BlockBuilderStrategy>(
+    input: BlockBuildInput<N::TxEssence>,
+    output: BlockBuildOutput,
+    cli: Arc<Cli>,
+    guest_elf: &'static [u8],
+) -> anyhow::Result<Option<(String, Receipt)>>
+where
+    N::TxEssence: 'static + Send + TryFrom<EthersTransaction> + Serialize + Deserialize<'static>,
+    <N::TxEssence as TryFrom<EthersTransaction>>::Error: Debug,
+{
     let compressed_output = output.with_state_hashed();
-    let result = match cli {
+    let result = match &*cli {
         Cli::Build(..) => None,
         Cli::Run(run_args) => {
             execute(
                 &input,
                 run_args.execution_po2,
                 run_args.profile,
-                guest_elf,
+                &guest_elf,
                 &compressed_output,
                 &cli.execution_tag(),
             );
             None
         }
         Cli::Prove(..) => {
             maybe_prove(
-                cli,
+                &cli,
                 &input,
-                guest_elf,
+                &guest_elf,
                 &compressed_output,
                 Default::default(),
             )
             .await
         }
         Cli::Verify(verify_args) => Some(
             verify_bonsai_receipt(
-                compute_image_id(guest_elf)?,
+                compute_image_id(&guest_elf)?,
                 &compressed_output,
                 verify_args.bonsai_receipt_uuid.clone(),
                 4,
             )
             .await?,
         ),
     };
-
     Ok(result)
 }
+
+/// Build a single block using the specified strategy.
+pub async fn build_block<N: BlockBuilderStrategy>(
+    cli: Arc<Cli>,
+    rpc_url: Option<String>,
+    chain_spec: Arc<ChainSpec>,
+    guest_elf: &'static [u8],
+) -> anyhow::Result<Vec<Option<(String, Receipt)>>>
+where
+    N::TxEssence:
+        'static + Send + Sync + TryFrom<EthersTransaction> + Serialize + Deserialize<'static>,
+    <N::TxEssence as TryFrom<EthersTransaction>>::Error: Debug,
+{
+    let build_args = cli.build_args().clone();
+
+    let semaphore = Arc::new(Semaphore::new(MAX_CONCURRENT_REQUESTS));
+    let mut join_handles = Vec::new();
+
+    let block_num = build_args.block_number;
+    // TODO semantics are a bit mixed with block count (was OP specific)
+    for num in block_num..block_num + build_args.block_count as u64 {
+        // Acquire permit before queueing job.
+        let semaphore = semaphore.clone();
+
+        // Clone variables needed.
+        let rpc_url = rpc_url.clone();
+        let cli = cli.clone();
+        let chain_spec = chain_spec.clone();
+
+        // Spawn blocking for
+        join_handles.push(tokio::spawn(async move {
+            // Acquire permit before sending request.
+            let _permit = semaphore.acquire().await.unwrap();
+
+            let (input, output) =
+                preflight_block::<N>(cli.build_args().clone(), num, rpc_url, chain_spec).await?;
+
+            drop(_permit);
+
+            // TODO this could be separated into a separate task, to make sure Bonsai also
+            //  doesn't get throttled, for now just going quick path of dropping permit after
+            //  preflight.
+            let result = execute_block::<N>(input, output, cli, guest_elf).await;
+
+            result
+        }));
+    }
+
+    // Collect responses from tasks.
+    let mut responses = Vec::new();
+    for jh in join_handles {
+        let response = jh.await?;
+        responses.push(response?);
+    }
+
+    Ok(responses)
+}