feat(tree): implement trie traits for in-memory overlay provider (#9522)

crates/engine/tree/src/database.rs

@@ -85,7 +85,8 @@ impl<DB: Database> DatabaseService<DB> {
             {
                 let trie_updates = block.trie_updates().clone();
                 let hashed_state = block.hashed_state();
-                HashedStateChanges(hashed_state).write_to_db(&provider_rw)?;
+                HashedStateChanges(&hashed_state.clone().into_sorted())
+                    .write_to_db(&provider_rw)?;
                 trie_updates.write_to_database(provider_rw.tx_ref())?;
             }
 

crates/engine/tree/src/test_utils.rs

@@ -11,7 +11,7 @@ use reth_prune_types::PruneModes;
 use reth_stages::{test_utils::TestStages, ExecOutput, StageError};
 use reth_stages_api::Pipeline;
 use reth_static_file::StaticFileProducer;
-use reth_trie::{updates::TrieUpdates, HashedPostStateSorted};
+use reth_trie::{updates::TrieUpdates, HashedPostState};
 use revm::db::BundleState;
 use std::{collections::VecDeque, ops::Range, sync::Arc};
 use tokio::sync::watch;
@@ -103,7 +103,7 @@ pub(crate) fn get_executed_block_with_number(block_number: BlockNumber) -> Execu
             block_number,
             vec![Requests::default()],
         )),
-        Arc::new(HashedPostStateSorted::default()),
+        Arc::new(HashedPostState::default()),
         Arc::new(TrieUpdates::default()),
     )
 }

crates/engine/tree/src/tree/memory_overlay.rs

@@ -12,14 +12,20 @@ use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
 pub struct MemoryOverlayStateProvider<H> {
     /// The collection of executed parent blocks.
     in_memory: Vec<ExecutedBlock>,
+    /// The collection of hashed state from in-memory blocks.
+    hashed_post_state: HashedPostState,
     /// Historical state provider for state lookups that are not found in in-memory blocks.
     historical: H,
 }
 
 impl<H> MemoryOverlayStateProvider<H> {
     /// Create new memory overlay state provider.
-    pub const fn new(in_memory: Vec<ExecutedBlock>, historical: H) -> Self {
-        Self { in_memory, historical }
+    pub fn new(in_memory: Vec<ExecutedBlock>, historical: H) -> Self {
+        let mut hashed_post_state = HashedPostState::default();
+        for block in &in_memory {
+            hashed_post_state.extend(block.hashed_state.as_ref().clone());
+        }
+        Self { in_memory, hashed_post_state, historical }
     }
 }
 
@@ -78,15 +84,21 @@ impl<H> StateRootProvider for MemoryOverlayStateProvider<H>
 where
     H: StateRootProvider + Send,
 {
+    // TODO: Currently this does not reuse available in-memory trie nodes.
     fn hashed_state_root(&self, hashed_state: &HashedPostState) -> ProviderResult<B256> {
-        todo!()
+        let mut state = self.hashed_post_state.clone();
+        state.extend(hashed_state.clone());
+        self.historical.hashed_state_root(&state)
     }
 
+    // TODO: Currently this does not reuse available in-memory trie nodes.
     fn hashed_state_root_with_updates(
         &self,
         hashed_state: &HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        todo!()
+        let mut state = self.hashed_post_state.clone();
+        state.extend(hashed_state.clone());
+        self.historical.hashed_state_root_with_updates(&state)
     }
 }
 
@@ -94,13 +106,16 @@ impl<H> StateProofProvider for MemoryOverlayStateProvider<H>
 where
     H: StateProofProvider + Send,
 {
+    // TODO: Currently this does not reuse available in-memory trie nodes.
     fn hashed_proof(
         &self,
         hashed_state: &HashedPostState,
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
-        todo!()
+        let mut state = self.hashed_post_state.clone();
+        state.extend(hashed_state.clone());
+        self.historical.hashed_proof(&state, address, slots)
     }
 }
 

crates/engine/tree/src/tree/mod.rs

@@ -32,7 +32,7 @@ use reth_rpc_types::{
     },
     ExecutionPayload,
 };
-use reth_trie::{updates::TrieUpdates, HashedPostState, HashedPostStateSorted};
+use reth_trie::{updates::TrieUpdates, HashedPostState};
 use std::{
     collections::{BTreeMap, HashMap},
     marker::PhantomData,
@@ -53,7 +53,7 @@ pub struct ExecutedBlock {
     block: Arc<SealedBlock>,
     senders: Arc<Vec<Address>>,
     execution_output: Arc<ExecutionOutcome>,
-    hashed_state: Arc<HashedPostStateSorted>,
+    hashed_state: Arc<HashedPostState>,
     trie: Arc<TrieUpdates>,
 }
 
@@ -62,7 +62,7 @@ impl ExecutedBlock {
         block: Arc<SealedBlock>,
         senders: Arc<Vec<Address>>,
         execution_output: Arc<ExecutionOutcome>,
-        hashed_state: Arc<HashedPostStateSorted>,
+        hashed_state: Arc<HashedPostState>,
         trie: Arc<TrieUpdates>,
     ) -> Self {
         Self { block, senders, execution_output, hashed_state, trie }
@@ -84,7 +84,7 @@ impl ExecutedBlock {
     }
 
     /// Returns a reference to the hashed state result of the execution outcome
-    pub(crate) fn hashed_state(&self) -> &HashedPostStateSorted {
+    pub(crate) fn hashed_state(&self) -> &HashedPostState {
         &self.hashed_state
     }
 
@@ -666,7 +666,7 @@ where
                 block_number,
                 vec![Requests::from(output.requests)],
             )),
-            hashed_state: Arc::new(hashed_state.into_sorted()),
+            hashed_state: Arc::new(hashed_state),
             trie: Arc::new(trie_output),
         };
         self.state.tree_state.insert_executed(executed);