feat: keccak zktrie

Cargo.toml

@@ -21,11 +21,12 @@ tiny-keccak = "2.0"
 
 # dependencies from scroll-tech
 poseidon-bn254 = { git = "https://github.com/scroll-tech/poseidon-bn254", branch = "master", features = ["bn254"] }
-zktrie-ng = { git = "https://github.com/scroll-tech/zktrie-ng", branch = "master", features = ["scroll"] }
+zktrie-ng = { git = "https://github.com/scroll-tech/zktrie-ng", branch = "feat/keccak", features = ["scroll"] }
 
 # binary dependencies
 anyhow = "1.0"
 async-channel = "2.2"
+bincode = { version = "=2.0.0-rc.3", features = ["serde"] }
 clap = "4"
 env_logger = "0.9"
 futures = "0.3"
@@ -102,9 +103,9 @@ alloy-primitives = { git = "https://github.com/scroll-tech/alloy-core", branch =
 alloy-sol-types = {git = "https://github.com/scroll-tech/alloy-core", branch = "v0.8.10" }
 
 # for local development
-# [patch."https://github.com/scroll-tech/revm"]
-# revm = { path = "../revm/crates/revm" }
-# revm-primitives = { path = "../revm/crates/primitives" }
+#[patch."https://github.com/scroll-tech/revm"]
+#revm = { path = "../revm/crates/revm" }
+#revm-primitives = { path = "../revm/crates/primitives" }
 
 #[profile.release]
 #debug-assertions = true

crates/bin/src/commands/run_file.rs

@@ -1,18 +1,20 @@
 use crate::utils;
 use anyhow::{anyhow, bail};
 use clap::Args;
+use sbv::primitives::zk_trie::hash::keccak::Keccak;
+use sbv::primitives::zk_trie::hash::HashSchemeKind;
 use sbv::{
     core::{BlockExecutionResult, ChunkInfo, EvmExecutorBuilder, HardforkConfig},
     primitives::{
         types::{BlockTrace, LegacyStorageTrace},
-        zk_trie::db::kv::HashMapDb,
+        zk_trie::{db::kv::HashMapDb, hash::poseidon::Poseidon},
         Block, B256,
     },
 };
 use serde::Deserialize;
 use std::panic::catch_unwind;
 use std::path::PathBuf;
-use tiny_keccak::{Hasher, Keccak};
+use tiny_keccak::Hasher;
 use tokio::task::JoinSet;
 
 #[derive(Args)]
@@ -23,6 +25,9 @@ pub struct RunFileCommand {
     /// Chunk mode
     #[arg(short, long)]
     chunk_mode: bool,
+    /// Hash scheme
+    #[arg(long, value_enum, default_value_t = HashSchemeKind::Poseidon)]
+    hash_scheme: HashSchemeKind,
 }
 
 impl RunFileCommand {
@@ -44,7 +49,7 @@ impl RunFileCommand {
         let mut tasks = JoinSet::new();
 
         for path in self.path.into_iter() {
-            tasks.spawn(run_trace(path, fork_config));
+            tasks.spawn(run_trace(path, fork_config, self.hash_scheme));
         }
 
         while let Some(task) = tasks.join_next().await {
@@ -80,31 +85,51 @@ impl RunFileCommand {
         }
 
         let fork_config = fork_config(traces[0].chain_id());
-        let (chunk_info, mut zktrie_db) = ChunkInfo::from_block_traces(&traces);
+        let (chunk_info, mut zktrie_db) = ChunkInfo::from_block_traces(&traces, self.hash_scheme);
         let mut code_db = HashMapDb::default();
 
-        let mut tx_bytes_hasher = Keccak::v256();
+        let mut tx_bytes_hasher = tiny_keccak::Keccak::v256();
 
-        let mut executor = EvmExecutorBuilder::new(&mut code_db, &mut zktrie_db)
+        let builder = EvmExecutorBuilder::new(&mut code_db, &mut zktrie_db)
             .hardfork_config(fork_config)
-            .chain_id(traces[0].chain_id())
-            .build(traces[0].root_before())?;
-        for trace in traces.iter() {
-            executor.insert_codes(trace)?;
-        }
-
-        for trace in traces.iter() {
-            let BlockExecutionResult { tx_rlps, .. } = executor.handle_block(trace)?;
-            for tx_rlp in tx_rlps {
-                tx_bytes_hasher.update(&tx_rlp);
+            .chain_id(traces[0].chain_id());
+        let post_state_root = match self.hash_scheme {
+            HashSchemeKind::Poseidon => {
+                let mut executor = builder
+                    .hash_scheme(Poseidon)
+                    .build(traces[0].root_before())?;
+                for trace in traces.iter() {
+                    executor.insert_codes(trace)?;
+                }
+
+                for trace in traces.iter() {
+                    let BlockExecutionResult { tx_rlps, .. } = executor.handle_block(trace)?;
+                    for tx_rlp in tx_rlps {
+                        tx_bytes_hasher.update(&tx_rlp);
+                    }
+                }
+
+                executor.commit_changes()?
             }
-        }
-
-        let post_state_root = executor.commit_changes()?;
+            HashSchemeKind::Keccak => {
+                let mut executor = builder.hash_scheme(Keccak).build(traces[0].root_before())?;
+                for trace in traces.iter() {
+                    executor.insert_codes(trace)?;
+                }
+
+                for trace in traces.iter() {
+                    let BlockExecutionResult { tx_rlps, .. } = executor.handle_block(trace)?;
+                    for tx_rlp in tx_rlps {
+                        tx_bytes_hasher.update(&tx_rlp);
+                    }
+                }
+
+                executor.commit_changes()?
+            }
+        };
         if post_state_root != chunk_info.post_state_root() {
             bail!("post state root mismatch");
         }
-        drop(executor);
 
         let mut tx_bytes_hash = B256::ZERO;
         tx_bytes_hasher.finalize(&mut tx_bytes_hash.0);
@@ -148,22 +173,24 @@ fn deserialize_may_wrapped<'de, T: Deserialize<'de>>(trace: &'de str) -> anyhow:
 async fn run_trace(
     path: PathBuf,
     fork_config: impl Fn(u64) -> HardforkConfig,
+    hash_scheme: HashSchemeKind,
 ) -> anyhow::Result<()> {
     let trace = read_block_trace(&path).await?;
     let fork_config = fork_config(trace.chain_id());
-    if let Err(e) =
-        tokio::task::spawn_blocking(move || catch_unwind(|| utils::verify(&trace, &fork_config)))
-            .await?
-            .map_err(|e| {
-                e.downcast_ref::<&str>()
+    if let Err(e) = tokio::task::spawn_blocking(move || {
+        catch_unwind(|| utils::verify(&trace, &fork_config, hash_scheme))
+    })
+    .await?
+    .map_err(|e| {
+        e.downcast_ref::<&str>()
+            .map(|s| anyhow!("task panics with: {s}"))
+            .or_else(|| {
+                e.downcast_ref::<String>()
                     .map(|s| anyhow!("task panics with: {s}"))
-                    .or_else(|| {
-                        e.downcast_ref::<String>()
-                            .map(|s| anyhow!("task panics with: {s}"))
-                    })
-                    .unwrap_or_else(|| anyhow!("task panics"))
             })
-            .and_then(|r| r.map_err(anyhow::Error::from))
+            .unwrap_or_else(|| anyhow!("task panics"))
+    })
+    .and_then(|r| r.map_err(anyhow::Error::from))
     {
         dev_error!(
             "Error occurs when verifying block ({}): {:?}",

crates/bin/src/commands/run_rpc.rs

@@ -3,6 +3,7 @@ use alloy::providers::{Provider, ProviderBuilder};
 use clap::Args;
 use futures::future::OptionFuture;
 use sbv::primitives::types::LegacyStorageTrace;
+use sbv::primitives::zk_trie::hash::HashSchemeKind;
 use sbv::{
     core::HardforkConfig,
     primitives::{types::BlockTrace, Block},
@@ -115,10 +116,12 @@ impl RunRpcCommand {
                         l2_trace.block_hash()
                     );
 
-                    tokio::task::spawn_blocking(move || utils::verify(&l2_trace, &fork_config))
-                        .await
-                        .expect("failed to spawn blocking task")
-                        .map_err(|e| (block_number, e.into()))?;
+                    tokio::task::spawn_blocking(move || {
+                        utils::verify(&l2_trace, &fork_config, HashSchemeKind::Poseidon)
+                    })
+                    .await
+                    .expect("failed to spawn blocking task")
+                    .map_err(|e| (block_number, e.into()))?;
                 }
                 Ok::<_, (u64, anyhow::Error)>(())
             });

crates/bin/src/utils.rs

@@ -1,22 +1,26 @@
 use sbv::primitives::zk_trie::db::NodeDb;
+use sbv::primitives::zk_trie::hash::keccak::Keccak;
+use sbv::primitives::zk_trie::hash::HashSchemeKind;
 use sbv::{
     core::{EvmExecutorBuilder, HardforkConfig, VerificationError},
-    primitives::{zk_trie::db::kv::HashMapDb, Block},
+    primitives::{zk_trie::db::kv::HashMapDb, zk_trie::hash::poseidon::Poseidon, Block},
 };
 
 pub fn verify<T: Block + Clone>(
     l2_trace: T,
     fork_config: &HardforkConfig,
+    hash_scheme: HashSchemeKind,
 ) -> Result<(), VerificationError> {
     measure_duration_millis!(
         total_block_verification_duration_milliseconds,
-        verify_inner(l2_trace, fork_config)
+        verify_inner(l2_trace, fork_config, hash_scheme)
     )
 }
 
 fn verify_inner<T: Block + Clone>(
     l2_trace: T,
     fork_config: &HardforkConfig,
+    hash_scheme: HashSchemeKind,
 ) -> Result<(), VerificationError> {
     dev_trace!("{l2_trace:#?}");
     let root_before = l2_trace.root_before();
@@ -42,35 +46,62 @@ fn verify_inner<T: Block + Clone>(
             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()
+                l2_trace
+                    .build_zktrie_db(&mut zktrie_db, hash_scheme)
+                    .unwrap()
             );
             zktrie_db
         },
         "build ZktrieState"
     );
     let mut code_db = HashMapDb::default();
 
-    let mut executor = EvmExecutorBuilder::new(&mut code_db, &mut zktrie_db)
+    let builder = EvmExecutorBuilder::new(&mut code_db, &mut zktrie_db)
         .hardfork_config(*fork_config)
-        .chain_id(l2_trace.chain_id())
-        .build(root_before)?;
+        .chain_id(l2_trace.chain_id());
 
-    executor.insert_codes(&l2_trace)?;
+    let revm_root_after = match hash_scheme {
+        HashSchemeKind::Poseidon => {
+            let mut executor = builder.hash_scheme(Poseidon).build(root_before)?;
 
-    // TODO: change to Result::inspect_err when sp1 toolchain >= 1.76
-    #[allow(clippy::map_identity)]
-    #[allow(clippy::manual_inspect)]
-    executor.handle_block(&l2_trace).map_err(|e| {
-        dev_error!(
-            "Error occurs when executing block #{}({:?}): {e:?}",
-            l2_trace.number(),
-            l2_trace.block_hash()
-        );
+            executor.insert_codes(&l2_trace)?;
 
-        update_metrics_counter!(verification_error);
-        e
-    })?;
-    let revm_root_after = executor.commit_changes()?;
+            // TODO: change to Result::inspect_err when sp1 toolchain >= 1.76
+            #[allow(clippy::map_identity)]
+            #[allow(clippy::manual_inspect)]
+            executor.handle_block(&l2_trace).map_err(|e| {
+                dev_error!(
+                    "Error occurs when executing block #{}({:?}): {e:?}",
+                    l2_trace.number(),
+                    l2_trace.block_hash()
+                );
+
+                update_metrics_counter!(verification_error);
+                e
+            })?;
+            executor.commit_changes()?
+        }
+        HashSchemeKind::Keccak => {
+            let mut executor = builder.hash_scheme(Keccak).build(root_before)?;
+
+            executor.insert_codes(&l2_trace)?;
+
+            // TODO: change to Result::inspect_err when sp1 toolchain >= 1.76
+            #[allow(clippy::map_identity)]
+            #[allow(clippy::manual_inspect)]
+            executor.handle_block(&l2_trace).map_err(|e| {
+                dev_error!(
+                    "Error occurs when executing block #{}({:?}): {e:?}",
+                    l2_trace.number(),
+                    l2_trace.block_hash()
+                );
+
+                update_metrics_counter!(verification_error);
+                e
+            })?;
+            executor.commit_changes()?
+        }
+    };
 
     #[cfg(feature = "profiling")]
     if let Ok(report) = guard.report().build() {