fix(trie): use block hash in OverlayStateProviderFactory

crates/engine/tree/src/tree/payload_validator.rs

@@ -15,7 +15,7 @@ use crate::tree::{
 use alloy_consensus::transaction::Either;
 use alloy_eips::{eip1898::BlockWithParent, NumHash};
 use alloy_evm::Evm;
-use alloy_primitives::{BlockNumber, B256};
+use alloy_primitives::B256;
 use reth_chain_state::{CanonicalInMemoryState, ExecutedBlock};
 use reth_consensus::{ConsensusError, FullConsensus};
 use reth_engine_primitives::{
@@ -33,8 +33,8 @@ use reth_primitives_traits::{
     AlloyBlockHeader, BlockTy, GotExpected, NodePrimitives, RecoveredBlock, SealedHeader,
 };
 use reth_provider::{
-    providers::OverlayStateProviderFactory, BlockExecutionOutput, BlockNumReader, BlockReader,
-    DBProvider, DatabaseProviderFactory, ExecutionOutcome, HashedPostStateProvider, ProviderError,
+    providers::OverlayStateProviderFactory, BlockExecutionOutput, BlockReader,
+    DatabaseProviderFactory, ExecutionOutcome, HashedPostStateProvider, ProviderError,
     PruneCheckpointReader, StageCheckpointReader, StateProvider, StateProviderFactory, StateReader,
     StateRootProvider, TrieReader,
 };
@@ -680,10 +680,7 @@ where
         hashed_state: &HashedPostState,
         state: &EngineApiTreeState<N>,
     ) -> Result<(B256, TrieUpdates), ParallelStateRootError> {
-        let provider = self.provider.database_provider_ro()?;
-
-        let (mut input, block_number) =
-            self.compute_trie_input(provider, parent_hash, state, None)?;
+        let (mut input, block_hash) = self.compute_trie_input(parent_hash, state, None)?;
 
         // Extend with block we are validating root for.
         input.append_ref(hashed_state);
@@ -693,7 +690,7 @@ where
         let (_, multiproof_config) = MultiProofConfig::from_input(input);
 
         let factory = OverlayStateProviderFactory::new(self.provider.clone())
-            .with_block_number(Some(block_number))
+            .with_block_hash(Some(block_hash))
             .with_trie_overlay(Some(multiproof_config.nodes_sorted))
             .with_hashed_state_overlay(Some(multiproof_config.state_sorted));
 
@@ -806,12 +803,8 @@ where
 
                 // Compute trie input
                 let trie_input_start = Instant::now();
-                let (trie_input, block_number) = self.compute_trie_input(
-                    self.provider.database_provider_ro()?,
-                    parent_hash,
-                    state,
-                    allocated_trie_input,
-                )?;
+                let (trie_input, block_hash) =
+                    self.compute_trie_input(parent_hash, state, allocated_trie_input)?;
 
                 self.metrics
                     .block_validation
@@ -827,7 +820,7 @@ where
                 // multiproofs.
                 let multiproof_provider_factory =
                     OverlayStateProviderFactory::new(self.provider.clone())
-                        .with_block_number(Some(block_number))
+                        .with_block_hash(Some(block_hash))
                         .with_trie_overlay(Some(multiproof_config.nodes_sorted))
                         .with_hashed_state_overlay(Some(multiproof_config.state_sorted));
 
@@ -976,38 +969,30 @@ where
         skip_all,
         fields(parent_hash)
     )]
-    fn compute_trie_input<TP: DBProvider + BlockNumReader + TrieReader>(
+    fn compute_trie_input(
         &self,
-        provider: TP,
         parent_hash: B256,
         state: &EngineApiTreeState<N>,
         allocated_trie_input: Option<TrieInput>,
-    ) -> ProviderResult<(TrieInput, BlockNumber)> {
+    ) -> ProviderResult<(TrieInput, B256)> {
         // get allocated trie input or use a default trie input
         let mut input = allocated_trie_input.unwrap_or_default();
 
-        let (historical, blocks) = state
-            .tree_state
-            .blocks_by_hash(parent_hash)
-            .map_or_else(|| (parent_hash.into(), vec![]), |(hash, blocks)| (hash.into(), blocks));
+        let (block_hash, blocks) =
+            state.tree_state.blocks_by_hash(parent_hash).unwrap_or_else(|| (parent_hash, vec![]));
 
         if blocks.is_empty() {
             debug!(target: "engine::tree::payload_validator", "Parent found on disk");
         } else {
-            debug!(target: "engine::tree::payload_validator", %historical, blocks = blocks.len(), "Parent found in memory");
+            debug!(target: "engine::tree::payload_validator", historical = ?block_hash, blocks = blocks.len(), "Parent found in memory");
         }
 
-        // Convert the historical block to the block number
-        let block_number = provider
-            .convert_hash_or_number(historical)?
-            .ok_or_else(|| ProviderError::BlockHashNotFound(historical.as_hash().unwrap()))?;
-
         // Extend with contents of parent in-memory blocks.
         input.extend_with_blocks(
             blocks.iter().rev().map(|block| (block.hashed_state(), block.trie_updates())),
         );
 
-        Ok((input, block_number))
+        Ok((input, block_hash))
     }
 }
 

