feat: stateful (#56)

* use zktrie-ng

* fix

* clean legacy code

* add more metrics

* range adjustment

* update grafana

* fix

* update

* tmp

* save

* save work

* fix changes

* done

* include genesis

* clippy

* back pressure

* fix

* fix

* fix debug

Cargo.toml

@@ -29,9 +29,11 @@ async-channel = "2.2"
 clap = "4"
 env_logger = "0.9"
 futures = "0.3"
+tokio-retry = "0.3"
 serde = "1.0"
 serde_json = "1.0"
 serde_with = "3.9"
+sled = "0.34"
 tokio = { version = "1", default-features = false }
 url = "2.5"
 
@@ -90,6 +92,7 @@ ruint = { git = "https://github.com/scroll-tech/uint.git", branch = "v1.12.3" }
 
 alloy = { git = "https://github.com/scroll-tech/alloy.git", branch = "v0.3.0" }
 alloy-eips = { git = "https://github.com/scroll-tech/alloy.git", branch = "v0.3.0" }
+alloy-serde = { git = "https://github.com/scroll-tech/alloy.git", branch = "v0.3.0" }
 
 alloy-eip2930 = { git = "https://github.com/scroll-tech/alloy-eips", branch = "v0.1.0" }
 alloy-eip7702 = { git = "https://github.com/scroll-tech/alloy-eips", branch = "v0.1.0" }

crates/bin/src/commands/run_file.rs

@@ -6,6 +6,7 @@ use sbv::{
     core::{ChunkInfo, EvmExecutorBuilder, HardforkConfig},
     primitives::{types::BlockTrace, zk_trie::db::kv::HashMapDb, Block, B256},
 };
+use serde::Deserialize;
 use std::panic::catch_unwind;
 use std::{cell::RefCell, path::PathBuf};
 use tiny_keccak::{Hasher, Keccak};
@@ -77,20 +78,23 @@ impl RunFileCommand {
 
         let fork_config = fork_config(traces[0].chain_id());
         let (chunk_info, mut zktrie_db) = ChunkInfo::from_block_traces(&traces);
+        let mut code_db = HashMapDb::default();
 
         let tx_bytes_hasher = RefCell::new(Keccak::v256());
 
-        let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), &mut zktrie_db)
+        let mut executor = EvmExecutorBuilder::new(&mut code_db, &mut zktrie_db)
             .hardfork_config(fork_config)
