make raiko call preflight v2 api

Signed-off-by: smtmfft <[email protected]>

core/Cargo.toml

@@ -64,6 +64,7 @@ clap = { workspace = true }
 assert_cmd = { workspace = true }
 rstest = { workspace = true }
 ethers-core = { workspace = true }
+env_logger = { workspace = true }
 
 [features]
 # powdr = ["dep:powdr"]

core/src/lib.rs

@@ -16,7 +16,7 @@ use tracing::{debug, error, info, warn};
 
 use crate::{
     interfaces::{ProofRequest, RaikoError, RaikoResult},
-    preflight::{preflight, PreflightData},
+    preflight::{preflight, preflight_v2, PreflightData},
     provider::BlockDataProvider,
 };
 
@@ -31,18 +31,21 @@ pub struct Raiko {
     l1_chain_spec: ChainSpec,
     taiko_chain_spec: ChainSpec,
     request: ProofRequest,
+    v2_preflight: bool,
 }
 
 impl Raiko {
     pub fn new(
         l1_chain_spec: ChainSpec,
         taiko_chain_spec: ChainSpec,
         request: ProofRequest,
+        v2_preflight: bool,
     ) -> Self {
         Self {
             l1_chain_spec,
             taiko_chain_spec,
             request,
+            v2_preflight,
         }
     }
 
@@ -63,15 +66,46 @@ impl Raiko {
     pub async fn generate_input<BDP: BlockDataProvider>(
         &self,
         provider: BDP,
+    ) -> RaikoResult<GuestInput> {
+        if self.v2_preflight {
+            self.generate_input_v2(provider).await
+        } else {
+            self.generate_input_v1(provider).await
+        }
+    }
+
+    // v1 uses the geth node & batched rpc calls
+    pub async fn generate_input_v1<BDP: BlockDataProvider>(
+        &self,
+        provider: BDP,
     ) -> RaikoResult<GuestInput> {
         //TODO: read fork from config
         let preflight_data = self.get_preflight_data();
-        info!("Generating input for block {}", self.request.block_number);
+        info!(
+            "Generating v1(common-batch) input for block {}",
+            self.request.block_number
+        );
         preflight(provider, preflight_data)
             .await
             .map_err(Into::<RaikoError>::into)
     }
 
+    // v2 uses the taiko node & single preflight rpc call together with the block rpc calls
+    pub async fn generate_input_v2<BDP: BlockDataProvider>(
+        &self,
+        provider: BDP,
+    ) -> RaikoResult<GuestInput> {
+        //TODO: read fork from config
+        let preflight_data = self.get_preflight_data();
+        info!(
+            "Generating v2(reth-preflight) input for block {}",
+            self.request.block_number
+        );
+        preflight_v2(provider, preflight_data)
+            .await
+            .map_err(Into::<RaikoError>::into)
+    }
+
     pub fn get_output(&self, input: &GuestInput) -> RaikoResult<GuestOutput> {
         let db = create_mem_db(&mut input.clone()).unwrap();
         let mut builder = RethBlockBuilder::new(input, db);
@@ -215,9 +249,14 @@ pub fn merge(a: &mut Value, b: &Value) {
 #[cfg(test)]
 mod tests {
     use crate::interfaces::aggregate_proofs;
-    use crate::{interfaces::ProofRequest, provider::rpc::RpcBlockDataProvider, ChainSpec, Raiko};
+    use crate::{
+        interfaces::ProofRequest,
+        provider::{BlockDataProviderType, RethPreflightBlockDataProvider, RpcBlockDataProvider},
+        ChainSpec, Raiko,
+    };
     use alloy_primitives::Address;
     use alloy_provider::Provider;
+    use env_logger;
     use raiko_lib::{
         consts::{Network, SupportedChainSpecs},
         input::{AggregationGuestInput, AggregationGuestOutput, BlobProofType},
@@ -288,10 +327,34 @@ mod tests {
         taiko_chain_spec: ChainSpec,
         proof_request: ProofRequest,
     ) -> Proof {
-        let provider =
-            RpcBlockDataProvider::new(&taiko_chain_spec.rpc, proof_request.block_number - 1)
-                .expect("Could not create RpcBlockDataProvider");
-        let raiko = Raiko::new(l1_chain_spec, taiko_chain_spec, proof_request.clone());
+        let v2_preflight = std::env::var("V2_PREFLIGHT")
+            .unwrap_or("false".to_string())
+            .parse()
+            .unwrap();
+        let provider = if v2_preflight {
+            let provider: RethPreflightBlockDataProvider = RethPreflightBlockDataProvider::new(
+                &taiko_chain_spec.rpc.clone(),
+                proof_request.block_number - 1,
+            )
+            .await
+            .expect("new RethPreflightBlockDataProvider should be ok");
+            BlockDataProviderType::PreflightRpc(provider)
+        } else {
+            let provider = RpcBlockDataProvider::new(
+                &taiko_chain_spec.rpc.clone(),
+                proof_request.block_number - 1,
+            )
+            .expect("new RpcBlockDataProvider should be ok");
+            BlockDataProviderType::CommonRpc(provider)
+        };
+
+        let raiko = Raiko::new(
+            l1_chain_spec,
+            taiko_chain_spec,
+            proof_request.clone(),
+            v2_preflight,
+        );
+
         let input = raiko
             .generate_input(provider)
             .await
@@ -335,6 +398,8 @@ mod tests {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn test_prove_block_taiko_a7() {
+        let _ = env_logger::builder().is_test(true).try_init();
+
         let proof_type = get_proof_type_from_env();
         let l1_network = Network::Holesky.to_string();
         let network = Network::TaikoA7.to_string();

core/src/provider/mod.rs

@@ -6,13 +6,16 @@ use std::collections::HashMap;
 
 use crate::{
     interfaces::{RaikoError, RaikoResult},
-    provider::rpc::RpcBlockDataProvider,
     MerkleProof,
 };
 
 pub mod db;
-pub mod rpc;
 pub mod preflight_rpc;
+pub mod rpc;
+
+pub use crate::provider::{
+    preflight_rpc::RethPreflightBlockDataProvider, rpc::RpcBlockDataProvider,
+};
 
 #[allow(async_fn_in_trait)]
 pub trait BlockDataProvider {
@@ -31,6 +34,63 @@ pub trait BlockDataProvider {
     ) -> RaikoResult<MerkleProof>;
 }
 
+pub enum BlockDataProviderType {
+    CommonRpc(RpcBlockDataProvider),
+    PreflightRpc(RethPreflightBlockDataProvider),
+}
+
+impl BlockDataProvider for BlockDataProviderType {
+    async fn get_blocks(&self, blocks_to_fetch: &[(u64, bool)]) -> RaikoResult<Vec<Block>> {
+        match self {
+            BlockDataProviderType::CommonRpc(provider) => {
+                provider.get_blocks(blocks_to_fetch).await
+            }
+            BlockDataProviderType::PreflightRpc(provider) => {
+                provider.get_blocks(blocks_to_fetch).await
+            }
+        }
+    }
+
+    async fn get_accounts(&self, accounts: &[Address]) -> RaikoResult<Vec<AccountInfo>> {
+        match self {
+            BlockDataProviderType::CommonRpc(provider) => provider.get_accounts(accounts).await,
+            BlockDataProviderType::PreflightRpc(provider) => provider.get_accounts(accounts).await,
+        }
+    }
+
+    async fn get_storage_values(&self, accounts: &[(Address, U256)]) -> RaikoResult<Vec<U256>> {
+        match self {
+            BlockDataProviderType::CommonRpc(provider) => {
+                provider.get_storage_values(accounts).await
+            }
+            BlockDataProviderType::PreflightRpc(provider) => {
+                provider.get_storage_values(accounts).await
+            }
+        }
+    }
+
+    async fn get_merkle_proofs(
+        &self,
+        block_number: u64,
+        accounts: HashMap<Address, Vec<U256>>,
+        offset: usize,
+        num_storage_proofs: usize,
+    ) -> RaikoResult<MerkleProof> {
+        match self {
+            BlockDataProviderType::CommonRpc(provider) => {
+                provider
+                    .get_merkle_proofs(block_number, accounts, offset, num_storage_proofs)
+                    .await
+            }
+            BlockDataProviderType::PreflightRpc(provider) => {
+                provider
+                    .get_merkle_proofs(block_number, accounts, offset, num_storage_proofs)
+                    .await
+            }
+        }
+    }
+}
+
 pub async fn get_task_data(
     network: &str,
     block_number: u64,
@@ -40,7 +100,7 @@ pub async fn get_task_data(
         .get_chain_spec(network)
         .ok_or_else(|| RaikoError::InvalidRequestConfig("Unsupported raiko network".to_string()))?;
     let provider = RpcBlockDataProvider::new(&taiko_chain_spec.rpc.clone(), block_number - 1)?;
-    let blocks = provider.get_blocks(&[(block_number, true)]).await?;
+    let blocks = provider.get_blocks(&[(block_number, false)]).await?;
     let block = blocks
         .first()
         .ok_or_else(|| RaikoError::RPC("No block for requested block number".to_string()))?;

host/src/cache.rs

@@ -2,7 +2,7 @@ use std::{fs::File, path::PathBuf};
 
 use raiko_core::{
     interfaces::RaikoError,
-    provider::{rpc::RpcBlockDataProvider, BlockDataProvider},
+    provider::{BlockDataProvider, BlockDataProviderType},
 };
 use raiko_lib::input::{get_input_path, GuestInput};
 use tracing::{debug, info};
@@ -42,7 +42,7 @@ pub fn set_input(
 
 pub async fn validate_input(
     cached_input: Option<GuestInput>,
-    provider: &RpcBlockDataProvider,
+    provider: &BlockDataProviderType,
 ) -> HostResult<GuestInput> {
     if let Some(cache_input) = cached_input {
         debug!("Using cached input");

host/src/proof.rs

@@ -8,7 +8,9 @@ use raiko_core::{
     interfaces::{
         aggregate_proofs, cancel_proof, AggregationOnlyRequest, ProofRequest, RaikoError,
     },
-    provider::{get_task_data, rpc::RpcBlockDataProvider},
+    provider::{
+        get_task_data, BlockDataProviderType, RethPreflightBlockDataProvider, RpcBlockDataProvider,
+    },
     Raiko,
 };
 use raiko_lib::{
@@ -416,16 +418,28 @@ pub async fn handle_proof(
     let total_time = Measurement::start("", false);
 
     let v2_preflight = std::env::var("V2_PREFLIGHT").is_ok();
+    let provider = if v2_preflight {
+        BlockDataProviderType::PreflightRpc(
+            RethPreflightBlockDataProvider::new(
+                &taiko_chain_spec.rpc.clone(),
+                proof_request.block_number - 1,
+            )
+            .await?,
+        )
+    } else {
+        BlockDataProviderType::CommonRpc(RpcBlockDataProvider::new(
+            &taiko_chain_spec.rpc.clone(),
+            proof_request.block_number - 1,
+        )?)
+    };
+
     let raiko = Raiko::new(
         l1_chain_spec.clone(),
         taiko_chain_spec.clone(),
         proof_request.clone(),
         v2_preflight,
     );
-    let provider = RpcBlockDataProvider::new(
-        &taiko_chain_spec.rpc.clone(),
-        proof_request.block_number - 1,
-    )?;
+
     let input = match cache::validate_input(cached_input, &provider).await {
         Ok(cache_input) => cache_input,
         Err(_) => {

lib/src/builder.rs

@@ -69,7 +69,7 @@ impl<DB: Database<Error = ProviderError> + DatabaseCommit + OptimisticDatabase>
 {
     /// Creates a new block builder.
     pub fn new(input: &GuestInput, db: DB) -> RethBlockBuilder<DB> {
-        // original tx comes from guest input L1 event, which is a super set of input.block.body
+        // original tx comes from guest input L1 event, which should be a super set of input.block.body
         let original_txs = input
             .get_original_txs()
             .expect("get_original_txs should success");