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, StateProviderBox, StateRootProvider, StorageRootProvider,
9};
10use reth_trie::{
11 updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
12 MultiProofTargets, StorageMultiProof, TrieInput,
13};
14use revm_database::BundleState;
15use std::{borrow::Cow, sync::OnceLock};
16
17#[expect(missing_debug_implementations)]
20pub struct MemoryOverlayStateProviderRef<
21 'a,
22 N: NodePrimitives = reth_ethereum_primitives::EthPrimitives,
23> {
24 pub(crate) historical: Box<dyn StateProvider + 'a>,
26 pub(crate) in_memory: Cow<'a, [ExecutedBlock<N>]>,
28 pub(crate) trie_input: OnceLock<TrieInput>,
30}
31
32impl<'a, N: NodePrimitives> MemoryOverlayStateProviderRef<'a, N> {
33 pub fn new(historical: Box<dyn StateProvider + 'a>, in_memory: Vec<ExecutedBlock<N>>) -> Self {
41 Self { historical, in_memory: Cow::Owned(in_memory), trie_input: OnceLock::new() }
42 }
43
44 pub fn boxed(self) -> Box<dyn StateProvider + 'a> {
46 Box::new(self)
47 }
48
49 fn trie_input(&self) -> &TrieInput {
51 self.trie_input.get_or_init(|| {
52 let bundles: Vec<_> =
53 self.in_memory.iter().rev().map(|block| block.trie_data()).collect();
54 TrieInput::from_blocks_sorted(
55 bundles.iter().map(|data| (data.hashed_state.as_ref(), data.trie_updates.as_ref())),
56 )
57 })
58 }
59
60 fn merged_hashed_storage(&self, address: Address, storage: HashedStorage) -> HashedStorage {
61 let state = &self.trie_input().state;
62 let mut hashed = state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
63 hashed.extend(&storage);
64 hashed
65 }
66}
67
68impl<N: NodePrimitives> BlockHashReader for MemoryOverlayStateProviderRef<'_, N> {
69 fn block_hash(&self, number: BlockNumber) -> ProviderResult<Option<B256>> {
70 for block in self.in_memory.iter() {
71 if block.recovered_block().number() == number {
72 return Ok(Some(block.recovered_block().hash()));
73 }
74 }
75
76 self.historical.block_hash(number)
77 }
78
79 fn canonical_hashes_range(
80 &self,
81 start: BlockNumber,
82 end: BlockNumber,
83 ) -> ProviderResult<Vec<B256>> {
84 let range = start..end;
85 let mut earliest_block_number = None;
86 let mut in_memory_hashes = Vec::with_capacity(range.size_hint().0);
87
88 for block in self.in_memory.iter() {
90 let block_num = block.recovered_block().number();
91 if range.contains(&block_num) {
92 in_memory_hashes.push(block.recovered_block().hash());
93 earliest_block_number = Some(block_num);
94 }
95 }
96
97 in_memory_hashes.reverse();
101
102 let mut hashes =
103 self.historical.canonical_hashes_range(start, earliest_block_number.unwrap_or(end))?;
104 hashes.append(&mut in_memory_hashes);
105 Ok(hashes)
106 }
107}
108
109impl<N: NodePrimitives> AccountReader for MemoryOverlayStateProviderRef<'_, N> {
110 fn basic_account(&self, address: &Address) -> ProviderResult<Option<Account>> {
111 for block in self.in_memory.iter() {
112 if let Some(account) = block.execution_output.account(address) {
113 return Ok(account);
114 }
115 }
116
117 self.historical.basic_account(address)
118 }
119}
120
121impl<N: NodePrimitives> StateRootProvider for MemoryOverlayStateProviderRef<'_, N> {
122 fn state_root(&self, state: HashedPostState) -> ProviderResult<B256> {
123 self.state_root_from_nodes(TrieInput::from_state(state))
124 }
125
126 fn state_root_from_nodes(&self, mut input: TrieInput) -> ProviderResult<B256> {
127 input.prepend_self(self.trie_input().clone());
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 input.prepend_self(self.trie_input().clone());
143 self.historical.state_root_from_nodes_with_updates(input)
144 }
145}
146
147impl<N: NodePrimitives> StorageRootProvider for MemoryOverlayStateProviderRef<'_, N> {
148 fn storage_root(&self, address: Address, storage: HashedStorage) -> ProviderResult<B256> {
150 let merged = self.merged_hashed_storage(address, storage);
151 self.historical.storage_root(address, merged)
152 }
153
154 fn storage_proof(
156 &self,
157 address: Address,
158 slot: B256,
159 storage: HashedStorage,
160 ) -> ProviderResult<reth_trie::StorageProof> {
161 let merged = self.merged_hashed_storage(address, storage);
162 self.historical.storage_proof(address, slot, merged)
163 }
164
165 fn storage_multiproof(
167 &self,
168 address: Address,
169 slots: &[B256],
170 storage: HashedStorage,
171 ) -> ProviderResult<StorageMultiProof> {
172 let merged = self.merged_hashed_storage(address, storage);
173 self.historical.storage_multiproof(address, slots, merged)
174 }
175}
176
177impl<N: NodePrimitives> StateProofProvider for MemoryOverlayStateProviderRef<'_, N> {
178 fn proof(
179 &self,
180 mut input: TrieInput,
181 address: Address,
182 slots: &[B256],
183 ) -> ProviderResult<AccountProof> {
184 input.prepend_self(self.trie_input().clone());
185 self.historical.proof(input, address, slots)
186 }
187
188 fn multiproof(
189 &self,
190 mut input: TrieInput,
191 targets: MultiProofTargets,
192 ) -> ProviderResult<MultiProof> {
193 input.prepend_self(self.trie_input().clone());
194 self.historical.multiproof(input, targets)
195 }
196
197 fn witness(&self, mut input: TrieInput, target: HashedPostState) -> ProviderResult<Vec<Bytes>> {
198 input.prepend_self(self.trie_input().clone());
199 self.historical.witness(input, target)
200 }
201}
202
203impl<N: NodePrimitives> HashedPostStateProvider for MemoryOverlayStateProviderRef<'_, N> {
204 fn hashed_post_state(&self, bundle_state: &BundleState) -> HashedPostState {
205 self.historical.hashed_post_state(bundle_state)
206 }
207}
208
209impl<N: NodePrimitives> StateProvider for MemoryOverlayStateProviderRef<'_, N> {
210 fn storage(
211 &self,
212 address: Address,
213 storage_key: StorageKey,
214 ) -> ProviderResult<Option<StorageValue>> {
215 for block in self.in_memory.iter() {
216 if let Some(value) = block.execution_output.storage(&address, storage_key.into()) {
217 return Ok(Some(value));
218 }
219 }
220
221 self.historical.storage(address, storage_key)
222 }
223}
224
225impl<N: NodePrimitives> BytecodeReader for MemoryOverlayStateProviderRef<'_, N> {
226 fn bytecode_by_hash(&self, code_hash: &B256) -> ProviderResult<Option<Bytecode>> {
227 for block in self.in_memory.iter() {
228 if let Some(contract) = block.execution_output.bytecode(code_hash) {
229 return Ok(Some(contract));
230 }
231 }
232
233 self.historical.bytecode_by_hash(code_hash)
234 }
235}
236
237#[expect(missing_debug_implementations)]
240pub struct MemoryOverlayStateProvider<N: NodePrimitives = reth_ethereum_primitives::EthPrimitives> {
241 pub(crate) historical: StateProviderBox,
243 pub(crate) in_memory: Vec<ExecutedBlock<N>>,
245 pub(crate) trie_input: OnceLock<TrieInput>,
247}
248
249impl<N: NodePrimitives> MemoryOverlayStateProvider<N> {
250 pub fn new(historical: StateProviderBox, in_memory: Vec<ExecutedBlock<N>>) -> Self {
258 Self { historical, in_memory, trie_input: OnceLock::new() }
259 }
260
261 #[inline(always)]
263 fn as_ref(&self) -> MemoryOverlayStateProviderRef<'_, N> {
264 MemoryOverlayStateProviderRef {
265 historical: Box::new(self.historical.as_ref()),
266 in_memory: Cow::Borrowed(&self.in_memory),
267 trie_input: self.trie_input.clone(),
268 }
269 }
270
271 pub fn boxed(self) -> StateProviderBox {
273 Box::new(self)
274 }
275}
276
277reth_storage_api::macros::delegate_provider_impls!(MemoryOverlayStateProvider<N> where [N: NodePrimitives]);