reth_chain_state/
memory_overlay.rs

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

reth_chain_state/
memory_overlay.rs
