feat(stateless): make witness generation conform to the draft specs

crates/chain-state/src/in_memory.rs

@@ -1169,6 +1169,7 @@ mod tests {
             &self,
             _input: TrieInput,
             _target: HashedPostState,
+            _mode: reth_trie::ExecutionWitnessMode,
         ) -> ProviderResult<Vec<Bytes>> {
             Ok(Vec::default())
         }

crates/chain-state/src/memory_overlay.rs

@@ -197,9 +197,14 @@ impl<N: NodePrimitives> StateProofProvider for MemoryOverlayStateProviderRef<'_,
         self.historical.multiproof(input, targets)
     }
 
-    fn witness(&self, mut input: TrieInput, target: HashedPostState) -> ProviderResult<Vec<Bytes>> {
+    fn witness(
+        &self,
+        mut input: TrieInput,
+        target: HashedPostState,
+        mode: reth_trie::ExecutionWitnessMode,
+    ) -> ProviderResult<Vec<Bytes>> {
         input.prepend_self(self.trie_input().clone());
-        self.historical.witness(input, target)
+        self.historical.witness(input, target, mode)
     }
 }
 

crates/engine/execution-cache/src/cached_state.rs

@@ -566,8 +566,9 @@ impl<S: StateProofProvider, const PREWARM: bool> StateProofProvider
         &self,
         input: TrieInput,
         target: HashedPostState,
+        mode: reth_trie::ExecutionWitnessMode,
     ) -> ProviderResult<Vec<alloy_primitives::Bytes>> {
-        self.state_provider.witness(input, target)
+        self.state_provider.witness(input, target, mode)
     }
 }
 

crates/engine/invalid-block-hooks/src/witness.rs

@@ -160,7 +160,11 @@ fn generate(
     hashed_state: reth_trie::HashedPostState,
     state_provider: Box<dyn StateProvider>,
 ) -> eyre::Result<ExecutionWitness> {
-    let state = state_provider.witness(Default::default(), hashed_state)?;
+    let state = state_provider.witness(
+        Default::default(),
+        hashed_state,
+        reth_trie::ExecutionWitnessMode::Legacy,
+    )?;
     Ok(ExecutionWitness {
         state,
         codes: codes.into_values().collect(),
@@ -239,6 +243,7 @@ where
                 DebugApiClient::<()>::debug_execution_witness(
                     healthy_node_client,
                     block_number.into(),
+                    None,
                 )
                 .await
             })?;

crates/engine/tree/src/tree/instrumented_state.rs

@@ -237,8 +237,9 @@ impl<S: StateProofProvider> StateProofProvider for InstrumentedStateProvider<S>
         &self,
         input: TrieInput,
         target: HashedPostState,
+        mode: reth_trie::ExecutionWitnessMode,
     ) -> ProviderResult<Vec<alloy_primitives::Bytes>> {
-        self.state_provider.witness(input, target)
+        self.state_provider.witness(input, target, mode)
     }
 }
 

crates/revm/src/test_utils.rs

@@ -141,7 +141,12 @@ impl StateProofProvider for StateProviderTest {
         unimplemented!("proof generation is not supported")
     }
 
-    fn witness(&self, _input: TrieInput, _target: HashedPostState) -> ProviderResult<Vec<Bytes>> {
+    fn witness(
+        &self,
+        _input: TrieInput,
+        _target: HashedPostState,
+        _mode: reth_trie::ExecutionWitnessMode,
+    ) -> ProviderResult<Vec<Bytes>> {
         unimplemented!("witness generation is not supported")
     }
 }

crates/revm/src/witness.rs

@@ -1,6 +1,6 @@
 use alloc::vec::Vec;
-use alloy_primitives::{keccak256, Bytes, B256};
-use reth_trie::{HashedPostState, HashedStorage};
+use alloy_primitives::{keccak256, map::B256Map, Bytes, B256};
+use reth_trie::{ExecutionWitnessMode, HashedPostState, HashedStorage};
 use revm::database::State;
 
 /// Tracks state changes during execution.
