refactor block builder

guests/eth-block/src/main.rs

@@ -16,7 +16,8 @@
 
 use risc0_zkvm::guest::env;
 use zeth_lib::{
-    block_builder::BlockBuilder, consts::ETH_MAINNET_CHAIN_SPEC, execution::EthTxExecStrategy, validation::Input,
+    block_builder::BlockBuilder, consts::ETH_MAINNET_CHAIN_SPEC, execution::EthTxExecStrategy,
+    mem_db::MemDb, validation::Input,
 };
 
 risc0_zkvm::guest::entry!(main);
@@ -25,9 +26,8 @@ pub fn main() {
     // Read the input previous block and transaction data
     let input: Input = env::read();
     // Build the resulting block
-    let output = BlockBuilder::from(input)
-        .with_chain_spec(ETH_MAINNET_CHAIN_SPEC.clone())
-        .initialize_evm_storage()
+    let output = BlockBuilder::<MemDb>::new(&ETH_MAINNET_CHAIN_SPEC, input)
+        .initialize_db()
         .expect("Failed to create in-memory evm storage")
         .initialize_header()
         .expect("Failed to create the initial block header fields")

host/src/main.rs

@@ -30,6 +30,7 @@ use zeth_lib::{
     block_builder::BlockBuilder,
     consts::{Network, ETH_MAINNET_CHAIN_SPEC},
     execution::EthTxExecStrategy,
+    mem_db::MemDb,
     validation::Input,
 };
 use zeth_primitives::BlockHash;
@@ -89,9 +90,8 @@ async fn main() -> Result<()> {
 
         info!("Running from memory ...");
 
-        let block_builder = BlockBuilder::from(input)
-            .with_chain_spec(ETH_MAINNET_CHAIN_SPEC.clone())
-            .initialize_evm_storage()
+        let block_builder = BlockBuilder::<MemDb>::new(&ETH_MAINNET_CHAIN_SPEC, input)
+            .initialize_db()
             .expect("Error initializing MemDb from Input")
             .initialize_header()
             .expect("Error creating initial block header")

lib/src/block_builder.rs

@@ -14,7 +14,7 @@
 
 use core::mem;
 
-use anyhow::{bail, Result};
+use anyhow::Result;
 use hashbrown::{hash_map, HashMap};
 use revm::primitives::{Account, AccountInfo, Address, Bytecode, B160, B256, U256};
 use zeth_primitives::{
@@ -29,7 +29,7 @@ use crate::{
     consts::ChainSpec,
     execution::TxExecStrategy,
     guest_mem_forget,
-    mem_db::{AccountState, DbAccount, MemDb},
+    mem_db::{AccountState, DbAccount},
     validation::{
         compute_base_fee, compute_block_number, verify_extra_data, verify_gas_limit,
         verify_parent_chain, verify_state_trie, verify_storage_trie, verify_timestamp, Input,
@@ -47,55 +47,37 @@ pub trait BlockBuilderDatabase: revm::Database + Sized {
     fn update(&mut self, address: Address, account: Account);
 }
 
-#[derive(Clone)]
-pub struct BlockBuilder<D> {
-    pub chain_spec: Option<ChainSpec>,
-    pub db: Option<D>,
-    pub header: Option<Header>,
-    pub input: Input,
+#[derive(Clone, Debug)]
+pub struct BlockBuilder<'a, D> {
+    pub(crate) chain_spec: &'a ChainSpec,
+    pub(crate) input: Input,
+    pub(crate) db: Option<D>,
+    pub(crate) header: Option<Header>,
 }
 
-impl From<Input> for BlockBuilder<MemDb> {
-    fn from(input: Input) -> Self {
-        BlockBuilder {
-            chain_spec: None,
-            db: None,
-            header: None,
-            input,
-        }
-    }
-}
-
-impl<D> BlockBuilder<D>
+impl<D> BlockBuilder<'_, D>
 where
     D: BlockBuilderDatabase,
     <D as revm::Database>::Error: std::fmt::Debug,
 {
-    pub fn new(chain_spec: Option<ChainSpec>, db: Option<D>, input: Input) -> Self {
+    /// Creates a new block builder.
+    pub fn new(chain_spec: &ChainSpec, input: Input) -> BlockBuilder<'_, D> {
         BlockBuilder {
             chain_spec,
-            db,
+            db: None,
             header: None,
             input,
         }
     }
 
-    /// Returns a reference to the database.
-    pub fn db(&self) -> Option<&D> {
-        self.db.as_ref()
-    }
-
-    /// Returns a mutable reference to the database.
-    pub fn mut_db(&mut self) -> Option<&mut D> {
-        self.db.as_mut()
-    }
-
-    pub fn with_chain_spec(mut self, chain_spec: ChainSpec) -> Self {
-        self.chain_spec = Some(chain_spec);
+    /// Sets the database.
+    pub fn with_db(mut self, db: D) -> Self {
+        self.db = Some(db);
         self
     }
 
-    pub fn initialize_evm_storage(mut self) -> Result<Self> {
+    /// Initializes the database from the input tries.
+    pub fn initialize_db(mut self) -> Result<Self> {
         verify_state_trie(
             &self.input.parent_state_trie,
             &self.input.parent_header.state_root,
@@ -164,22 +146,20 @@ where
         Ok(self)
     }
 
+    /// Initializes the header. This must be called before executing transactions.
     pub fn initialize_header(mut self) -> Result<Self> {
         // Verify current block
         verify_gas_limit(self.input.gas_limit, self.input.parent_header.gas_limit)?;
         verify_timestamp(self.input.timestamp, self.input.parent_header.timestamp)?;
         verify_extra_data(&self.input.extra_data)?;
         // Initialize result header
-        let Some(ref chain_spec) = self.chain_spec else {
-            bail!("Missing ChainSpec");
-        };
         self.header = Some(Header {
             // Initialize fields that we can compute from the parent
             parent_hash: self.input.parent_header.hash(),
             number: compute_block_number(&self.input.parent_header)?,
             base_fee_per_gas: compute_base_fee(
                 &self.input.parent_header,
-                chain_spec.gas_constants(),
+                self.chain_spec.gas_constants(),
             )?,
             // Initialize metadata from input
             beneficiary: self.input.beneficiary,
@@ -193,10 +173,12 @@ where
         Ok(self)
     }
 
+    /// Executes the transactions.
     pub fn execute_transactions<T: TxExecStrategy>(self) -> Result<Self> {
         T::execute_transactions(self)
     }
 
+    /// Builds the block and returns the header.
     pub fn build(
         mut self,
         mut debug_storage_tries: Option<&mut HashMap<Address, MptNode>>,
@@ -271,4 +253,14 @@ where
 
         Ok(header)
     }
+
+    /// Returns a reference to the database.
+    pub fn db(&self) -> Option<&D> {
+        self.db.as_ref()
+    }
+
+    /// Returns a mutable reference to the database.
+    pub fn mut_db(&mut self) -> Option<&mut D> {
+        self.db.as_mut()
+    }
 }

lib/src/execution.rs

@@ -56,10 +56,7 @@ impl TxExecStrategy for EthTxExecStrategy {
             .header
             .as_mut()
             .expect("Header is not initialized");
-        let Some(ref chain_spec) = block_builder.chain_spec else {
-            bail!("Undefined ChainSpec");
-        };
-        let spec_id = compute_spec_id(chain_spec, header.number)?;
+        let spec_id = compute_spec_id(block_builder.chain_spec, header.number)?;
 
         #[cfg(not(target_os = "zkvm"))]
         {
@@ -83,7 +80,7 @@ impl TxExecStrategy for EthTxExecStrategy {
         let mut evm = EVM::new();
 
         evm.env.cfg = CfgEnv {
-            chain_id: U256::from(chain_spec.chain_id()),
+            chain_id: U256::from(block_builder.chain_spec.chain_id()),
             spec_id,
             ..Default::default()
         };

lib/src/host/mod.rs

@@ -120,13 +120,10 @@ pub fn get_initial_data(
     };
 
     // Create the block builder, run the transactions and extract the DB
-    let mut builder = BlockBuilder::new(
-        Some(ETH_MAINNET_CHAIN_SPEC.clone()),
-        Some(provider_db),
-        input,
-    )
-    .initialize_header()?
-    .execute_transactions::<EthTxExecStrategy>()?;
+    let mut builder = BlockBuilder::new(&ETH_MAINNET_CHAIN_SPEC, input)
+        .with_db(provider_db)
+        .initialize_header()?
+        .execute_transactions::<EthTxExecStrategy>()?;
     let provider_db = builder.mut_db().unwrap();
 
     info!("Gathering inclusion proofs ...");

testeth/Cargo.toml

@@ -10,7 +10,7 @@ hashbrown = "0.13"
 revm = { version = "3.3", default-features = false, features = ["std", "serde"] }
 serde = "1.0"
 serde_json = "1.0"
-serde_with = { version = "3.1", features = ["base64"] }
+serde_with = "3.1"
 stacker = "0.1"
 zeth-guests = { path = "../guests" }
 zeth-lib = { path = "../lib" }

testeth/src/lib.rs

@@ -25,7 +25,7 @@ use zeth_lib::{
         provider_db::ProviderDb,
         Init,
     },
-    mem_db::{DbAccount, MemDb},
+    mem_db::DbAccount,
     validation::Input,
 };
 use zeth_primitives::{
@@ -298,14 +298,14 @@ fn proof_internal(node: &MptNode, key_nibs: &[u8]) -> Result<Vec<Vec<u8>>, anyho
 /// The size of the stack to use for the EVM.
 pub const BIG_STACK_SIZE: usize = 8 * 1024 * 1024;
 
-pub fn new_builder(
-    chain_spec: ChainSpec,
+pub fn create_input(
+    chain_spec: &ChainSpec,
     state: TestState,
     parent_header: Header,
     header: Header,
     transactions: Vec<TestTransaction>,
     withdrawals: Vec<Withdrawal>,
-) -> BlockBuilder<MemDb> {
+) -> Input {
     // create the provider DB
     let provider_db = ProviderDb::new(
         Box::new(TestProvider {
@@ -329,7 +329,8 @@ pub fn new_builder(
     };
 
     // create and run the block builder once to create the initial DB
-    let builder = BlockBuilder::new(Some(chain_spec.clone()), Some(provider_db), input)
+    let builder = BlockBuilder::new(chain_spec, input)
+        .with_db(provider_db)
         .initialize_header()
         .unwrap();
     // execute the transactions with a larger stack
@@ -342,7 +343,7 @@ pub fn new_builder(
     let fini_proofs = get_proofs(provider_db, provider_db.get_latest_db().storage_keys()).unwrap();
     let ancestor_headers = provider_db.get_ancestor_headers().unwrap();
 
-    let input: Input = Init {
+    Init {
         db: provider_db.get_initial_db().clone(),
         init_block: parent_header,
         init_proofs,
@@ -352,8 +353,5 @@ pub fn new_builder(
         fini_proofs,
         ancestor_headers,
     }
-    .into();
-
-    let builder: BlockBuilder<MemDb> = input.into();
-    builder.with_chain_spec(chain_spec)
+    .into()
 }

testeth/testguest/src/main.rs

@@ -16,9 +16,9 @@
 
 use risc0_zkvm::guest::env;
 use zeth_lib::{
-    block_builder::BlockBuilder, execution::EthTxExecStrategy, validation::Input,
+    block_builder::BlockBuilder, consts::ChainSpec, execution::EthTxExecStrategy, mem_db::MemDb,
+    validation::Input,
 };
-use zeth_lib::consts::ChainSpec;
 
 risc0_zkvm::guest::entry!(main);
 
@@ -28,9 +28,8 @@ pub fn main() {
     // Read the input previous block and transaction data
     let input: Input = env::read();
     // Build the resulting block
-    let output = BlockBuilder::from(input)
-        .with_chain_spec(chain_spec)
-        .initialize_evm_storage()
+    let output = BlockBuilder::<MemDb>::new(&chain_spec, input)
+        .initialize_db()
         .expect("Failed to create in-memory evm storage")
         .initialize_header()
         .expect("Failed to create the initial block header fields")

testeth/tests/evm.rs

@@ -17,7 +17,9 @@ use std::{fs::File, io::BufReader, path::PathBuf};
 use revm::primitives::SpecId;
 use rstest::rstest;
 use serde_json::Value;
-use zeth_lib::{consts::ChainSpec, execution::EthTxExecStrategy};
+use zeth_lib::{
+    block_builder::BlockBuilder, consts::ChainSpec, execution::EthTxExecStrategy, mem_db::MemDb,
+};
 use zeth_primitives::block::Header;
 use zeth_testeth::*;
 
@@ -68,18 +70,19 @@ fn evm(
             let expected_header: Header = block_header.clone().into();
             assert_eq!(&expected_header.hash(), &block_header.hash);
 
-            let builder = new_builder(
-                chain_spec.clone(),
+            let input = create_input(
+                &chain_spec,
                 state,
                 parent_header.clone(),
                 expected_header.clone(),
                 block.transactions,
                 block.withdrawals.unwrap_or_default(),
-            )
-            .initialize_evm_storage()
-            .unwrap()
-            .initialize_header()
-            .unwrap();
+            );
+            let builder = BlockBuilder::<MemDb>::new(&chain_spec, input)
+                .initialize_db()
+                .unwrap()
+                .initialize_header()
+                .unwrap();
             // execute the transactions with a larger stack
             let builder = stacker::grow(BIG_STACK_SIZE, move || {
                 builder.execute_transactions::<EthTxExecStrategy>().unwrap()