-            .with_hooks(&traces[0], |hooks| {
+            .chain_id(traces[0].chain_id())
+            .with_hooks(traces[0].root_before(), |hooks| {
                 hooks.add_tx_rlp_handler(|_, rlp| {
                     tx_bytes_hasher.borrow_mut().update(rlp);
                 });
             })?;
-        executor.handle_block(&traces[0])?;
+        for trace in traces.iter() {
+            executor.insert_codes(trace)?;
+        }
 
-        for trace in traces[1..].iter() {
-            executor.update_db(trace)?;
+        for trace in traces.iter() {
             executor.handle_block(trace)?;
         }
 
@@ -116,22 +120,27 @@ async fn read_block_trace(path: &PathBuf) -> anyhow::Result<BlockTrace> {
 }
 
 fn deserialize_block_trace(trace: &str) -> anyhow::Result<BlockTrace> {
+    let block_trace = deserialize_may_wrapped::<BlockTrace>(trace)?;
+    if block_trace.storage_trace.flatten_proofs.is_empty() {
+        dev_warn!("Storage trace is empty, try to deserialize as legacy storage trace");
+        let legacy_trace = deserialize_may_wrapped::<BlockTrace<LegacyStorageTrace>>(trace)?;
+        return Ok(legacy_trace.into());
+    }
+    Ok(block_trace)
+}
+fn deserialize_may_wrapped<'de, T: Deserialize<'de>>(trace: &'de str) -> anyhow::Result<T> {
     // Try to deserialize `BlockTrace` from JSON. In case of failure, try to
     // deserialize `BlockTrace` from a JSON-RPC response that has the actual block
     // trace nested in the value of the key "result".
-    Ok(
-        serde_json::from_str::<BlockTrace<LegacyStorageTrace>>(trace)
-            .or_else(|_| {
-                #[derive(serde::Deserialize, Default, Debug, Clone)]
-                pub struct BlockTraceJsonRpcResult {
-                    pub result: BlockTrace<LegacyStorageTrace>,
-                }
-                Ok::<_, serde_json::Error>(
-                    serde_json::from_str::<BlockTraceJsonRpcResult>(trace)?.result,
-                )
-            })?
-            .into(),
-    )
+    Ok(serde_json::from_str::<T>(trace).or_else(|_| {
+        #[derive(serde::Deserialize, Default, Debug, Clone)]
+        pub struct BlockTraceJsonRpcResult<T> {
+            pub result: T,
+        }
+        Ok::<_, serde_json::Error>(
+            serde_json::from_str::<BlockTraceJsonRpcResult<T>>(trace)?.result,
+        )
+    })?)
 }
 
 async fn run_trace(

crates/bin/src/utils.rs

@@ -19,6 +19,7 @@ fn verify_inner<T: Block + Clone>(
     fork_config: &HardforkConfig,
 ) -> Result<(), VerificationError> {
     dev_trace!("{l2_trace:#?}");
+    let root_before = l2_trace.root_before();
     let root_after = l2_trace.root_after();
 
     // or with v2 trace
@@ -47,10 +48,14 @@ fn verify_inner<T: Block + Clone>(
         },
         "build ZktrieState"
     );
+    let mut code_db = HashMapDb::default();
 
-    let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), &mut zktrie_db)
+    let mut executor = EvmExecutorBuilder::new(&mut code_db, &mut zktrie_db)
         .hardfork_config(*fork_config)
-        .build(&l2_trace)?;
+        .chain_id(l2_trace.chain_id())
+        .build(root_before)?;
+
+    executor.insert_codes(&l2_trace)?;
 
     // TODO: change to Result::inspect_err when sp1 toolchain >= 1.76
     #[allow(clippy::map_identity)]

crates/core/Cargo.toml

@@ -17,6 +17,8 @@ hex.workspace = true
 once_cell.workspace = true
 revm.workspace = true
 rkyv.workspace = true
+serde.workspace = true
+serde_json.workspace = true
 thiserror.workspace = true
 tiny-keccak.workspace = true
 

crates/core/src/chunk.rs

@@ -143,21 +143,24 @@ mod tests {
 
         let fork_config = HardforkConfig::default_from_chain_id(traces[0].chain_id());
         let (chunk_info, mut zktrie_db) = ChunkInfo::from_block_traces(&traces);
+        let mut code_db = HashMapDb::default();
 
         let tx_bytes_hasher = RefCell::new(Keccak::v256());
 
-        let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), &mut zktrie_db)
+        let mut executor = EvmExecutorBuilder::new(&mut code_db, &mut zktrie_db)
             .hardfork_config(fork_config)
-            .with_hooks(&traces[0], |hooks| {
+            .chain_id(traces[0].chain_id())
+            .with_hooks(traces[0].root_before(), |hooks| {
                 hooks.add_tx_rlp_handler(|_, rlp| {
                     tx_bytes_hasher.borrow_mut().update(rlp);
                 });
             })
             .unwrap();
-        executor.handle_block(&traces[0]).unwrap();
+        for trace in traces.iter() {
+            executor.insert_codes(trace).unwrap();
+        }
 
-        for trace in traces[1..].iter() {
-            executor.update_db(trace).unwrap();
+        for trace in traces.iter() {
             executor.handle_block(trace).unwrap();
         }
 

crates/core/src/database.rs

@@ -3,12 +3,13 @@ use revm::{
     db::DatabaseRef,
     primitives::{AccountInfo, Address, Bytecode, B256, U256},
 };
-use sbv_primitives::zk_trie::db::NodeDb;
-use sbv_primitives::zk_trie::hash::ZkHash;
 use sbv_primitives::{
     zk_trie::{
-        db::kv::{KVDatabase, KVDatabaseItem},
-        hash::{key_hasher::NoCacheHasher, poseidon::Poseidon},
+        db::{
+            kv::{KVDatabase, KVDatabaseItem},
+            NodeDb,
+        },
+        hash::{key_hasher::NoCacheHasher, poseidon::Poseidon, ZkHash},
         scroll_types::Account,
         trie::ZkTrie,
     },
@@ -21,7 +22,7 @@ type Result<T, E = DatabaseError> = std::result::Result<T, E>;
 /// A database that consists of account and storage information.
 pub struct EvmDatabase<'a, CodeDb, ZkDb> {
     /// Map of code hash to bytecode.
-    code_db: CodeDb,
+    pub(crate) code_db: &'a mut CodeDb,
     /// Storage root cache, avoid re-query account when storage root is needed
     storage_root_caches: RefCell<HashMap<Address, ZkHash>>,
     /// Storage trie cache, avoid re-creating trie for the same account.
@@ -44,23 +45,12 @@ impl<CodeDb, Db> fmt::Debug for EvmDatabase<'_, CodeDb, Db> {
 }
 
 impl<'a, CodeDb: KVDatabase, ZkDb: KVDatabase + 'static> EvmDatabase<'a, CodeDb, ZkDb> {
-    /// Initialize an EVM database from a block trace.
-    pub fn new<T: Block>(
-        l2_trace: T,
-        mut code_db: CodeDb,
+    /// Initialize an EVM database from a zkTrie root.
+    pub fn new_from_root(
+        committed_zktrie_root: B256,
+        code_db: &'a mut CodeDb,
         zktrie_db: &'a mut NodeDb<ZkDb>,
     ) -> Result<Self> {
-        cycle_tracker_start!("insert CodeDB");
-        for code in l2_trace.codes() {
-            let hash = revm::primitives::keccak256(code);
-            code_db
-                .or_put(hash.as_slice(), code)
-                .map_err(DatabaseError::code_db)?;
-        }
-        cycle_tracker_end!("insert CodeDB");
-
-        let committed_zktrie_root = l2_trace.root_before();
-
         let zktrie = ZkTrie::new_with_root(zktrie_db, NoCacheHasher, committed_zktrie_root)
             .map_err(DatabaseError::zk_trie)?;
 
@@ -113,11 +103,14 @@ impl<'a, CodeDb: KVDatabase, ZkDb: KVDatabase + 'static> EvmDatabase<'a, CodeDb,
     }
 
     /// Update the database with a new block trace.
-    pub fn update<T: Block>(&mut self, l2_trace: T) -> Result<()> {
-        measure_duration_millis!(update_db_duration_milliseconds, self.update_inner(l2_trace))
+    pub fn insert_codes<T: Block>(&mut self, l2_trace: T) -> Result<()> {
+        measure_duration_millis!(
+            update_db_duration_milliseconds,
+            self.insert_codes_inner(l2_trace)
+        )
     }
 
-    fn update_inner<T: Block>(&mut self, l2_trace: T) -> Result<()> {
+    fn insert_codes_inner<T: Block>(&mut self, l2_trace: T) -> Result<()> {
         cycle_tracker_start!("insert CodeDB");
         for code in l2_trace.codes() {
             let hash = revm::primitives::keccak256(code);
@@ -126,13 +119,6 @@ impl<'a, CodeDb: KVDatabase, ZkDb: KVDatabase + 'static> EvmDatabase<'a, CodeDb,
                 .map_err(DatabaseError::code_db)?;
         }
         cycle_tracker_end!("insert CodeDB");
-
-        self.zktrie = cycle_track!(
-            ZkTrie::new_with_root(self.zktrie_db, NoCacheHasher, l2_trace.root_before(),),
-            "ZkTrie::new_with_root"
-        )
-        .map_err(DatabaseError::zk_trie)?;
-
         Ok(())
     }
 }
@@ -161,6 +147,18 @@ impl<CodeDb: KVDatabase, ZkDb: KVDatabase + 'static> DatabaseRef for EvmDatabase
             .get(&account.code_hash)
             .map_err(DatabaseError::code_db)?
             .map(|v| Bytecode::new_legacy(v.into_bytes().into()));
+        if let Some(ref code) = info.code {
+            debug_assert_eq!(
+                info.code_hash,
+                code.hash_slow(),
+                "code hash mismatch for account {address:?}",
+            );
+            assert_eq!(
+                info.code_size,
+                code.original_bytes().len(),
+                "code size mismatch for account {address:?}",
+            );
+        }
 
         Ok(Some(info))
     }
@@ -208,7 +206,9 @@ impl<CodeDb: KVDatabase, ZkDb: KVDatabase + 'static> DatabaseRef for EvmDatabase
 
                 ZkTrie::new_with_root(self.zktrie_db, NoCacheHasher, *storage_root)
                     .inspect_err(|e| {
-                        dev_warn!("storage trie associated with account({address}) not found: {e}")
+                        let backtrace = std::backtrace::Backtrace::force_capture();
+                        dev_warn!("storage trie associated with account({address}) not found: {e}\n{backtrace}");
+
                     })
                     .ok()
             });