refactor: [phase2] remove mpt-zktrie (#45)

* remove mpt-zktrie

* rename

Author: lightsing

Cargo.lock

@@ -20,7 +20,7 @@ thiserror = "1.0"
 tiny-keccak = "2.0"
 
 # dependencies from scroll-tech
-mpt-zktrie = { git = "https://github.com/scroll-tech/zkevm-circuits", branch = "feat/rkyv" }
+poseidon-bn254 = { git = "https://github.com/scroll-tech/poseidon-bn254", branch = "master" }
 zktrie = { git = "https://github.com/scroll-tech/zktrie.git", branch = "main", features= ["rs_zktrie"] }
 
 # binary dependencies
@@ -86,7 +86,6 @@ missing-debug-implementations = "deny"
 
 [patch.crates-io]
 ethers-core = { git = "https://github.com/scroll-tech/ethers-rs.git", branch = "feat/rkyv" }
-ethers-signers  = { git = "https://github.com/scroll-tech/ethers-rs.git", branch = "feat/rkyv" }
 primitive-types = { git = "https://github.com/scroll-tech/parity-common.git", branch = "feat/rkyv" }
 ethereum-types = { git = "https://github.com/scroll-tech/parity-common.git", branch = "feat/rkyv" }
 ff = { git = "https://github.com/scroll-tech/ff", branch = "feat/sp1" }
@@ -107,13 +106,10 @@ alloy-primitives = { git = "https://github.com/scroll-tech/alloy-core", branch =
 # ethers-providers = { path = "../ethers-rs/ethers-providers" }
 # ethers = { path = "../ethers-rs/ethers" }
 # ethers-etherscan = { path = "../ethers-rs/ethers-etherscan" }
-# ethers-signers  = { path = "../ethers-rs/ethers-signers" }
 # primitive-types = { path = "../parity-common/primitive-types" }
 # ethereum-types = { path = "../parity-common/ethereum-types" }
 
 # for local development
-# [patch."https://github.com/scroll-tech/zkevm-circuits"]
-# mpt-zktrie = { path = "../zkevm-circuits/zktrie" }
 # [patch."https://github.com/scroll-tech/revm"]
 # revm = { path = "../revm/crates/revm" }
 

Cargo.toml

@@ -27,8 +27,6 @@ url.workspace = true
 
 sbv.workspace = true
 
-mpt-zktrie.workspace = true
-
 pprof = { workspace = true, optional = true }
 tracing-subscriber = { workspace = true, optional = true }
 

crates/bin/Cargo.toml

@@ -74,11 +74,11 @@ impl RunFileCommand {
         }
 
         let fork_config = fork_config(traces[0].chain_id());
-        let (chunk_info, mut zktrie_state) = ChunkInfo::from_block_traces(&traces);
+        let (chunk_info, zktrie_db) = ChunkInfo::from_block_traces(&traces);
 
         let tx_bytes_hasher = RefCell::new(Keccak::v256());
 
-        let mut executor = EvmExecutorBuilder::new(&zktrie_state)
+        let mut executor = EvmExecutorBuilder::new(zktrie_db.clone())
             .hardfork_config(fork_config)
             .with_execute_hooks(|hooks| {
                 hooks.add_tx_rlp_handler(|_, rlp| {
@@ -93,7 +93,7 @@ impl RunFileCommand {
             executor.handle_block(trace)?;
         }
 
-        let post_state_root = executor.commit_changes(&mut zktrie_state);
+        let post_state_root = executor.commit_changes(&zktrie_db);
         if post_state_root != chunk_info.post_state_root() {
             bail!("post state root mismatch");
         }

crates/bin/src/commands/run_file.rs

@@ -1,8 +1,9 @@
-use mpt_zktrie::ZktrieState;
+use sbv::primitives::zk_trie::ZkMemoryDb;
 use sbv::{
     core::{EvmExecutorBuilder, HardforkConfig, VerificationError},
     primitives::Block,
 };
+use std::rc::Rc;
 
 pub fn verify<T: Block + Clone>(
     l2_trace: T,
@@ -36,20 +37,19 @@ fn verify_inner<T: Block + Clone>(
         .build()
         .unwrap();
 
-    let mut zktrie_state = cycle_track!(
+    let zktrie_db = cycle_track!(
         {
-            let old_root = l2_trace.root_before();
-            let mut zktrie_state = ZktrieState::construct(old_root.0.into());
+            let mut zktrie_db = ZkMemoryDb::new();
             measure_duration_histogram!(
-                build_zktrie_state_duration_microseconds,
-                l2_trace.build_zktrie_state(&mut zktrie_state)
+                build_zktrie_db_duration_microseconds,
+                l2_trace.build_zktrie_db(&mut zktrie_db)
             );
-            zktrie_state
+            Rc::new(zktrie_db)
         },
         "build ZktrieState"
     );
 
-    let mut executor = EvmExecutorBuilder::new(&zktrie_state)
+    let mut executor = EvmExecutorBuilder::new(zktrie_db.clone())
         .hardfork_config(*fork_config)
         .build(&l2_trace)?;
 
@@ -66,7 +66,7 @@ fn verify_inner<T: Block + Clone>(
         update_metrics_counter!(verification_error);
         e
     })?;
-    let revm_root_after = executor.commit_changes(&mut zktrie_state);
+    let revm_root_after = executor.commit_changes(&zktrie_db);
 
     #[cfg(feature = "profiling")]
     if let Ok(report) = guard.report().build() {

crates/bin/src/utils.rs

@@ -20,10 +20,6 @@ rkyv.workspace = true
 thiserror.workspace = true
 tiny-keccak.workspace = true
 
-# dependencies from scroll-tech
-mpt-zktrie.workspace = true
-zktrie.workspace = true
-
 sbv-primitives.workspace = true
 sbv-utils.workspace = true
 

crates/core/Cargo.toml

@@ -1,6 +1,6 @@
-use mpt_zktrie::ZktrieState;
 use revm::primitives::B256;
-use sbv_primitives::Block;
+use sbv_primitives::{zk_trie::ZkMemoryDb, Block};
+use std::rc::Rc;
 use tiny_keccak::{Hasher, Keccak};
 
 /// A chunk is a set of continuous blocks.
@@ -22,7 +22,7 @@ pub struct ChunkInfo {
 
 impl ChunkInfo {
     /// Construct by block traces
-    pub fn from_block_traces<T: Block>(traces: &[T]) -> (Self, ZktrieState) {
+    pub fn from_block_traces<T: Block>(traces: &[T]) -> (Self, Rc<ZkMemoryDb>) {
         let chain_id = traces.first().unwrap().chain_id();
         let prev_state_root = traces
             .first()
@@ -41,13 +41,14 @@ impl ChunkInfo {
         let mut data_hash = B256::ZERO;
         data_hasher.finalize(&mut data_hash.0);
 
-        let mut zktrie_state = ZktrieState::construct(prev_state_root.0.into());
+        let mut zktrie_db = ZkMemoryDb::new();
         for trace in traces.iter() {
             measure_duration_histogram!(
-                build_zktrie_state_duration_microseconds,
-                trace.build_zktrie_state(&mut zktrie_state)
+                build_zktrie_db_duration_microseconds,
+                trace.build_zktrie_db(&mut zktrie_db)
             );
         }
+        let zktrie_db = Rc::new(zktrie_db);
 
         let info = ChunkInfo {
             chain_id,
@@ -57,7 +58,7 @@ impl ChunkInfo {
             data_hash,
         };
 
-        (info, zktrie_state)
+        (info, zktrie_db)
     }
 
     /// Public input hash for a given chunk is defined as
@@ -137,18 +138,17 @@ mod tests {
         });
 
         let fork_config = HardforkConfig::default_from_chain_id(traces[0].chain_id());
-        let (chunk_info, mut zktrie_state) = ChunkInfo::from_block_traces(&traces);
+        let (chunk_info, zktrie_db) = ChunkInfo::from_block_traces(&traces);
 
         let tx_bytes_hasher = RefCell::new(Keccak::v256());
 
-        let mut executor = EvmExecutorBuilder::new(&zktrie_state)
+        let mut executor = EvmExecutorBuilder::new(zktrie_db.clone())
             .hardfork_config(fork_config)
             .with_execute_hooks(|hooks| {
                 hooks.add_tx_rlp_handler(|_, rlp| {
                     tx_bytes_hasher.borrow_mut().update(rlp);
                 });
             })
-            .zktrie_state(&zktrie_state)
             .build(&traces[0])
             .unwrap();
         executor.handle_block(&traces[0]).unwrap();
@@ -158,7 +158,7 @@ mod tests {
             executor.handle_block(trace).unwrap();
         }
 
-        let post_state_root = executor.commit_changes(&mut zktrie_state);
+        let post_state_root = executor.commit_changes(&zktrie_db);
         assert_eq!(post_state_root, chunk_info.post_state_root);
         drop(executor); // drop executor to release Rc<Keccek>
 

crates/core/src/chunk.rs

@@ -1,16 +1,16 @@
 use crate::error::ZkTrieError;
-use mpt_zktrie::state::StorageData;
-use mpt_zktrie::{AccountData, ZktrieState};
 use once_cell::sync::Lazy;
-use revm::db::AccountState;
 use revm::{
-    db::DatabaseRef,
+    db::{AccountState, DatabaseRef},
     primitives::{AccountInfo, Address, Bytecode, B256, U256},
 };
-use sbv_primitives::Block;
+use sbv_primitives::{
+    init_hash_scheme,
+    zk_trie::{SharedMemoryDb, ZkMemoryDb, ZkTrie},
+    Block,
+};
 use std::rc::Rc;
 use std::{cell::RefCell, collections::HashMap, convert::Infallible, fmt};
-use zktrie::{SharedMemoryDb, ZkMemoryDb, ZkTrie};
 
 type Result<T, E = ZkTrieError> = std::result::Result<T, E>;
 
@@ -26,36 +26,38 @@ pub struct ReadOnlyDB {
     /// Storage trie cache, avoid re-creating trie for the same account.
     /// Need to invalidate before `update`, otherwise the trie root may be outdated.
     storage_trie_refs: RefCell<HashMap<Address, Lazy<ZkTrie<SharedMemoryDb>, StorageTrieLazyFn>>>,
-    /// Current zkTrie root based on the block trace.
-    zktrie_root: B256,
+    /// Current uncommitted zkTrie root based on the block trace.
+    committed_zktrie_root: B256,
     /// The underlying zkTrie database.
     zktrie_db: Rc<ZkMemoryDb>,
-    /// Current view of zkTrie database with `zktrie_root`.
+    /// Current view of zkTrie database.
     zktrie_db_ref: ZkTrie<SharedMemoryDb>,
 }
 
 impl fmt::Debug for ReadOnlyDB {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.debug_struct("ReadOnlyDB")
             .field("code_db", &self.code_db.len())
-            .field("zktrie_root", &self.zktrie_root)
+            .field("committed_zktrie_root", &self.committed_zktrie_root)
             .finish()
     }
 }
 
 impl ReadOnlyDB {
     /// Initialize an EVM database from a block trace.
-    pub fn new<T: Block>(l2_trace: T, zktrie_state: &ZktrieState) -> Result<Self> {
+    pub fn new<T: Block>(l2_trace: T, zktrie_db: &Rc<ZkMemoryDb>) -> Result<Self> {
         let size_hint = l2_trace.codes().len();
-        Self::new_with_size_hint(l2_trace, zktrie_state, size_hint)
+        Self::new_with_size_hint(l2_trace, zktrie_db, size_hint)
     }
 
     /// Initialize an EVM database from a block trace with size hint of code database.
     pub fn new_with_size_hint<T: Block>(
         l2_trace: T,
-        zktrie_state: &ZktrieState,
+        zktrie_db: &Rc<ZkMemoryDb>,
         size_hint: usize,
     ) -> Result<Self> {
+        init_hash_scheme();
+
         cycle_tracker_start!("insert CodeDB");
         let mut code_db = HashMap::with_capacity(size_hint);
         for code in l2_trace.codes() {
@@ -67,17 +69,16 @@ impl ReadOnlyDB {
         }
         cycle_tracker_end!("insert CodeDB");
 
-        let zktrie_root = l2_trace.root_before().0.into();
+        let uncommitted_zktrie_root = l2_trace.root_before();
 
         Ok(ReadOnlyDB {
             code_db,
             prev_storage_roots: Default::default(),
             storage_trie_refs: Default::default(),
-            zktrie_root,
-            zktrie_db: zktrie_state.zk_db.clone(),
-            zktrie_db_ref: zktrie_state
-                .zk_db
-                .new_ref_trie(&zktrie_root.0)
+            committed_zktrie_root: uncommitted_zktrie_root,
+            zktrie_db: zktrie_db.clone(),
+            zktrie_db_ref: zktrie_db
+                .new_ref_trie(&uncommitted_zktrie_root.0)
                 .ok_or(ZkTrieError::ZkTrieRootNotFound)?,
         })
     }
@@ -106,6 +107,18 @@ impl ReadOnlyDB {
             .unwrap_or_default()
     }
 
+    /// Get the zkTrie root.
+    #[inline]
+    pub(crate) fn committed_zktrie_root(&self) -> B256 {
+        self.committed_zktrie_root
+    }
+
+    /// Get the zkTrie root.
+    #[inline]
+    pub(crate) fn updated_committed_zktrie_root(&mut self, new_root: B256) {
+        self.committed_zktrie_root = new_root;
+    }
+
     /// Update the database with a new block trace.
     pub fn update<T: Block>(&mut self, l2_trace: T) -> Result<()> {
         measure_duration_histogram!(update_db_duration_microseconds, self.update_inner(l2_trace))
@@ -122,11 +135,9 @@ impl ReadOnlyDB {
         }
         cycle_tracker_end!("insert CodeDB");
 
-        self.zktrie_root = l2_trace.root_before().0.into();
-
         self.zktrie_db_ref = self
             .zktrie_db
-            .new_ref_trie(&self.zktrie_root.0)
+            .new_ref_trie(&l2_trace.root_before().0)
             .ok_or(ZkTrieError::ZkTrieRootNotFound)?;
 
         Ok(())
@@ -154,11 +165,15 @@ impl DatabaseRef for ReadOnlyDB {
         Ok(self
             .zktrie_db_ref
             .get_account(address.as_slice())
-            .map(AccountData::from)
             .map(|account_data| {
-                let code_hash = B256::from(account_data.keccak_code_hash.0);
-
-                let storage_root = account_data.storage_root;
+                let code_size =
+                    u64::from_be_bytes((&account_data[0][16..24]).try_into().unwrap()) as usize;
+                let nonce = u64::from_be_bytes((&account_data[0][24..]).try_into().unwrap());
+                let balance = U256::from_be_bytes(account_data[1]);
+                let code_hash = B256::from(account_data[3]);
+                let poseidon_code_hash = B256::from(account_data[4]);
+
+                let storage_root = B256::from(account_data[2]);
                 self.prev_storage_roots
                     .borrow_mut()
                     .entry(address)
@@ -174,11 +189,11 @@ impl DatabaseRef for ReadOnlyDB {
                     })),
                 );
                 AccountInfo {
-                    balance: U256::from_limbs(account_data.balance.0),
-                    nonce: account_data.nonce,
-                    code_size: account_data.code_size as usize,
+                    balance,
+                    nonce,
+                    code_size,
                     code_hash,
-                    poseidon_code_hash: B256::from(account_data.poseidon_code_hash.0),
+                    poseidon_code_hash,
                     code: self.code_db.get(&code_hash).cloned(),
                 }
             }))
@@ -208,8 +223,7 @@ impl DatabaseRef for ReadOnlyDB {
                 let storage_root = self
                     .zktrie_db_ref
                     .get_account(address.as_slice())
-                    .map(AccountData::from)
-                    .map(|account_data| account_data.storage_root)
+                    .map(|account_data| B256::from(account_data[2]))
                     .unwrap_or_default();
                 let zktrie_db = self.zktrie_db.clone();
                 Lazy::new(Box::new(move || {
@@ -222,8 +236,7 @@ impl DatabaseRef for ReadOnlyDB {
 
         Ok(trie
             .get_store(&index.to_be_bytes::<32>())
-            .map(StorageData::from)
-            .map(|val| U256::from_limbs(val.as_ref().0))
+            .map(|store_data| U256::from_be_bytes(store_data))
             .unwrap_or_default())
     }