fix ChunkInfo

Author: lightsing

crates/bin/src/commands/run_file.rs

@@ -1,6 +1,8 @@
 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::{
@@ -12,7 +14,7 @@ use sbv::{
 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,32 +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())
-            .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);
+            .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);
@@ -149,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,4 +1,5 @@
 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},
@@ -8,16 +9,18 @@ use sbv::{
 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();
@@ -44,7 +47,7 @@ fn verify_inner<T: Block + Clone>(
             measure_duration_millis!(
                 build_zktrie_db_duration_milliseconds,
                 l2_trace
-                    .build_zktrie_db(&mut zktrie_db, HashSchemeKind::Poseidon)
+                    .build_zktrie_db(&mut zktrie_db, hash_scheme)
                     .unwrap()
             );
             zktrie_db
@@ -53,28 +56,52 @@ fn verify_inner<T: Block + Clone>(
     );
     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())
-        .hash_scheme(Poseidon)
-        .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() {

crates/core/src/chunk.rs

@@ -23,7 +23,10 @@ pub struct ChunkInfo {
 
 impl ChunkInfo {
     /// Construct by block traces
-    pub fn from_block_traces<T: Block>(traces: &[T]) -> (Self, NodeDb<HashMapDb>) {
+    pub fn from_block_traces<T: Block>(
+        traces: &[T],
+        hash_scheme: HashSchemeKind,
+    ) -> (Self, NodeDb<HashMapDb>) {
         let chain_id = traces.first().unwrap().chain_id();
         let prev_state_root = traces
             .first()
@@ -49,9 +52,7 @@ impl ChunkInfo {
         for trace in traces.iter() {
             measure_duration_millis!(
                 build_zktrie_db_duration_milliseconds,
-                trace
-                    .build_zktrie_db(&mut zktrie_db, HashSchemeKind::Poseidon)
-                    .unwrap()
+                trace.build_zktrie_db(&mut zktrie_db, hash_scheme).unwrap()
             );
         }
         cycle_tracker_end!("Block::build_zktrie_db");