@@ -30,21 +30,36 @@ pub struct ExecutionWitnessRecord {
 }
 
 impl ExecutionWitnessRecord {
-    /// Records the state after execution.
-    pub fn record_executed_state<DB>(&mut self, statedb: &State<DB>) {
-        self.codes = statedb
-            .cache
-            .contracts
-            .values()
-            .map(|code| 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.values().map(|code| code.original_bytes()),
-            )
-            .collect();
+    /// Records the state after execution using the given witness generation mode.
+    pub fn record_executed_state<DB>(&mut self, statedb: &State<DB>, mode: ExecutionWitnessMode) {
+        self.codes = match mode {
+            ExecutionWitnessMode::Legacy => statedb
+                .cache
+                .contracts
+                .values()
+                .map(|code| 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.values().map(|code| code.original_bytes()),
+                )
+                .collect(),
+            ExecutionWitnessMode::Canonical => {
+                let mut accessed_codes = B256Map::default();
+                for code in statedb.cache.contracts.values() {
+                    let code = code.original_bytes();
+                    if code.is_empty() {
+                        continue;
+                    }
+                    accessed_codes.entry(keccak256(&code)).or_insert(code);
+                }
+                let mut codes: Vec<_> = accessed_codes.into_values().collect();
+                codes.sort_unstable();
+                codes
+            }
+        };
 
         for (address, account) in &statedb.cache.accounts {
             let hashed_address = keccak256(address);
@@ -75,9 +90,9 @@ impl ExecutionWitnessRecord {
     }
 
     /// Creates the record from the state after execution.
-    pub fn from_executed_state<DB>(state: &State<DB>) -> Self {
+    pub fn from_executed_state<DB>(state: &State<DB>, mode: ExecutionWitnessMode) -> Self {
         let mut record = Self::default();
-        record.record_executed_state(state);
+        record.record_executed_state(state, mode);
         record
     }
 
@@ -93,6 +108,7 @@ impl ExecutionWitnessRecord {
         state_provider: &SP,
         headers_provider: &HP,
         block_number: u64,
+        mode: ExecutionWitnessMode,
     ) -> reth_storage_errors::provider::ProviderResult<alloy_rpc_types_debug::ExecutionWitness>
     where
         SP: reth_storage_api::StateProofProvider + ?Sized,
@@ -101,7 +117,7 @@ impl ExecutionWitnessRecord {
     {
         let Self { hashed_state, codes, keys, lowest_block_number } = self;
 
-        let state = state_provider.witness(Default::default(), hashed_state)?;
+        let state = state_provider.witness(Default::default(), hashed_state, mode)?;
         let mut exec_witness =
             alloy_rpc_types_debug::ExecutionWitness { state, codes, keys, ..Default::default() };
 

crates/rpc/rpc-api/Cargo.toml

@@ -16,7 +16,7 @@ workspace = true
 reth-rpc-eth-api.workspace = true
 reth-engine-primitives.workspace = true
 reth-network-peers.workspace = true
-reth-trie-common.workspace = true
+reth-trie-common = { workspace = true, features = ["serde"] }
 reth-chain-state.workspace = true
 
 # ethereum

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

@@ -9,7 +9,7 @@ use alloy_rpc_types_trace::geth::{
     BlockTraceResult, GethDebugTracingCallOptions, GethDebugTracingOptions, GethTrace, TraceResult,
 };
 use jsonrpsee::{core::RpcResult, proc_macros::rpc};
-use reth_trie_common::{updates::TrieUpdates, HashedPostState};
+use reth_trie_common::{updates::TrieUpdates, ExecutionWitnessMode, HashedPostState};
 
 /// Debug rpc interface.
 #[cfg_attr(not(feature = "client"), rpc(server, namespace = "debug"))]
@@ -140,21 +140,27 @@ pub trait DebugApi<TxReq: RpcObject> {
     /// to their preimages that were required during the execution of the block, including during
     /// state root recomputation.
     ///
-    /// The first argument is the block number or tag.
+    /// The first argument is the block number or tag. The optional second argument selects the
+    /// witness generation mode and defaults to `legacy`.
     #[method(name = "executionWitness")]
-    async fn debug_execution_witness(&self, block: BlockNumberOrTag)
-        -> RpcResult<ExecutionWitness>;
+    async fn debug_execution_witness(
+        &self,
+        block: BlockNumberOrTag,
+        mode: Option<ExecutionWitnessMode>,
+    ) -> RpcResult<ExecutionWitness>;
 
     /// The `debug_executionWitnessByBlockHash` 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 argument is the block hash.
+    /// The first argument is the block hash. The optional second argument selects the witness
+    /// generation mode and defaults to `legacy`.
     #[method(name = "executionWitnessByBlockHash")]
     async fn debug_execution_witness_by_block_hash(
         &self,
         hash: B256,
+        mode: Option<ExecutionWitnessMode>,
     ) -> RpcResult<ExecutionWitness>;
 
     /// Re-executes a block and returns the Block Access List (BAL) as defined in EIP-7928.

crates/rpc/rpc-eth-types/src/cache/db.rs

@@ -101,8 +101,9 @@ impl reth_storage_api::StateProofProvider for StateProviderTraitObjWrapper {
         &self,
         input: reth_trie::TrieInput,
         target: reth_trie::HashedPostState,
+        mode: reth_trie::ExecutionWitnessMode,
     ) -> reth_errors::ProviderResult<Vec<alloy_primitives::Bytes>> {
-        self.0.witness(input, target)
+        self.0.witness(input, target, mode)
     }
 }
 

crates/rpc/rpc/src/debug.rs

@@ -33,11 +33,10 @@ use reth_rpc_eth_types::EthApiError;
 use reth_rpc_server_types::{result::internal_rpc_err, ToRpcResult};
 use reth_storage_api::{
     BlockIdReader, BlockReaderIdExt, HashedPostStateProvider, HeaderProvider, ProviderBlock,
-    ReceiptProviderIdExt, StateProofProvider, StateProviderFactory, StateRootProvider,
-    TransactionVariant,
+    ReceiptProviderIdExt, StateProviderFactory, StateRootProvider, TransactionVariant,
 };
 use reth_tasks::{pool::BlockingTaskGuard, Runtime};
-use reth_trie_common::{updates::TrieUpdates, HashedPostState};
+use reth_trie_common::{updates::TrieUpdates, ExecutionWitnessMode, HashedPostState};
 use revm::{database::states::bundle_state::BundleRetention, DatabaseCommit};
 use revm_inspectors::tracing::{DebugInspector, TransactionContext};
 use serde::{Deserialize, Serialize};
@@ -498,6 +497,7 @@ where
     pub async fn debug_execution_witness_by_block_hash(
         &self,
         hash: B256,
+        mode: Option<ExecutionWitnessMode>,
     ) -> Result<ExecutionWitness, Eth::Error> {
         let this = self.clone();
         let block = this
@@ -506,7 +506,7 @@ where
             .await?
             .ok_or(EthApiError::HeaderNotFound(hash.into()))?;
 
-        self.debug_execution_witness_for_block(block).await
+        self.debug_execution_witness_for_block(block, mode.unwrap_or_default()).await
     }
 
     /// The `debug_executionWitness` method allows for re-execution of a block with the purpose of
@@ -516,6 +516,7 @@ where
     pub async fn debug_execution_witness(
         &self,
         block_id: BlockNumberOrTag,
+        mode: Option<ExecutionWitnessMode>,
     ) -> Result<ExecutionWitness, Eth::Error> {
         let this = self.clone();
         let block = this
@@ -524,67 +525,33 @@ where
             .await?
             .ok_or(EthApiError::HeaderNotFound(block_id.into()))?;
 
-        self.debug_execution_witness_for_block(block).await
+        self.debug_execution_witness_for_block(block, mode.unwrap_or_default()).await
     }
 
     /// Generates an execution witness, using the given recovered block.
     pub async fn debug_execution_witness_for_block(
         &self,
         block: Arc<RecoveredBlock<ProviderBlock<Eth::Provider>>>,
+        mode: ExecutionWitnessMode,
     ) -> Result<ExecutionWitness, Eth::Error> {
         let block_number = block.header().number();
-
-        let (mut exec_witness, lowest_block_number) = self
-            .eth_api()
+        self.eth_api()
             .spawn_with_state_at_block(block.parent_hash(), move |eth_api, mut db| {
                 let block_executor = eth_api.evm_config().executor(&mut db);
 
                 let mut witness_record = ExecutionWitnessRecord::default();
 
                 let _ = block_executor
                     .execute_with_state_closure(&block, |statedb: &State<_>| {
-                        witness_record.record_executed_state(statedb);
+                        witness_record.record_executed_state(statedb, mode);
                     })
                     .map_err(|err| EthApiError::Internal(err.into()))?;
 
-                let ExecutionWitnessRecord { hashed_state, codes, keys, lowest_block_number } =
-                    witness_record;
-
-                let state = db
-                    .database
-                    .0
-                    .witness(Default::default(), hashed_state)
-                    .map_err(EthApiError::from)?;
-                Ok((
-                    ExecutionWitness { state, codes, keys, ..Default::default() },
-                    lowest_block_number,
-                ))
+                Ok(witness_record
+                    .into_execution_witness(&db.database.0, eth_api.provider(), block_number, mode)
+                    .map_err(EthApiError::from)?)
             })
-            .await?;
-
-        let smallest = match lowest_block_number {
-            Some(smallest) => smallest,
-            None => {
-                // Return only the parent header, if there were no calls to the
-                // BLOCKHASH opcode.
-                block_number.saturating_sub(1)
-            }
-        };
-
-        let range = smallest..block_number;
-        exec_witness.headers = self
-            .provider()
-            .headers_range(range)
-            .map_err(EthApiError::from)?
-            .into_iter()
-            .map(|header| {
-                let mut serialized_header = Vec::new();
-                header.encode(&mut serialized_header);
-                serialized_header.into()
-            })
-            .collect();
-
-        Ok(exec_witness)
+            .await
     }
 
     /// Returns the code associated with a given hash at the specified block ID. If no code is
@@ -846,18 +813,20 @@ where
     async fn debug_execution_witness(
         &self,
         block: BlockNumberOrTag,
+        mode: Option<ExecutionWitnessMode>,
     ) -> RpcResult<ExecutionWitness> {
         let _permit = self.acquire_trace_permit().await;
-        Self::debug_execution_witness(self, block).await.map_err(Into::into)
+        Self::debug_execution_witness(self, block, mode).await.map_err(Into::into)
     }
 
     /// Handler for `debug_executionWitnessByBlockHash`
     async fn debug_execution_witness_by_block_hash(
         &self,
         hash: B256,
+        mode: Option<ExecutionWitnessMode>,
     ) -> RpcResult<ExecutionWitness> {
         let _permit = self.acquire_trace_permit().await;
-        Self::debug_execution_witness_by_block_hash(self, hash).await.map_err(Into::into)
+        Self::debug_execution_witness_by_block_hash(self, hash, mode).await.map_err(Into::into)
     }
 
     async fn debug_get_block_access_list(&self, _block_id: BlockId) -> RpcResult<BlockAccessList> {