feat(rpc): debug_executionWitness (#9249)

Co-authored-by: Roman Krasiuk <rokrassyuk@gmail.com>

crates/rpc/rpc-api/src/debug.rs

@@ -7,6 +7,7 @@ use reth_rpc_types::{
     },
     Bundle, RichBlock, StateContext, TransactionRequest,
 };
+use std::collections::HashMap;
 
 /// Debug rpc interface.
 #[cfg_attr(not(feature = "client"), rpc(server, namespace = "debug"))]
@@ -132,6 +133,18 @@ pub trait DebugApi {
         opts: Option<GethDebugTracingCallOptions>,
     ) -> RpcResult<Vec<Vec<GethTrace>>>;
 
+    /// The `debug_executionWitness` method allows for re-execution of a block with the purpose of
+    /// generating an execution witness. The witness comprises of a map of all hashed trie nodes
+    /// to their preimages that were required during the execution of the block, including during
+    /// state root recomputation.
+    ///
+    /// The first and only argument is the block number or block hash.
+    #[method(name = "executionWitness")]
+    async fn debug_execution_witness(
+        &self,
+        block: BlockNumberOrTag,
+    ) -> RpcResult<HashMap<B256, Bytes>>;
+
     /// Sets the logging backtrace location. When a backtrace location is set and a log message is
     /// emitted at that location, the stack of the goroutine executing the log statement will
     /// be printed to stderr.

crates/rpc/rpc/Cargo.toml

@@ -34,6 +34,7 @@ reth-rpc-eth-types.workspace = true
 reth-rpc-server-types.workspace = true
 reth-node-api.workspace = true
 reth-network-types.workspace = true
+reth-trie.workspace = true
 
 # eth
 alloy-dyn-abi.workspace = true

crates/rpc/rpc/src/debug.rs

@@ -1,16 +1,14 @@
-use std::sync::Arc;
-
 use alloy_rlp::{Decodable, Encodable};
 use async_trait::async_trait;
 use jsonrpsee::core::RpcResult;
 use reth_chainspec::EthereumHardforks;
-use reth_evm::ConfigureEvmEnv;
+use reth_evm::{system_calls::pre_block_beacon_root_contract_call, ConfigureEvmEnv};
 use reth_primitives::{
     Address, Block, BlockId, BlockNumberOrTag, Bytes, TransactionSignedEcRecovered, B256, U256,
 };
 use reth_provider::{
-    BlockReaderIdExt, ChainSpecProvider, EvmEnvProvider, HeaderProvider, StateProviderFactory,
-    TransactionVariant,
+    BlockReaderIdExt, ChainSpecProvider, EvmEnvProvider, HeaderProvider, StateProofProvider,
+    StateProviderFactory, TransactionVariant,
 };
 use reth_revm::database::StateProviderDatabase;
 use reth_rpc_api::DebugApiServer;
@@ -29,14 +27,18 @@ use reth_rpc_types::{
     BlockError, Bundle, RichBlock, StateContext, TransactionRequest,
 };
 use reth_tasks::pool::BlockingTaskGuard;
+use reth_trie::{HashedPostState, HashedStorage};
 use revm::{
-    db::CacheDB,
+    db::{states::bundle_state::BundleRetention, CacheDB},
     primitives::{db::DatabaseCommit, BlockEnv, CfgEnvWithHandlerCfg, Env, EnvWithHandlerCfg},
+    StateBuilder,
 };
 use revm_inspectors::tracing::{
     js::{JsInspector, TransactionContext},
     FourByteInspector, MuxInspector, TracingInspector, TracingInspectorConfig,
 };
+use revm_primitives::{keccak256, HashMap};
+use std::sync::Arc;
 use tokio::sync::{AcquireError, OwnedSemaphorePermit};
 
 /// `debug` API implementation.
@@ -550,6 +552,103 @@ where
             .await
     }
 
