reth_chain_state/
memory_overlay.rs

1use super::ExecutedBlock;
2use alloy_consensus::BlockHeader;
3use alloy_primitives::{keccak256, Address, BlockNumber, Bytes, StorageKey, StorageValue, B256};
4use reth_errors::ProviderResult;
5use reth_primitives_traits::{Account, Bytecode, NodePrimitives};
6use reth_storage_api::{
7    AccountReader, BlockHashReader, BytecodeReader, HashedPostStateProvider, StateProofProvider,
8    StateProvider, StateRootProvider, StorageRootProvider,
9};
10use reth_trie::{
11    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
12    MultiProofTargets, StorageMultiProof, TrieInput,
13};
14use revm_database::BundleState;
15use std::sync::OnceLock;
16
17/// A state provider that stores references to in-memory blocks along with their state as well as a
18/// reference of the historical state provider for fallback lookups.
19#[expect(missing_debug_implementations)]
20pub struct MemoryOverlayStateProviderRef<
21    'a,
22    N: NodePrimitives = reth_ethereum_primitives::EthPrimitives,
23> {
24    /// Historical state provider for state lookups that are not found in memory blocks.
25    pub(crate) historical: Box<dyn StateProvider + 'a>,
26    /// The collection of executed parent blocks. Expected order is newest to oldest.
27    pub(crate) in_memory: Vec<ExecutedBlock<N>>,
28    /// Lazy-loaded in-memory trie data.
29    pub(crate) trie_input: OnceLock<TrieInput>,
30}
31
32/// A state provider that stores references to in-memory blocks along with their state as well as
33/// the historical state provider for fallback lookups.
34pub type MemoryOverlayStateProvider<N> = MemoryOverlayStateProviderRef<'static, N>;
35
36impl<'a, N: NodePrimitives> MemoryOverlayStateProviderRef<'a, N> {
37    /// Create new memory overlay state provider.
38    ///
39    /// ## Arguments
40    ///
41    /// - `in_memory` - the collection of executed ancestor blocks in reverse.
42    /// - `historical` - a historical state provider for the latest ancestor block stored in the
43    ///   database.
44    pub fn new(historical: Box<dyn StateProvider + 'a>, in_memory: Vec<ExecutedBlock<N>>) -> Self {
45        Self { historical, in_memory, trie_input: OnceLock::new() }
46    }
47
48    /// Turn this state provider into a state provider
49    pub fn boxed(self) -> Box<dyn StateProvider + 'a> {
50        Box::new(self)
51    }
52
53    /// Return lazy-loaded trie state aggregated from in-memory blocks.
54    fn trie_input(&self) -> &TrieInput {
55        self.trie_input.get_or_init(|| {
56            let bundles: Vec<_> =
57                self.in_memory.iter().rev().map(|block| block.trie_data()).collect();
58            TrieInput::from_blocks_sorted(
59                bundles.iter().map(|data| (data.hashed_state.as_ref(), data.trie_updates.as_ref())),
60            )
61        })
62    }
63
64    fn merged_hashed_storage(&self, address: Address, storage: HashedStorage) -> HashedStorage {
65        let state = &self.trie_input().state;
66        let mut hashed = state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
67        hashed.extend(&storage);
68        hashed
69    }
70}
71
72impl<N: NodePrimitives> BlockHashReader for MemoryOverlayStateProviderRef<'_, N> {
73    fn block_hash(&self, number: BlockNumber) -> ProviderResult<Option<B256>> {
74        for block in &self.in_memory {
75            if block.recovered_block().number() == number {
76                return Ok(Some(block.recovered_block().hash()));
77            }
78        }
79
80        self.historical.block_hash(number)
81    }
82
83    fn canonical_hashes_range(
84        &self,
85        start: BlockNumber,
86        end: BlockNumber,
87    ) -> ProviderResult<Vec<B256>> {
88        let range = start..end;
89        let mut earliest_block_number = None;
90        let mut in_memory_hashes = Vec::with_capacity(range.size_hint().0);
91
92        // iterate in ascending order (oldest to newest = low to high)
93        for block in &self.in_memory {
94            let block_num = block.recovered_block().number();
95            if range.contains(&block_num) {
96                in_memory_hashes.push(block.recovered_block().hash());
97                earliest_block_number = Some(block_num);
98            }
99        }
100
101        // `self.in_memory` stores executed blocks in ascending order (oldest to newest).
102        // However, `in_memory_hashes` should be constructed in descending order (newest to oldest),
103        // so we reverse the vector after collecting the hashes.
104        in_memory_hashes.reverse();
105
106        let mut hashes =
107            self.historical.canonical_hashes_range(start, earliest_block_number.unwrap_or(end))?;
108        hashes.append(&mut in_memory_hashes);
109        Ok(hashes)
110    }
111}
112
113impl<N: NodePrimitives> AccountReader for MemoryOverlayStateProviderRef<'_, N> {
114    fn basic_account(&self, address: &Address) -> ProviderResult<Option<Account>> {
115        for block in &self.in_memory {
116            if let Some(account) = block.execution_output.account(address) {
117                return Ok(account);
118            }
119        }
120
121        self.historical.basic_account(address)
122    }
123}
124
125impl<N: NodePrimitives> StateRootProvider for MemoryOverlayStateProviderRef<'_, N> {
126    fn state_root(&self, state: HashedPostState) -> ProviderResult<B256> {
127        self.state_root_from_nodes(TrieInput::from_state(state))
128    }
129
130    fn state_root_from_nodes(&self, mut input: TrieInput) -> ProviderResult<B256> {
131        input.prepend_self(self.trie_input().clone());
132        self.historical.state_root_from_nodes(input)
133    }
134
135    fn state_root_with_updates(
136        &self,
137        state: HashedPostState,
138    ) -> ProviderResult<(B256, TrieUpdates)> {
139        self.state_root_from_nodes_with_updates(TrieInput::from_state(state))
140    }
141
142    fn state_root_from_nodes_with_updates(
143        &self,
144        mut input: TrieInput,
145    ) -> ProviderResult<(B256, TrieUpdates)> {
146        input.prepend_self(self.trie_input().clone());
147        self.historical.state_root_from_nodes_with_updates(input)
148    }
149}
150
151impl<N: NodePrimitives> StorageRootProvider for MemoryOverlayStateProviderRef<'_, N> {
152    // TODO: Currently this does not reuse available in-memory trie nodes.
153    fn storage_root(&self, address: Address, storage: HashedStorage) -> ProviderResult<B256> {
154        let merged = self.merged_hashed_storage(address, storage);
155        self.historical.storage_root(address, merged)
156    }
157
158    // TODO: Currently this does not reuse available in-memory trie nodes.
159    fn storage_proof(
160        &self,
161        address: Address,
162        slot: B256,
163        storage: HashedStorage,
164    ) -> ProviderResult<reth_trie::StorageProof> {
165        let merged = self.merged_hashed_storage(address, storage);
166        self.historical.storage_proof(address, slot, merged)
167    }
168
169    // TODO: Currently this does not reuse available in-memory trie nodes.
170    fn storage_multiproof(
171        &self,
172        address: Address,
173        slots: &[B256],
174        storage: HashedStorage,
175    ) -> ProviderResult<StorageMultiProof> {
176        let merged = self.merged_hashed_storage(address, storage);
177        self.historical.storage_multiproof(address, slots, merged)
178    }
179}
180
181impl<N: NodePrimitives> StateProofProvider for MemoryOverlayStateProviderRef<'_, N> {
182    fn proof(
183        &self,
184        mut input: TrieInput,
185        address: Address,
186        slots: &[B256],
187    ) -> ProviderResult<AccountProof> {
188        input.prepend_self(self.trie_input().clone());
189        self.historical.proof(input, address, slots)
190    }
191
192    fn multiproof(
193        &self,
194        mut input: TrieInput,
195        targets: MultiProofTargets,
196    ) -> ProviderResult<MultiProof> {
197        input.prepend_self(self.trie_input().clone());
198        self.historical.multiproof(input, targets)
199    }
200
201    fn witness(&self, mut input: TrieInput, target: HashedPostState) -> ProviderResult<Vec<Bytes>> {
202        input.prepend_self(self.trie_input().clone());
203        self.historical.witness(input, target)
204    }
205}
206
207impl<N: NodePrimitives> HashedPostStateProvider for MemoryOverlayStateProviderRef<'_, N> {
208    fn hashed_post_state(&self, bundle_state: &BundleState) -> HashedPostState {
209        self.historical.hashed_post_state(bundle_state)
210    }
211}
212
213impl<N: NodePrimitives> StateProvider for MemoryOverlayStateProviderRef<'_, N> {
214    fn storage(
215        &self,
216        address: Address,
217        storage_key: StorageKey,
218    ) -> ProviderResult<Option<StorageValue>> {
219        for block in &self.in_memory {
220            if let Some(value) = block.execution_output.storage(&address, storage_key.into()) {
221                return Ok(Some(value));
222            }
223        }
224
225        self.historical.storage(address, storage_key)
226    }
227}
228
229impl<N: NodePrimitives> BytecodeReader for MemoryOverlayStateProviderRef<'_, N> {
230    fn bytecode_by_hash(&self, code_hash: &B256) -> ProviderResult<Option<Bytecode>> {
231        for block in &self.in_memory {
232            if let Some(contract) = block.execution_output.bytecode(code_hash) {
233                return Ok(Some(contract));
234            }
235        }
236
237        self.historical.bytecode_by_hash(code_hash)
238    }
239}
reth_chain_state/memory_overlay.rs

reth_chain_state/
memory_overlay.rs
