feat: zktrie-ng rkyv node (#60)

* update with zktrie-ng * update zktrie-ng * fix dep version * remove sp1 related changes on master * fix clippy * fix tests
scroll-tech · Oct 8, 2024 · 06f85ee · 06f85ee
1 parent a615e96
commit 06f85ee
Show file tree

Hide file tree

Showing 8 changed files with 314 additions and 309 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/bin/src/commands/run_file.rs b/crates/bin/src/commands/run_file.rs
@@ -3,9 +3,8 @@ use anyhow::bail;
 use clap::Args;
 use sbv::{
     core::{ChunkInfo, EvmExecutorBuilder, HardforkConfig},
-    primitives::{types::BlockTrace, zk_trie::db::HashMapDb, Block, B256},
+    primitives::{types::BlockTrace, zk_trie::db::kv::HashMapDb, Block, B256},
 };
-use std::rc::Rc;
 use std::{cell::RefCell, path::PathBuf};
 use tiny_keccak::{Hasher, Keccak};
 use tokio::task::JoinSet;
@@ -75,12 +74,11 @@ impl RunFileCommand {
         }
 
         let fork_config = fork_config(traces[0].chain_id());
-        let (chunk_info, zktrie_db) = ChunkInfo::from_block_traces(&traces);
-        let zktrie_db = Rc::new(RefCell::new(zktrie_db));
+        let (chunk_info, mut zktrie_db) = ChunkInfo::from_block_traces(&traces);
 
         let tx_bytes_hasher = RefCell::new(Keccak::v256());
 
-        let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), zktrie_db.clone())
+        let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), &mut zktrie_db)
             .hardfork_config(fork_config)
             .with_hooks(&traces[0], |hooks| {
                 hooks.add_tx_rlp_handler(|_, rlp| {
@@ -94,7 +92,7 @@ impl RunFileCommand {
             executor.handle_block(trace)?;
         }
 
-        let post_state_root = executor.commit_changes(zktrie_db.clone())?;
+        let post_state_root = executor.commit_changes()?;
         if post_state_root != chunk_info.post_state_root() {
             bail!("post state root mismatch");
         }

diff --git a/crates/bin/src/utils.rs b/crates/bin/src/utils.rs
@@ -1,9 +1,8 @@
+use sbv::primitives::zk_trie::db::NodeDb;
 use sbv::{
     core::{EvmExecutorBuilder, HardforkConfig, VerificationError},
-    primitives::{zk_trie::db::HashMapDb, Block},
+    primitives::{zk_trie::db::kv::HashMapDb, Block},
 };
-use std::cell::RefCell;
-use std::rc::Rc;
 
 pub fn verify<T: Block + Clone>(
     l2_trace: T,
@@ -37,19 +36,19 @@ fn verify_inner<T: Block + Clone>(
         .build()
         .unwrap();
 
-    let zktrie_db = cycle_track!(
+    let mut zktrie_db = cycle_track!(
         {
-            let mut zktrie_db = HashMapDb::default();
+            let mut zktrie_db = NodeDb::new(HashMapDb::default());
             measure_duration_millis!(
                 build_zktrie_db_duration_milliseconds,
                 l2_trace.build_zktrie_db(&mut zktrie_db).unwrap()
             );
-            Rc::new(RefCell::new(zktrie_db))
+            zktrie_db
         },
         "build ZktrieState"
     );
 
-    let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), zktrie_db.clone())
+    let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), &mut zktrie_db)
         .hardfork_config(*fork_config)
         .build(&l2_trace)?;
 
@@ -66,7 +65,7 @@ fn verify_inner<T: Block + Clone>(
         update_metrics_counter!(verification_error);
         e
     })?;
-    let revm_root_after = executor.commit_changes(zktrie_db.clone())?;
+    let revm_root_after = executor.commit_changes()?;
 
     #[cfg(feature = "profiling")]
     if let Ok(report) = guard.report().build() {

diff --git a/crates/core/src/chunk.rs b/crates/core/src/chunk.rs
@@ -1,5 +1,6 @@
 use revm::primitives::B256;
-use sbv_primitives::{zk_trie::db::HashMapDb, Block};
+use sbv_primitives::zk_trie::db::NodeDb;
+use sbv_primitives::{zk_trie::db::kv::HashMapDb, Block};
 use tiny_keccak::{Hasher, Keccak};
 
 /// A chunk is a set of continuous blocks.
@@ -21,7 +22,7 @@ pub struct ChunkInfo {
 
 impl ChunkInfo {
     /// Construct by block traces
-    pub fn from_block_traces<T: Block>(traces: &[T]) -> (Self, HashMapDb) {
+    pub fn from_block_traces<T: Block>(traces: &[T]) -> (Self, NodeDb<HashMapDb>) {
         let chain_id = traces.first().unwrap().chain_id();
         let prev_state_root = traces
             .first()
@@ -40,7 +41,7 @@ impl ChunkInfo {
         let mut data_hash = B256::ZERO;
         data_hasher.finalize(&mut data_hash.0);
 
-        let mut zktrie_db = HashMapDb::default();
+        let mut zktrie_db = NodeDb::new(HashMapDb::default());
         for trace in traces.iter() {
             measure_duration_millis!(
                 build_zktrie_db_duration_milliseconds,
@@ -116,7 +117,6 @@ mod tests {
     use revm::primitives::b256;
     use sbv_primitives::types::BlockTrace;
     use std::cell::RefCell;
-    use std::rc::Rc;
 
     const TRACES_STR: [&str; 4] = [
         include_str!("../../../testdata/mainnet_blocks/8370400.json"),
@@ -138,12 +138,11 @@ mod tests {
         });
 
         let fork_config = HardforkConfig::default_from_chain_id(traces[0].chain_id());
-        let (chunk_info, zktrie_db) = ChunkInfo::from_block_traces(&traces);
-        let zktrie_db = Rc::new(RefCell::new(zktrie_db));
+        let (chunk_info, mut zktrie_db) = ChunkInfo::from_block_traces(&traces);
 
         let tx_bytes_hasher = RefCell::new(Keccak::v256());
 
-        let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), zktrie_db.clone())
+        let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), &mut zktrie_db)
             .hardfork_config(fork_config)
             .with_hooks(&traces[0], |hooks| {
                 hooks.add_tx_rlp_handler(|_, rlp| {
@@ -158,7 +157,7 @@ mod tests {
             executor.handle_block(trace).unwrap();
         }
 
-        let post_state_root = executor.commit_changes(zktrie_db.clone()).unwrap();
+        let post_state_root = executor.commit_changes().unwrap();
         assert_eq!(post_state_root, chunk_info.post_state_root);
         drop(executor); // drop executor to release Rc<Keccek>
 

diff --git a/crates/core/src/database.rs b/crates/core/src/database.rs
@@ -1,12 +1,13 @@
 use crate::error::DatabaseError;
-use once_cell::sync::Lazy;
 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::{KVDatabase, KVDatabaseItem},
+        db::kv::{KVDatabase, KVDatabaseItem},
         hash::{key_hasher::NoCacheHasher, poseidon::Poseidon},
         scroll_types::Account,
         trie::ZkTrie,
@@ -17,38 +18,38 @@ use std::{cell::RefCell, collections::HashMap, fmt};
 
 type Result<T, E = DatabaseError> = std::result::Result<T, E>;
 
-type StorageTrieLazyFn<Db> = Box<dyn FnOnce() -> ZkTrie<Poseidon, Db>>;
-type LazyStorageTrie<Db> = Lazy<ZkTrie<Poseidon, Db>, StorageTrieLazyFn<Db>>;
-
 /// A database that consists of account and storage information.
-pub struct EvmDatabase<CodeDb, ZkDb> {
+pub struct EvmDatabase<'a, CodeDb, ZkDb> {
     /// Map of code hash to bytecode.
     code_db: CodeDb,
-    /// The initial storage roots of accounts, used for after commit.
-    /// Need to be updated after zkTrie commit.
-    prev_storage_roots: RefCell<HashMap<Address, B256>>,
+    /// 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.
-    /// Need to invalidate before `update`, otherwise the trie root may be outdated.
-    storage_trie_refs: RefCell<HashMap<Address, LazyStorageTrie<ZkDb>>>,
+    /// Need to invalidate before `update`, otherwise the trie root may be outdated
+    storage_trie_caches: RefCell<HashMap<ZkHash, ZkTrie<Poseidon>>>,
     /// Current uncommitted zkTrie root based on the block trace.
     committed_zktrie_root: B256,
     /// The underlying zkTrie database.
-    zktrie_db: ZkDb,
+    pub(crate) zktrie_db: &'a mut NodeDb<ZkDb>,
     /// Current view of zkTrie database.
-    zktrie: ZkTrie<Poseidon, ZkDb>,
+    zktrie: ZkTrie<Poseidon>,
 }
 
-impl<CodeDb, Db> fmt::Debug for EvmDatabase<CodeDb, Db> {
+impl<CodeDb, Db> fmt::Debug for EvmDatabase<'_, CodeDb, Db> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.debug_struct("EvmDatabase")
             .field("committed_zktrie_root", &self.committed_zktrie_root)
             .finish()
     }
 }
 
-impl<CodeDb: KVDatabase, ZkDb: KVDatabase + Clone + 'static> EvmDatabase<CodeDb, ZkDb> {
+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, zktrie_db: ZkDb) -> Result<Self> {
+    pub fn new<T: Block>(
+        l2_trace: T,
+        mut code_db: 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);
@@ -60,43 +61,45 @@ impl<CodeDb: KVDatabase, ZkDb: KVDatabase + Clone + 'static> EvmDatabase<CodeDb,
 
         let committed_zktrie_root = l2_trace.root_before();
 
-        let zktrie = ZkTrie::new_with_root(zktrie_db.clone(), NoCacheHasher, committed_zktrie_root)
+        let zktrie = ZkTrie::new_with_root(zktrie_db, NoCacheHasher, committed_zktrie_root)
             .map_err(DatabaseError::zk_trie)?;
 
         Ok(EvmDatabase {
             code_db,
-            prev_storage_roots: Default::default(),
-            storage_trie_refs: Default::default(),
+            storage_root_caches: Default::default(),
+            storage_trie_caches: Default::default(),
             committed_zktrie_root,
             zktrie_db,
             zktrie,
         })
     }
 
-    /// Set the previous storage root of an account.
-    ///
-    /// Should be updated after commit.
-    #[inline]
-    pub(crate) fn set_prev_storage_root(
-        &self,
-        address: Address,
-        storage_root: B256,
-    ) -> Option<B256> {
-        self.prev_storage_roots
-            .borrow_mut()
-            .insert(address, storage_root)
-    }
-
     /// Get the previous storage root of an account.
     #[inline]
     pub(crate) fn prev_storage_root(&self, address: &Address) -> B256 {
-        self.prev_storage_roots
+        self.storage_root_caches
             .borrow()
             .get(address)
             .copied()
             .unwrap_or_default()
     }
 
+    #[inline]
+    pub(crate) fn update_storage_root_cache(&self, address: Address, storage_root: ZkTrie) {
+        let new_root = *storage_root.root().unwrap_ref();
+        let old = self
+            .storage_root_caches
+            .borrow_mut()
+            .insert(address, new_root);
+
+        let mut storage_trie_caches = self.storage_trie_caches.borrow_mut();
+        if let Some(old) = old {
+            storage_trie_caches.remove(&old);
+        }
+
+        storage_trie_caches.insert(new_root, storage_root);
+    }
+
     /// Get the committed zkTrie root.
     #[inline]
     pub(crate) fn committed_zktrie_root(&self) -> B256 {
@@ -125,11 +128,7 @@ impl<CodeDb: KVDatabase, ZkDb: KVDatabase + Clone + 'static> EvmDatabase<CodeDb,
         cycle_tracker_end!("insert CodeDB");
 
         self.zktrie = cycle_track!(
-            ZkTrie::new_with_root(
-                self.zktrie_db.clone(),
-                NoCacheHasher,
-                l2_trace.root_before(),
-            ),
+            ZkTrie::new_with_root(self.zktrie_db, NoCacheHasher, l2_trace.root_before(),),
             "ZkTrie::new_with_root"
         )
         .map_err(DatabaseError::zk_trie)?;
@@ -138,36 +137,23 @@ impl<CodeDb: KVDatabase, ZkDb: KVDatabase + Clone + 'static> EvmDatabase<CodeDb,
     }
 }
 
-impl<CodeDb: KVDatabase, ZkDb: KVDatabase + Clone + 'static> DatabaseRef
-    for EvmDatabase<CodeDb, ZkDb>
-{
+impl<CodeDb: KVDatabase, ZkDb: KVDatabase + 'static> DatabaseRef for EvmDatabase<'_, CodeDb, ZkDb> {
     type Error = DatabaseError;
 
     /// Get basic account information.
     fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>> {
         let Some(account) = measure_duration_micros!(
             zktrie_get_duration_microseconds,
-            self.zktrie.get::<Account, _>(address)
+            self.zktrie.get::<_, Account, _>(self.zktrie_db, address)
         )
         .map_err(DatabaseError::zk_trie)?
         else {
             return Ok(None);
         };
 
-        self.prev_storage_roots
+        self.storage_root_caches
             .borrow_mut()
-            .entry(address)
-            .or_insert(account.storage_root);
-        let zktrie_db = self.zktrie_db.clone();
-        self.storage_trie_refs
-            .borrow_mut()
-            .entry(address)
-            .or_insert_with(|| {
-                Lazy::new(Box::new(move || {
-                    ZkTrie::new_with_root(zktrie_db.clone(), NoCacheHasher, account.storage_root)
-                        .expect("storage trie associated with account not found")
-                }))
-            });
+            .insert(address, account.storage_root);
 
         let mut info = AccountInfo::from(account);
         info.code = self
@@ -201,32 +187,35 @@ impl<CodeDb: KVDatabase, ZkDb: KVDatabase + Clone + 'static> DatabaseRef
     /// Get storage value of address at index.
     fn storage_ref(&self, address: Address, index: U256) -> Result<U256> {
         dev_trace!("get storage of {:?} at index {:?}", address, index);
-        let mut storage_trie_refs = self.storage_trie_refs.borrow_mut();
-        let trie = storage_trie_refs
+        let storage_root = *self
+            .storage_root_caches
+            .borrow_mut()
             .entry(address)
             .or_insert_with_key(|address| {
-                let storage_root = measure_duration_micros!(
-                    zktrie_get_duration_microseconds,
-                    self.zktrie.get::<Account, _>(address)
-                )
-                .expect("unexpected zktrie error")
-                .map(|acc| acc.storage_root)
-                .unwrap_or_default();
+                self.zktrie
+                    .get::<_, Account, _>(self.zktrie_db, address)
+                    .expect("unexpected zktrie error")
+                    .map(|acc| acc.storage_root)
+                    .unwrap_or_default()
+            });
+
+        let mut storage_trie_caches = self.storage_trie_caches.borrow_mut();
+
+        let trie = storage_trie_caches
+            .entry(storage_root)
+            .or_insert_with_key(|storage_root| {
                 dev_debug!("storage root of {:?} is {:?}", address, storage_root);
 
-                let zktrie_db = self.zktrie_db.clone();
-                Lazy::new(Box::new(move || {
-                    ZkTrie::new_with_root(zktrie_db.clone(), NoCacheHasher, storage_root)
-                        .expect("storage trie associated with account not found")
-                }))
+                ZkTrie::new_with_root(self.zktrie_db, NoCacheHasher, *storage_root)
+                    .expect("storage trie associated with account not found")
             });
 
         #[cfg(debug_assertions)]
         {
             let current_root = trie.root().unwrap_ref();
             let expected_root = self
                 .zktrie
-                .get::<Account, _>(address)
+                .get::<_, Account, _>(self.zktrie_db, address)
                 .expect("unexpected zktrie error")
                 .map(|acc| acc.storage_root)
                 .unwrap_or_default();
@@ -235,7 +224,7 @@ impl<CodeDb: KVDatabase, ZkDb: KVDatabase + Clone + 'static> DatabaseRef
 
         Ok(measure_duration_micros!(
             zktrie_get_duration_microseconds,
-            trie.get::<U256, _>(index.to_be_bytes::<32>())
+            trie.get::<_, U256, _>(self.zktrie_db, index.to_be_bytes::<32>())
         )
         .map_err(DatabaseError::zk_trie)?
         .unwrap_or_default())