+    /// The `debug_executionWitness` method allows for re-execution of a block with the purpose of
+    /// generating an execution witness. The witness comprises of a map of all hashed trie nodes
+    /// to their preimages that were required during the execution of the block, including during
+    /// state root recomputation.
+    pub async fn debug_execution_witness(
+        &self,
+        block_id: BlockNumberOrTag,
+    ) -> Result<HashMap<B256, Bytes>, Eth::Error> {
+        let ((cfg, block_env, _), maybe_block) = futures::try_join!(
+            self.inner.eth_api.evm_env_at(block_id.into()),
+            self.inner.eth_api.block_with_senders(block_id.into()),
+        )?;
+        let block = maybe_block.ok_or(EthApiError::UnknownBlockNumber)?;
+
+        let this = self.clone();
+
+        self.inner
+            .eth_api
+            .spawn_with_state_at_block(block.parent_hash.into(), move |state| {
+                let evm_config = Call::evm_config(this.eth_api()).clone();
+                let mut db = StateBuilder::new()
+                    .with_database(StateProviderDatabase::new(state))
+                    .with_bundle_update()
+                    .build();
+
+                pre_block_beacon_root_contract_call(
+                    &mut db,
+                    &evm_config,
+                    &this.inner.provider.chain_spec(),
+                    &cfg,
+                    &block_env,
+                    block.timestamp,
+                    block.number,
+                    block.parent_beacon_block_root,
+                )
+                .map_err(|err| EthApiError::Internal(err.into()))?;
+
+                // Re-execute all of the transactions in the block to load all touched accounts into
+                // the cache DB.
+                for tx in block.raw_transactions() {
+                    let tx_envelope = TransactionSignedEcRecovered::decode(&mut tx.as_ref())
+                        .map_err(|_| EthApiError::FailedToDecodeSignedTransaction)?;
+                    let env = EnvWithHandlerCfg {
+                        env: Env::boxed(
+                            cfg.cfg_env.clone(),
+                            block_env.clone(),
+                            evm_config.tx_env(&tx_envelope),
+                        ),
+                        handler_cfg: cfg.handler_cfg,
+                    };
+
+                    let (res, _) = this.inner.eth_api.transact(&mut db, env)?;
+                    db.commit(res.state);
+                }
+
+                // Merge all state transitions
+                db.merge_transitions(BundleRetention::Reverts);
+
+                // Take the bundle state
+                let bundle_state = db.take_bundle();
+
+                // Grab all account proofs for the data accessed during block execution.
+                //
+                // Note: We grab *all* accounts in the cache here, as the `BundleState` prunes
+                // referenced accounts + storage slots.
+                let mut hashed_state = HashedPostState::from_bundle_state(&bundle_state.state);
+                for (address, account) in db.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 {
+                        for (slot, value) in account.storage {
+                            let hashed_slot = keccak256(B256::from(slot));
+                            storage.storage.insert(hashed_slot, value);
+                        }
+                    }
+                }
+
+                // Generate an execution witness for the aggregated state of accessed accounts.
+                // Destruct the cache database to retrieve the state provider.
+                let state_provider = db.database.into_inner();
+                let witness = state_provider
+                    .witness(HashedPostState::default(), hashed_state)
+                    .map_err(Into::into)?;
+                Ok(witness)
+            })
+            .await
+    }
+
     /// Executes the configured transaction with the environment on the given database.
     ///
     /// Returns the trace frame and the state that got updated after executing the transaction.
@@ -806,6 +905,15 @@ where
             .map_err(Into::into)
     }
 
+    /// Handler for `debug_executionWitness`
+    async fn debug_execution_witness(
+        &self,
+        block: BlockNumberOrTag,
+    ) -> RpcResult<HashMap<B256, Bytes>> {
+        let _permit = self.acquire_trace_permit().await;
+        Self::debug_execution_witness(self, block).await.map_err(Into::into)
+    }
+
     /// Handler for `debug_traceCall`
     async fn debug_trace_call(
         &self,