diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs index a2dd656af..28f20607a 100644 --- a/crates/chain-state/src/in_memory.rs +++ b/crates/chain-state/src/in_memory.rs @@ -493,7 +493,7 @@ impl CanonicalInMemoryState { Vec::new() }; - MemoryOverlayStateProvider::new(in_memory, historical) + MemoryOverlayStateProvider::new(historical, in_memory) } /// Returns an iterator over all canonical blocks in the in-memory state, from newest to oldest. diff --git a/crates/chain-state/src/memory_overlay.rs b/crates/chain-state/src/memory_overlay.rs index c4a0f2888..4ba10d34a 100644 --- a/crates/chain-state/src/memory_overlay.rs +++ b/crates/chain-state/src/memory_overlay.rs @@ -11,20 +11,18 @@ use reth_trie::{ prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, }; -use std::collections::HashMap; +use std::{collections::HashMap, sync::OnceLock}; /// A state provider that stores references to in-memory blocks along with their state as well as /// the historical state provider for fallback lookups. #[allow(missing_debug_implementations)] pub struct MemoryOverlayStateProvider { - /// The collection of executed parent blocks. Expected order is newest to oldest. - pub(crate) in_memory: Vec, - /// The collection of hashed state from in-memory blocks. - pub(crate) hashed_post_state: HashedPostState, - /// The collection of aggregated in-memory trie updates. - pub(crate) trie_updates: TrieUpdates, /// Historical state provider for state lookups that are not found in in-memory blocks. pub(crate) historical: Box, + /// The collection of executed parent blocks. Expected order is newest to oldest. + pub(crate) in_memory: Vec, + /// Lazy-loaded in-memory trie data. + pub(crate) trie_state: OnceLock, } impl MemoryOverlayStateProvider { @@ -35,20 +33,29 @@ impl MemoryOverlayStateProvider { /// - `in_memory` - the collection of executed ancestor blocks in reverse. /// - `historical` - a historical state provider for the latest ancestor block stored in the /// database. - pub fn new(in_memory: Vec, historical: Box) -> Self { - let mut hashed_post_state = HashedPostState::default(); - let mut trie_updates = TrieUpdates::default(); - for block in in_memory.iter().rev() { - hashed_post_state.extend_ref(block.hashed_state.as_ref()); - trie_updates.extend_ref(block.trie.as_ref()); - } - Self { in_memory, hashed_post_state, trie_updates, historical } + pub fn new(historical: Box, in_memory: Vec) -> Self { + Self { historical, in_memory, trie_state: OnceLock::new() } } /// Turn this state provider into a [`StateProviderBox`] pub fn boxed(self) -> StateProviderBox { Box::new(self) } + + /// Return lazy-loaded trie state aggregated from in-memory blocks. + fn trie_state(&self) -> MemoryOverlayTrieState { + self.trie_state + .get_or_init(|| { + let mut hashed_state = HashedPostState::default(); + let mut trie_nodes = TrieUpdates::default(); + for block in self.in_memory.iter().rev() { + hashed_state.extend_ref(block.hashed_state.as_ref()); + trie_nodes.extend_ref(block.trie.as_ref()); + } + MemoryOverlayTrieState { trie_nodes, hashed_state } + }) + .clone() + } } impl BlockHashReader for MemoryOverlayStateProvider { @@ -105,14 +112,13 @@ impl StateRootProvider for MemoryOverlayStateProvider { fn hashed_state_root_from_nodes( &self, nodes: TrieUpdates, - hashed_state: HashedPostState, + state: HashedPostState, prefix_sets: TriePrefixSetsMut, ) -> ProviderResult { - let mut trie_nodes = self.trie_updates.clone(); + let MemoryOverlayTrieState { mut trie_nodes, mut hashed_state } = self.trie_state(); trie_nodes.extend(nodes); - let mut state = self.hashed_post_state.clone(); - state.extend(hashed_state); - self.historical.hashed_state_root_from_nodes(trie_nodes, state, prefix_sets) + hashed_state.extend(state); + self.historical.hashed_state_root_from_nodes(trie_nodes, hashed_state, prefix_sets) } fn hashed_state_root_with_updates( @@ -130,14 +136,17 @@ impl StateRootProvider for MemoryOverlayStateProvider { fn hashed_state_root_from_nodes_with_updates( &self, nodes: TrieUpdates, - hashed_state: HashedPostState, + state: HashedPostState, prefix_sets: TriePrefixSetsMut, ) -> ProviderResult<(B256, TrieUpdates)> { - let mut trie_nodes = self.trie_updates.clone(); + let MemoryOverlayTrieState { mut trie_nodes, mut hashed_state } = self.trie_state(); trie_nodes.extend(nodes); - let mut state = self.hashed_post_state.clone(); - state.extend(hashed_state); - self.historical.hashed_state_root_from_nodes_with_updates(trie_nodes, state, prefix_sets) + hashed_state.extend(state); + self.historical.hashed_state_root_from_nodes_with_updates( + trie_nodes, + hashed_state, + prefix_sets, + ) } // TODO: Currently this does not reuse available in-memory trie nodes. @@ -153,13 +162,13 @@ impl StateRootProvider for MemoryOverlayStateProvider { impl StateProofProvider for MemoryOverlayStateProvider { fn hashed_proof( &self, - hashed_state: HashedPostState, + state: HashedPostState, address: Address, slots: &[B256], ) -> ProviderResult { - let mut state = self.hashed_post_state.clone(); - state.extend(hashed_state); - self.historical.hashed_proof(state, address, slots) + let MemoryOverlayTrieState { mut hashed_state, .. } = self.trie_state(); + hashed_state.extend(state); + self.historical.hashed_proof(hashed_state, address, slots) } // TODO: Currently this does not reuse available in-memory trie nodes. @@ -168,9 +177,9 @@ impl StateProofProvider for MemoryOverlayStateProvider { overlay: HashedPostState, target: HashedPostState, ) -> ProviderResult> { - let mut state = self.hashed_post_state.clone(); - state.extend(overlay); - self.historical.witness(state, target) + let MemoryOverlayTrieState { mut hashed_state, .. } = self.trie_state(); + hashed_state.extend(overlay); + self.historical.witness(hashed_state, target) } } @@ -199,3 +208,12 @@ impl StateProvider for MemoryOverlayStateProvider { self.historical.bytecode_by_hash(code_hash) } } + +/// The collection of data necessary for trie-related operations for [`MemoryOverlayStateProvider`]. +#[derive(Clone, Debug)] +pub(crate) struct MemoryOverlayTrieState { + /// The collection of aggregated in-memory trie updates. + pub(crate) trie_nodes: TrieUpdates, + /// The collection of hashed state from in-memory blocks. + pub(crate) hashed_state: HashedPostState, +} diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs index 71d93f3a6..4eeca2a35 100644 --- a/crates/engine/tree/src/tree/mod.rs +++ b/crates/engine/tree/src/tree/mod.rs @@ -1232,7 +1232,7 @@ where trace!(target: "engine", %hash, "found canonical state for block in memory"); // the block leads back to the canonical chain let historical = self.provider.state_by_block_hash(historical)?; - return Ok(Some(Box::new(MemoryOverlayStateProvider::new(blocks, historical)))) + return Ok(Some(Box::new(MemoryOverlayStateProvider::new(historical, blocks)))) } // the hash could belong to an unknown block or a persisted block