[feat] refactor and use trace v2 (#17)

* upgrade revm to v37

* suppress log

* bps counter

* fix

* less log

* noticeable log

* short hash

* upgrade revm and rust toolchain

* remove legacy feature gate

* upgrade to v40

* refactor and use trace v2

* do this later

Author: lightsing

src/bin/trace-verifier/utils.rs

@@ -1,6 +1,6 @@
-use eth_types::l2_types::BlockTrace;
+use eth_types::l2_types::{BlockTrace, BlockTraceV2};
 use eth_types::ToWord;
-use stateless_block_verifier::{EvmExecutor, HardforkConfig};
+use stateless_block_verifier::{utils, EvmExecutorBuilder, HardforkConfig};
 use std::sync::atomic::AtomicUsize;
 use std::sync::{LazyLock, Mutex};
 use std::time::Instant;
@@ -17,6 +17,8 @@ pub fn verify(
     trace!("{:#?}", l2_trace);
     let root_after = l2_trace.storage_trace.root_after.to_word();
 
+    let v2_trace = BlockTraceV2::from(l2_trace.clone());
+
     let now = Instant::now();
 
     #[cfg(feature = "profiling")]
@@ -26,8 +28,17 @@ pub fn verify(
         .build()
         .unwrap();
 
-    let mut executor = EvmExecutor::new(&l2_trace, &fork_config, disable_checks);
-    let revm_root_after = executor.handle_block(&l2_trace).to_word();
+    let mut executor = EvmExecutorBuilder::new()
+        .hardfork_config(*fork_config)
+        .with_execute_hooks(|hooks| {
+            if !disable_checks {
+                hooks.add_post_tx_execution_handler(move |executor, tx_id| {
+                    utils::post_check(executor.db(), &l2_trace.execution_results[tx_id]);
+                })
+            }
+        })
+        .build(&v2_trace);
+    let revm_root_after = executor.handle_block(&v2_trace).to_word();
 
     #[cfg(feature = "profiling")]
     if let Ok(report) = guard.report().build() {

src/database.rs

@@ -1,6 +1,6 @@
 use crate::utils::{collect_account_proofs, collect_storage_proofs};
 use eth_types::{
-    l2_types::{trace::collect_codes, BlockTrace},
+    l2_types::BlockTraceV2,
     state_db::{self, CodeDB, StateDB},
     ToWord, H160,
 };
@@ -20,7 +20,7 @@ pub struct ReadOnlyDB {
 
 impl ReadOnlyDB {
     /// Initialize an EVM database from a block trace.
-    pub fn new(l2_trace: &BlockTrace) -> Self {
+    pub fn new(l2_trace: &BlockTraceV2) -> Self {
         let mut sdb = StateDB::new();
         for parsed in
             ZktrieState::parse_account_from_proofs(collect_account_proofs(&l2_trace.storage_trace))
@@ -39,8 +39,11 @@ impl ReadOnlyDB {
         }
 
         let mut code_db = CodeDB::new();
-        for (hash, code) in collect_codes(l2_trace, Some(&sdb)).unwrap() {
-            code_db.insert_with_hash(hash, code);
+        for code_trace in l2_trace.codes.iter() {
+            // FIXME: use this later
+            // let hash = code_db.insert(code_trace.code.to_vec());
+            // assert_eq!(hash, code_trace.hash);
+            code_db.insert_with_hash(code_trace.hash, code_trace.code.to_vec());
         }
 
         ReadOnlyDB { code_db, sdb }

src/executor/builder.rs

@@ -0,0 +1,88 @@
+use crate::executor::hooks::ExecuteHooks;
+use crate::utils::{collect_account_proofs, collect_storage_proofs};
+use crate::{EvmExecutor, HardforkConfig, ReadOnlyDB};
+use eth_types::l2_types::{BlockTrace, BlockTraceV2};
+use mpt_zktrie::ZktrieState;
+use revm::db::CacheDB;
+
+/// Builder for EVM executor.
+#[derive(Debug)]
+pub struct EvmExecutorBuilder<H> {
+    hardfork_config: H,
+    execute_hooks: ExecuteHooks,
+}
+
+impl Default for EvmExecutorBuilder<()> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl EvmExecutorBuilder<()> {
+    /// Create a new builder.
+    pub fn new() -> Self {
+        Self {
+            hardfork_config: (),
+            execute_hooks: ExecuteHooks::default(),
+        }
+    }
+}
+
+impl<H1> EvmExecutorBuilder<H1> {
+    /// Set hardfork config.
+    pub fn hardfork_config<H2>(self, hardfork_config: H2) -> EvmExecutorBuilder<H2> {
+        EvmExecutorBuilder {
+            hardfork_config,
+            execute_hooks: self.execute_hooks,
+        }
+    }
+
+    /// Modify execute hooks.
+    pub fn with_execute_hooks(mut self, modify: impl FnOnce(&mut ExecuteHooks)) -> Self {
+        modify(&mut self.execute_hooks);
+        self
+    }
+}
+
+impl EvmExecutorBuilder<HardforkConfig> {
+    /// Initialize an EVM executor from a legacy block trace as the initial state.
+    pub fn build_legacy(self, l2_trace: &BlockTrace) -> EvmExecutor {
+        let v2_trace = BlockTraceV2::from(l2_trace.clone());
+        self.build(&v2_trace)
+    }
+
+    /// Initialize an EVM executor from a block trace as the initial state.
+    pub fn build(self, l2_trace: &BlockTraceV2) -> EvmExecutor {
+        let block_number = l2_trace.header.number.unwrap().as_u64();
+        let spec_id = self.hardfork_config.get_spec_id(block_number);
+        trace!("use spec id {:?}", spec_id);
+
+        let mut db = CacheDB::new(ReadOnlyDB::new(l2_trace));
+        self.hardfork_config.migrate(block_number, &mut db).unwrap();
+
+        let old_root = l2_trace.storage_trace.root_before;
+        let zktrie_state = ZktrieState::from_trace_with_additional(
+            old_root,
+            collect_account_proofs(&l2_trace.storage_trace),
+            collect_storage_proofs(&l2_trace.storage_trace),
+            l2_trace
+                .storage_trace
+                .deletion_proofs
+                .iter()
+                .map(|s| s.as_ref()),
+        )
+        .unwrap();
+        let root = *zktrie_state.root();
+        debug!("building partial statedb done, root {}", hex::encode(root));
+
+        let mem_db = zktrie_state.into_inner();
+        let zktrie = mem_db.new_trie(&root).unwrap();
+
+        EvmExecutor {
+            db,
+            zktrie,
+            spec_id,
+            hooks: self.execute_hooks,
+        }
+    }
+}

src/executor/hooks.rs

@@ -0,0 +1,43 @@
+use crate::EvmExecutor;
+use std::fmt::{Debug, Formatter};
+
+/// Post transaction execution handler.
+pub type PostTxExecutionHandler = dyn Fn(&EvmExecutor, usize) + Send + Sync + 'static;
+
+/// Hooks for the EVM executor.
+#[derive(Default)]
+pub struct ExecuteHooks {
+    post_tx_execution_handlers: Vec<Box<PostTxExecutionHandler>>,
+}
+
+impl ExecuteHooks {
+    /// Create a new hooks.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Add a post transaction execution handler.
+    pub fn add_post_tx_execution_handler<F>(&mut self, handler: F)
+    where
+        F: Fn(&EvmExecutor, usize) + Send + Sync + 'static,
+    {
+        self.post_tx_execution_handlers.push(Box::new(handler));
+    }
+
+    pub(crate) fn post_tx_execution(&self, executor: &EvmExecutor, tx_index: usize) {
+        for handler in &self.post_tx_execution_handlers {
+            handler(executor, tx_index);
+        }
+    }
+}
+
+impl Debug for ExecuteHooks {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ExecuteHooks")
+            .field(
+                "post_tx_execution_handlers",
+                &self.post_tx_execution_handlers.len(),
+            )
+            .finish()
+    }
+}

src/executor.rs renamed to src/executor/mod.rs

@@ -1,68 +1,33 @@
-use crate::{
-    database::ReadOnlyDB,
-    utils::{collect_account_proofs, collect_storage_proofs},
-    HardforkConfig,
-};
-use eth_types::{
-    geth_types::TxType,
-    l2_types::{BlockTrace, ExecutionResult},
-    H160, H256, U256,
-};
-use log::Level;
-use mpt_zktrie::{AccountData, ZktrieState};
+use crate::database::ReadOnlyDB;
+use eth_types::{geth_types::TxType, l2_types::BlockTraceV2, H160, H256, U256};
+use mpt_zktrie::AccountData;
 use revm::{
     db::CacheDB,
     primitives::{AccountInfo, BlockEnv, Env, SpecId, TxEnv},
-    DatabaseRef,
 };
 use std::fmt::Debug;
 use zktrie::ZkTrie;
 
+mod builder;
+/// Execute hooks
+pub mod hooks;
+pub use builder::EvmExecutorBuilder;
+
 /// EVM executor that handles the block.
 pub struct EvmExecutor {
     db: CacheDB<ReadOnlyDB>,
     zktrie: ZkTrie,
     spec_id: SpecId,
-    disable_checks: bool,
+    hooks: hooks::ExecuteHooks,
 }
 impl EvmExecutor {
-    /// Initialize an EVM executor from a block trace as the initial state.
-    pub fn new(l2_trace: &BlockTrace, fork_config: &HardforkConfig, disable_checks: bool) -> Self {
-        let block_number = l2_trace.header.number.unwrap().as_u64();
-        let spec_id = fork_config.get_spec_id(block_number);
-        trace!("use spec id {:?}", spec_id);
-
-        let mut db = CacheDB::new(ReadOnlyDB::new(l2_trace));
-        fork_config.migrate(block_number, &mut db).unwrap();
-
-        let old_root = l2_trace.storage_trace.root_before;
-        let zktrie_state = ZktrieState::from_trace_with_additional(
-            old_root,
-            collect_account_proofs(&l2_trace.storage_trace),
-            collect_storage_proofs(&l2_trace.storage_trace),
-            l2_trace
-                .storage_trace
-                .deletion_proofs
-                .iter()
-                .map(|s| s.as_ref()),
-        )
-        .unwrap();
-        let root = *zktrie_state.root();
-        debug!("building partial statedb done, root {}", hex::encode(root));
-
-        let mem_db = zktrie_state.into_inner();
-        let zktrie = mem_db.new_trie(&root).unwrap();
-
-        Self {
-            db,
-            zktrie,
-            spec_id,
-            disable_checks,
-        }
+    /// Get reference to the DB
+    pub fn db(&self) -> &CacheDB<ReadOnlyDB> {
+        &self.db
     }
 
     /// Handle a block.
-    pub fn handle_block(&mut self, l2_trace: &BlockTrace) -> H256 {
+    pub fn handle_block(&mut self, l2_trace: &BlockTraceV2) -> H256 {
         debug!("handle block {:?}", l2_trace.header.number.unwrap());
         let mut env = Box::<Env>::default();
         env.cfg.chain_id = l2_trace.chain_id;
@@ -100,14 +65,8 @@ impl EvmExecutor {
                 let result = revm.transact_commit().unwrap(); // TODO: handle error
                 trace!("{result:#?}");
             }
+            self.hooks.post_tx_execution(self, idx);
             debug!("handle {idx}th tx done");
-
-            if !self.disable_checks {
-                if let Some(exec) = l2_trace.execution_results.get(idx) {
-                    debug!("post check {idx}th tx");
-                    self.post_check(exec);
-                }
-            }
         }
         self.commit_changes();
         H256::from(self.zktrie.root())
@@ -256,56 +215,6 @@ impl EvmExecutor {
             }
         }
     }
-
-    fn post_check(&mut self, exec: &ExecutionResult) {
-        for account_post_state in exec.account_after.iter() {
-            let local_acc = self
-                .db
-                .basic_ref(account_post_state.address.0.into())
-                .unwrap()
-                .unwrap();
-            if log_enabled!(Level::Trace) {
-                let mut local_acc = local_acc.clone();
-                local_acc.code = None;
-                trace!("local acc {local_acc:?}, trace acc {account_post_state:?}");
-            }
-            let local_balance = U256(*local_acc.balance.as_limbs());
-            if local_balance != account_post_state.balance {
-                let post = account_post_state.balance;
-                error!(
-                    "incorrect balance, local {:#x} {} post {:#x} (diff {}{:#x})",
-                    local_balance,
-                    if local_balance < post { "<" } else { ">" },
-                    post,
-                    if local_balance < post { "-" } else { "+" },
-                    if local_balance < post {
-                        post - local_balance
-                    } else {
-                        local_balance - post
-                    }
-                )
-            }
-            if local_acc.nonce != account_post_state.nonce {
-                error!("incorrect nonce")
-            }
-            let p_hash = account_post_state.poseidon_code_hash;
-            if p_hash.is_zero() {
-                if !local_acc.is_empty() {
-                    error!("incorrect poseidon_code_hash")
-                }
-            } else if local_acc.code_hash.0 != p_hash.0 {
-                error!("incorrect poseidon_code_hash")
-            }
-            let k_hash = account_post_state.keccak_code_hash;
-            if k_hash.is_zero() {
-                if !local_acc.is_empty() {
-                    error!("incorrect keccak_code_hash")
-                }
-            } else if local_acc.keccak_code_hash.0 != k_hash.0 {
-                error!("incorrect keccak_code_hash")
-            }
-        }
-    }
 }
 
 impl Debug for EvmExecutor {

src/lib.rs

@@ -7,8 +7,9 @@ extern crate log;
 mod database;
 mod executor;
 mod hardfork;
-mod utils;
+/// Utilities
+pub mod utils;
 
 pub use database::ReadOnlyDB;
-pub use executor::EvmExecutor;
+pub use executor::{hooks, EvmExecutor, EvmExecutorBuilder};
 pub use hardfork::HardforkConfig;