feat: use zktrie-ng (#54)

* use zktrie-ng

* fix

* clean legacy code

* add more metrics

* range adjustment

* update grafana

* fix

* update

* fix

Author: lightsing

Cargo.lock

@@ -21,7 +21,7 @@ tiny-keccak = "2.0"
 
 # dependencies from scroll-tech
 poseidon-bn254 = { git = "https://github.com/scroll-tech/poseidon-bn254", branch = "master", features = ["bn254"] }
-zktrie = { git = "https://github.com/scroll-tech/zktrie.git", branch = "main", features= ["rs_zktrie"] }
+zktrie-ng = { git = "https://github.com/scroll-tech/zktrie-ng", branch = "master", features = ["scroll"] }
 
 # binary dependencies
 anyhow = "1.0"

Cargo.toml

@@ -3,8 +3,9 @@ use anyhow::bail;
 use clap::Args;
 use sbv::{
     core::{ChunkInfo, EvmExecutorBuilder, HardforkConfig},
-    primitives::{types::BlockTrace, Block, B256},
+    primitives::{types::BlockTrace, zk_trie::db::HashMapDb, Block, B256},
 };
+use std::rc::Rc;
 use std::{cell::RefCell, path::PathBuf};
 use tiny_keccak::{Hasher, Keccak};
 use tokio::task::JoinSet;
@@ -75,25 +76,25 @@ 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 tx_bytes_hasher = RefCell::new(Keccak::v256());
 
-        let mut executor = EvmExecutorBuilder::new(zktrie_db.clone())
+        let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), zktrie_db.clone())
             .hardfork_config(fork_config)
-            .with_execute_hooks(|hooks| {
+            .with_hooks(&traces[0], |hooks| {
                 hooks.add_tx_rlp_handler(|_, rlp| {
                     tx_bytes_hasher.borrow_mut().update(rlp);
                 });
-            })
-            .build(&traces[0])?;
+            })?;
         executor.handle_block(&traces[0])?;
 
         for trace in traces[1..].iter() {
             executor.update_db(trace)?;
             executor.handle_block(trace)?;
         }
 
-        let post_state_root = executor.commit_changes(&zktrie_db);
+        let post_state_root = executor.commit_changes(zktrie_db.clone())?;
         if post_state_root != chunk_info.post_state_root() {
             bail!("post state root mismatch");
         }

crates/bin/src/commands/run_file.rs

@@ -5,7 +5,6 @@ extern crate sbv;
 
 use clap::Parser;
 use sbv::core::HardforkConfig;
-use sbv::primitives::init_hash_scheme;
 
 #[cfg(feature = "dev")]
 use tracing_subscriber::EnvFilter;
@@ -43,8 +42,6 @@ async fn main() -> anyhow::Result<()> {
         )
         .init();
 
-    init_hash_scheme();
-
     let cmd = Cli::parse();
 
     #[cfg(feature = "metrics")]

crates/bin/src/main.rs

@@ -1,8 +1,8 @@
-use sbv::primitives::zk_trie::ZkMemoryDb;
 use sbv::{
     core::{EvmExecutorBuilder, HardforkConfig, VerificationError},
-    primitives::Block,
+    primitives::{zk_trie::db::HashMapDb, Block},
 };
+use std::cell::RefCell;
 use std::rc::Rc;
 
 pub fn verify<T: Block + Clone>(
@@ -39,17 +39,17 @@ fn verify_inner<T: Block + Clone>(
 
     let zktrie_db = cycle_track!(
         {
-            let mut zktrie_db = ZkMemoryDb::new();
+            let mut zktrie_db = HashMapDb::default();
             measure_duration_millis!(
                 build_zktrie_db_duration_milliseconds,
-                l2_trace.build_zktrie_db(&mut zktrie_db)
+                l2_trace.build_zktrie_db(&mut zktrie_db).unwrap()
             );
-            Rc::new(zktrie_db)
+            Rc::new(RefCell::new(zktrie_db))
         },
         "build ZktrieState"
     );
 
-    let mut executor = EvmExecutorBuilder::new(zktrie_db.clone())
+    let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), 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(&zktrie_db);
+    let revm_root_after = executor.commit_changes(zktrie_db.clone())?;
 
     #[cfg(feature = "profiling")]
     if let Ok(report) = guard.report().build() {

crates/bin/src/utils.rs

@@ -33,7 +33,7 @@ tracing-subscriber.workspace = true
 [features]
 debug-account = ["sbv-utils/debug-account"]
 debug-storage = ["sbv-utils/debug-storage"]
-dev = ["sbv-utils/dev"]
+dev = ["sbv-primitives/dev", "sbv-utils/dev"]
 metrics = ["sbv-utils/metrics"]
 
 # sp1 related

crates/core/Cargo.toml

@@ -1,6 +1,5 @@
 use revm::primitives::B256;
-use sbv_primitives::{zk_trie::ZkMemoryDb, Block};
-use std::rc::Rc;
+use sbv_primitives::{zk_trie::db::HashMapDb, Block};
 use tiny_keccak::{Hasher, Keccak};
 
 /// A chunk is a set of continuous blocks.
@@ -22,7 +21,7 @@ pub struct ChunkInfo {
 
 impl ChunkInfo {
     /// Construct by block traces
-    pub fn from_block_traces<T: Block>(traces: &[T]) -> (Self, Rc<ZkMemoryDb>) {
+    pub fn from_block_traces<T: Block>(traces: &[T]) -> (Self, HashMapDb) {
         let chain_id = traces.first().unwrap().chain_id();
         let prev_state_root = traces
             .first()
@@ -41,14 +40,13 @@ impl ChunkInfo {
         let mut data_hash = B256::ZERO;
         data_hasher.finalize(&mut data_hash.0);
 
-        let mut zktrie_db = ZkMemoryDb::new();
+        let mut zktrie_db = HashMapDb::default();
         for trace in traces.iter() {
             measure_duration_millis!(
                 build_zktrie_db_duration_milliseconds,
-                trace.build_zktrie_db(&mut zktrie_db)
+                trace.build_zktrie_db(&mut zktrie_db).unwrap()
             );
         }
-        let zktrie_db = Rc::new(zktrie_db);
 
         let info = ChunkInfo {
             chain_id,
@@ -118,6 +116,7 @@ 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"),
@@ -140,17 +139,17 @@ 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 tx_bytes_hasher = RefCell::new(Keccak::v256());
 
-        let mut executor = EvmExecutorBuilder::new(zktrie_db.clone())
+        let mut executor = EvmExecutorBuilder::new(HashMapDb::default(), zktrie_db.clone())
             .hardfork_config(fork_config)
-            .with_execute_hooks(|hooks| {
+            .with_hooks(&traces[0], |hooks| {
                 hooks.add_tx_rlp_handler(|_, rlp| {
                     tx_bytes_hasher.borrow_mut().update(rlp);
                 });
             })
-            .build(&traces[0])
             .unwrap();
         executor.handle_block(&traces[0]).unwrap();
 
@@ -159,7 +158,7 @@ mod tests {
             executor.handle_block(trace).unwrap();
         }
 
-        let post_state_root = executor.commit_changes(&zktrie_db);
+        let post_state_root = executor.commit_changes(zktrie_db.clone()).unwrap();
         assert_eq!(post_state_root, chunk_info.post_state_root);
         drop(executor); // drop executor to release Rc<Keccek>