chore: extract witness recorder helper type (#12566)

crates/revm/Cargo.toml

@@ -44,6 +44,7 @@ std = [
     "alloy-consensus/std",
     "reth-primitives-traits/std",
 ]
+witness = ["dep:reth-trie"]
 test-utils = [
 	"dep:reth-trie",
 	"reth-primitives/test-utils",

crates/revm/src/lib.rs

@@ -29,3 +29,7 @@ pub use revm::{self, *};
 
 /// Either type for flexible usage of different database types in the same context.
 pub mod either;
+
+/// Helper types for execution witness generation.
+#[cfg(feature = "witness")]
+pub mod witness;

crates/revm/src/witness.rs

@@ -0,0 +1,76 @@
+use alloy_primitives::{keccak256, map::B256HashMap, Bytes, B256};
+use reth_trie::{HashedPostState, HashedStorage};
+use revm::State;
+
+/// Tracks state changes during execution.
+#[derive(Debug, Clone, Default)]
+pub struct ExecutionWitnessRecord {
+    /// Records all state changes
+    pub hashed_state: HashedPostState,
+    /// Map of all contract codes (created / accessed) to their preimages that were required during
+    /// the execution of the block, including during state root recomputation.
+    ///
+    /// `keccak(bytecodes) => bytecodes`
+    pub codes: B256HashMap<Bytes>,
+    /// Map of all hashed account and storage keys (addresses and slots) to their preimages
+    /// (unhashed account addresses and storage slots, respectively) that were required during
+    /// the execution of the block. during the execution of the block.
+    ///
+    /// `keccak(address|slot) => address|slot`
+    pub keys: B256HashMap<Bytes>,
+}
+
+impl ExecutionWitnessRecord {
+    /// Records the state after execution.
+    pub fn record_executed_state<DB>(&mut self, statedb: &State<DB>) {
+        self.codes = statedb
+            .cache
+            .contracts
+            .iter()
+            .map(|(hash, code)| (*hash, code.original_bytes()))
+            .chain(
+                // cache state does not have all the contracts, especially when
+                // a contract is created within the block
+                // the contract only exists in bundle state, therefore we need
+                // to include them as well
+                statedb
+                    .bundle_state
+                    .contracts
+                    .iter()
+                    .map(|(hash, code)| (*hash, code.original_bytes())),
+            )
+            .collect();
+
+        for (address, account) in &statedb.cache.accounts {
+            let hashed_address = keccak256(address);
+            self.hashed_state
+                .accounts
+                .insert(hashed_address, account.account.as_ref().map(|a| a.info.clone().into()));
+
+            let storage = self
+                .hashed_state
+                .storages
+                .entry(hashed_address)
+                .or_insert_with(|| HashedStorage::new(account.status.was_destroyed()));
+
+            if let Some(account) = &account.account {
+                self.keys.insert(hashed_address, address.to_vec().into());
+
+                for (slot, value) in &account.storage {
+                    let slot = B256::from(*slot);
+                    let hashed_slot = keccak256(slot);
+                    storage.storage.insert(hashed_slot, *value);
+
+                    self.keys.insert(hashed_slot, slot.into());
+                }
+            }
+        }
+    }
+
+    /// Creates the record from the state after execution.
+    pub fn from_executed_state<DB>(state: &State<DB>) -> Self {
+        let mut record = Self::default();
+        record.record_executed_state(state);
+        record
+    }
+}

crates/rpc/rpc/Cargo.toml

@@ -23,7 +23,7 @@ reth-provider.workspace = true
 reth-transaction-pool.workspace = true
 reth-network-api.workspace = true
 reth-rpc-engine-api.workspace = true
-reth-revm.workspace = true
+reth-revm = { workspace = true, features = ["witness"] }
 reth-tasks = { workspace = true, features = ["rayon"] }
 reth-consensus-common.workspace = true
 reth-rpc-types-compat.workspace = true

crates/rpc/rpc/src/debug.rs

@@ -23,7 +23,7 @@ use reth_provider::{
     BlockReaderIdExt, ChainSpecProvider, HeaderProvider, StateProofProvider, StateProviderFactory,
     TransactionVariant,
 };
-use reth_revm::database::StateProviderDatabase;
+use reth_revm::{database::StateProviderDatabase, witness::ExecutionWitnessRecord};
 use reth_rpc_api::DebugApiServer;
 use reth_rpc_eth_api::{
     helpers::{EthApiSpec, EthTransactions, TraceExt},
@@ -32,7 +32,6 @@ use reth_rpc_eth_api::{
 use reth_rpc_eth_types::{EthApiError, StateCacheDb};
 use reth_rpc_server_types::{result::internal_rpc_err, ToRpcResult};
 use reth_tasks::pool::BlockingTaskGuard;
-use reth_trie::{HashedPostState, HashedStorage};
 use revm::{
     db::{CacheDB, State},
     primitives::{db::DatabaseCommit, BlockEnv, CfgEnvWithHandlerCfg, Env, EnvWithHandlerCfg},
@@ -40,7 +39,6 @@ use revm::{
 use revm_inspectors::tracing::{
     FourByteInspector, MuxInspector, TracingInspector, TracingInspectorConfig, TransactionContext,
 };
-use revm_primitives::{keccak256, HashMap};
 use std::sync::Arc;
 use tokio::sync::{AcquireError, OwnedSemaphorePermit};
 
@@ -613,60 +611,19 @@ where
                 let db = StateProviderDatabase::new(&state_provider);
                 let block_executor = this.inner.block_executor.executor(db);
 
-                let mut hashed_state = HashedPostState::default();
+                let mut witness_record = ExecutionWitnessRecord::default();
-                let mut keys = HashMap::default();
-                let mut codes = HashMap::default();
 
                 let _ = block_executor
                     .execute_with_state_closure(
                         (&(*block).clone().unseal(), block.difficulty).into(),
                         |statedb: &State<_>| {
-                            codes = statedb
+                            witness_record.record_executed_state(statedb);
-                                .cache
-                                .contracts
-                                .iter()
-                                .map(|(hash, code)| (*hash, code.original_bytes()))
-                                .chain(
-                                    // cache state does not have all the contracts, especially when
-                                    // a contract is created within the block
-                                    // the contract only exists in bundle state, therefore we need
-                                    // to include them as well
-                                    statedb
-                                        .bundle_state
-                                        .contracts
-                                        .iter()
-                                        .map(|(hash, code)| (*hash, code.original_bytes())),
-                                )
-                                .collect();
-
-                            for (address, account) in &statedb.cache.accounts {
-                                let hashed_address = keccak256(address);
-                                hashed_state.accounts.insert(
-                                    hashed_address,
-                                    account.account.as_ref().map(|a| a.info.clone().into()),
-                                );
-
-                                let storage =
-                                    hashed_state.storages.entry(hashed_address).or_insert_with(
-                                        || HashedStorage::new(account.status.was_destroyed()),
-                                    );
-
-                                if let Some(account) = &account.account {
-                                    keys.insert(hashed_address, address.to_vec().into());
-
-                                    for (slot, value) in &account.storage {
-                                        let slot = B256::from(*slot);
-                                        let hashed_slot = keccak256(slot);
-                                        storage.storage.insert(hashed_slot, *value);
-
-                                        keys.insert(hashed_slot, slot.into());
-                                    }
-                                }
-                            }
                         },
                     )
                     .map_err(|err| EthApiError::Internal(err.into()))?;
 
+                let ExecutionWitnessRecord { hashed_state, codes, keys } = witness_record;
+
                 let state =
                     state_provider.witness(Default::default(), hashed_state).map_err(Into::into)?;
                 Ok(ExecutionWitness { state: state.into_iter().collect(), codes, keys })