crates/storage/provider/src/providers/state/overlay.rs

@@ -4,8 +4,8 @@ use reth_errors::ProviderError;
 use reth_prune_types::PruneSegment;
 use reth_stages_types::StageId;
 use reth_storage_api::{
-    DBProvider, DatabaseProviderFactory, DatabaseProviderROFactory, PruneCheckpointReader,
-    StageCheckpointReader, TrieReader,
+    BlockNumReader, DBProvider, DatabaseProviderFactory, DatabaseProviderROFactory,
+    PruneCheckpointReader, StageCheckpointReader, TrieReader,
 };
 use reth_trie::{
     hashed_cursor::{HashedCursorFactory, HashedPostStateCursorFactory},
@@ -27,8 +27,8 @@ use tracing::debug;
 pub struct OverlayStateProviderFactory<F> {
     /// The underlying database provider factory
     factory: F,
-    /// Optional block number for collecting reverts
-    block_number: Option<BlockNumber>,
+    /// Optional block hash for collecting reverts
+    block_hash: Option<B256>,
     /// Optional trie overlay
     trie_overlay: Option<Arc<TrieUpdatesSorted>>,
     /// Optional hashed state overlay
@@ -38,27 +38,27 @@ pub struct OverlayStateProviderFactory<F> {
 impl<F> OverlayStateProviderFactory<F> {
     /// Create a new overlay state provider factory
     pub const fn new(factory: F) -> Self {
-        Self { factory, block_number: None, trie_overlay: None, hashed_state_overlay: None }
+        Self { factory, block_hash: None, trie_overlay: None, hashed_state_overlay: None }
     }
 
-    /// Set the block number for collecting reverts. All state will be reverted to the point
+    /// Set the block hash for collecting reverts. All state will be reverted to the point
     /// _after_ this block has been processed.
-    pub const fn with_block_number(mut self, block_number: Option<BlockNumber>) -> Self {
-        self.block_number = block_number;
+    pub const fn with_block_hash(mut self, block_hash: Option<B256>) -> Self {
+        self.block_hash = block_hash;
         self
     }
 
     /// Set the trie overlay.
     ///
-    /// This overlay will be applied on top of any reverts applied via `with_block_number`.
+    /// This overlay will be applied on top of any reverts applied via `with_block_hash`.
     pub fn with_trie_overlay(mut self, trie_overlay: Option<Arc<TrieUpdatesSorted>>) -> Self {
         self.trie_overlay = trie_overlay;
         self
     }
 
     /// Set the hashed state overlay
     ///
-    /// This overlay will be applied on top of any reverts applied via `with_block_number`.
+    /// This overlay will be applied on top of any reverts applied via `with_block_hash`.
     pub fn with_hashed_state_overlay(
         mut self,
         hashed_state_overlay: Option<Arc<HashedPostStateSorted>>,
@@ -127,7 +127,7 @@ where
 impl<F> DatabaseProviderROFactory for OverlayStateProviderFactory<F>
 where
     F: DatabaseProviderFactory,
-    F::Provider: TrieReader + StageCheckpointReader + PruneCheckpointReader,
+    F::Provider: TrieReader + StageCheckpointReader + PruneCheckpointReader + BlockNumReader,
 {
     type Provider = OverlayStateProvider<F::Provider>;
 
@@ -136,8 +136,13 @@ where
         // Get a read-only provider
         let provider = self.factory.database_provider_ro()?;
 
-        // If block_number is provided, collect reverts
-        let (trie_updates, hashed_state) = if let Some(from_block) = self.block_number {
+        // If block_hash is provided, collect reverts
+        let (trie_updates, hashed_state) = if let Some(block_hash) = self.block_hash {
+            // Convert block hash to block number
+            let from_block = provider
+                .convert_hash_or_number(block_hash.into())?
+                .ok_or_else(|| ProviderError::BlockHashNotFound(block_hash))?;
+
             // Validate that we have sufficient changesets for the requested block
             self.validate_changesets_availability(&provider, from_block)?;
 
@@ -162,6 +167,7 @@ where
 
             debug!(
                 target: "providers::state::overlay",
+                ?block_hash,
                 ?from_block,
                 num_trie_updates = ?trie_updates_mut.total_len(),
                 num_state_updates = ?hashed_state_mut.total_len(),
@@ -170,7 +176,7 @@ where
 
             (Arc::new(trie_updates_mut), Arc::new(hashed_state_mut))
         } else {
-            // If no block_number, use overlays directly or defaults
+            // If no block_hash, use overlays directly or defaults
             let trie_updates =
                 self.trie_overlay.clone().unwrap_or_else(|| Arc::new(TrieUpdatesSorted::default()));
             let hashed_state = self