Skip to main content

reth_chain_state/
lazy_overlay.rs

1//! Lazy overlay computation for trie input.
2//!
3//! This module provides [`LazyOverlay`], a type that computes the [`TrieInputSorted`]
4//! lazily on first access. This allows execution to start before the trie overlay
5//! is fully computed.
6
7use crate::{EthPrimitives, ExecutedBlock};
8use alloy_primitives::B256;
9use reth_primitives_traits::{
10    dashmap::{self, DashMap},
11    AlloyBlockHeader, NodePrimitives,
12};
13use reth_trie::{updates::TrieUpdatesSorted, HashedPostStateSorted, TrieInputSorted};
14use std::sync::Arc;
15use tracing::{debug, trace};
16
17/// Inputs captured for lazy overlay computation.
18#[derive(Clone)]
19struct LazyOverlayInputs<N: NodePrimitives = EthPrimitives> {
20    /// In-memory blocks from tip to anchor child.
21    ///
22    /// Blocks must be provided in reverse chain order (newest to oldest).
23    blocks: Vec<ExecutedBlock<N>>,
24}
25
26/// Lazily computed trie overlay.
27///
28/// Captures the inputs needed to compute a [`TrieInputSorted`] and defers the actual
29/// computation until first access.
30///
31/// Blocks must be provided in reverse chain order (newest to oldest), so the first block is the
32/// chain tip and the last block is the oldest in-memory block in the chain segment.
33///
34/// # Fast Path vs Slow Path
35///
36/// - **Fast path**: If the tip block's cached `anchored_trie_input` is ready and its `anchor_hash`
37///   matches our expected anchor, we can reuse it directly (O(1)).
38/// - **Slow path**: Otherwise, we merge all ancestor blocks' trie data into a new overlay.
39#[derive(Clone)]
40pub struct LazyOverlay<N: NodePrimitives = EthPrimitives> {
41    /// Computed results, cached by requested anchor hash.
42    inner: Arc<DashMap<B256, Arc<TrieInputSorted>>>,
43    /// Inputs for lazy computation.
44    inputs: LazyOverlayInputs<N>,
45}
46
47impl<N: NodePrimitives> std::fmt::Debug for LazyOverlay<N> {
48    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49        f.debug_struct("LazyOverlay")
50            .field(
51                "oldest_block_parent_hash",
52                &self.inputs.blocks.last().map(|block| block.recovered_block().parent_hash()),
53            )
54            .field("num_blocks", &self.inputs.blocks.len())
55            .field("cached_anchors", &self.inner.len())
56            .finish()
57    }
58}
59
60impl<N: NodePrimitives> LazyOverlay<N> {
61    /// Create a new lazy overlay from in-memory blocks.
62    ///
63    /// # Arguments
64    ///
65    /// * `blocks` - Executed blocks in reverse chain order (newest to oldest)
66    pub fn new(blocks: Vec<ExecutedBlock<N>>) -> Self {
67        debug_assert!(
68            blocks.windows(2).all(|window| {
69                window[0].recovered_block().parent_hash() == window[1].recovered_block().hash()
70            }),
71            "LazyOverlay blocks must be ordered newest to oldest along a single chain"
72        );
73
74        Self { inner: Default::default(), inputs: LazyOverlayInputs { blocks } }
75    }
76
77    /// Returns the number of in-memory blocks this overlay covers.
78    pub const fn num_blocks(&self) -> usize {
79        self.inputs.blocks.len()
80    }
81
82    /// Returns the oldest anchor hash this overlay can serve.
83    ///
84    /// This is the parent hash of the oldest block in the stored newest-to-oldest chain segment.
85    pub fn anchor_hash(&self) -> Option<B256> {
86        self.inputs.blocks.last().map(|block| block.recovered_block().parent_hash())
87    }
88
89    /// Returns true if there are no blocks in the overlay, or if one of the blocks has the given
90    /// hash as a parent hash.
91    pub fn has_anchor_hash(&self, hash: B256) -> bool {
92        self.inputs.blocks.is_empty() ||
93            self.inputs.blocks.iter().any(|b| b.recovered_block().parent_hash() == hash)
94    }
95
96    #[cfg(test)]
97    /// Returns true if the overlay has already been computed for the requested anchor.
98    pub fn is_computed(&self, anchor_hash: B256) -> bool {
99        self.inner.contains_key(&anchor_hash)
100    }
101
102    /// Returns the computed trie input for the requested anchor, computing it if necessary.
103    ///
104    /// The first call triggers computation (which may block waiting for deferred data).
105    /// Subsequent calls for the same anchor return the cached result immediately.
106    pub fn get(&self, anchor_hash: B256) -> Arc<TrieInputSorted> {
107        match self.inner.entry(anchor_hash) {
108            dashmap::Entry::Occupied(entry) => Arc::clone(entry.get()),
109            dashmap::Entry::Vacant(entry) => {
110                let input = self.compute(anchor_hash);
111                entry.insert(Arc::clone(&input));
112                input
113            }
114        }
115    }
116
117    /// Returns the overlay as (nodes, state) tuple for use with `OverlayStateProviderFactory`.
118    pub fn as_overlay(
119        &self,
120        anchor_hash: B256,
121    ) -> (Arc<TrieUpdatesSorted>, Arc<HashedPostStateSorted>) {
122        let input = self.get(anchor_hash);
123        (Arc::clone(&input.nodes), Arc::clone(&input.state))
124    }
125
126    /// Compute the trie input overlay.
127    fn compute(&self, anchor_hash: B256) -> Arc<TrieInputSorted> {
128        let blocks = &self.inputs.blocks;
129        if blocks.is_empty() {
130            return Default::default()
131        }
132
133        let Some(last_index) =
134            blocks.iter().position(|block| block.recovered_block().parent_hash() == anchor_hash)
135        else {
136            panic!(
137                "LazyOverlay does not contain a block whose parent hash matches requested anchor {anchor_hash}"
138            );
139        };
140        let blocks = &blocks[..=last_index];
141
142        // Fast path: Check if tip block's overlay is ready and anchor matches.
143        // The tip block (first in list) has the cumulative overlay from all ancestors up to the
144        // requested anchor.
145        if let Some(tip) = blocks.first() {
146            let data = tip.trie_data();
147            if let Some(anchored) = &data.anchored_trie_input {
148                if anchored.anchor_hash == anchor_hash {
149                    trace!(target: "chain_state::lazy_overlay", %anchor_hash, "Reusing tip block's cached overlay (fast path)");
150                    return Arc::clone(&anchored.trie_input);
151                }
152                debug!(
153                    target: "chain_state::lazy_overlay",
154                    computed_anchor = %anchored.anchor_hash,
155                    %anchor_hash,
156                    "Anchor mismatch, falling back to merge"
157                );
158            }
159        }
160
161        // Slow path: Merge the prefix of blocks from the tip back to the requested anchor.
162        debug!(
163            target: "chain_state::lazy_overlay",
164            %anchor_hash,
165            num_blocks = blocks.len(),
166            "Merging blocks (slow path)"
167        );
168        Arc::new(Self::merge_blocks(blocks))
169    }
170
171    /// Merge all blocks' trie data into a single [`TrieInputSorted`].
172    ///
173    /// Blocks are ordered newest to oldest.
174    fn merge_blocks(blocks: &[ExecutedBlock<N>]) -> TrieInputSorted {
175        if blocks.is_empty() {
176            return TrieInputSorted::default();
177        }
178
179        let state = HashedPostStateSorted::merge_batch(
180            blocks.iter().map(|block| block.trie_data().hashed_state),
181        );
182        let nodes = TrieUpdatesSorted::merge_batch(
183            blocks.iter().map(|block| block.trie_data().trie_updates),
184        );
185
186        TrieInputSorted { state, nodes, prefix_sets: Default::default() }
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193    use crate::{test_utils::TestBlockBuilder, ComputedTrieData, EthPrimitives, ExecutedBlock};
194    use alloy_primitives::U256;
195    use reth_primitives_traits::Account;
196    use reth_trie::{updates::TrieUpdatesSorted, HashedPostState, HashedStorage};
197    use std::sync::Arc;
198
199    fn with_unique_state(
200        block: &ExecutedBlock<EthPrimitives>,
201        id: u8,
202    ) -> ExecutedBlock<EthPrimitives> {
203        let hashed_address = B256::with_last_byte(id);
204        let hashed_slot = B256::with_last_byte(id.saturating_add(32));
205        let hashed_state = HashedPostState::default()
206            .with_accounts([(hashed_address, Some(Account::default()))])
207            .with_storages([(
208                hashed_address,
209                HashedStorage::from_iter(false, [(hashed_slot, U256::from(id))]),
210            )])
211            .into_sorted();
212
213        ExecutedBlock::new(
214            Arc::clone(&block.recovered_block),
215            Arc::clone(&block.execution_output),
216            ComputedTrieData::without_trie_input(
217                Arc::new(hashed_state),
218                Arc::new(TrieUpdatesSorted::default()),
219            ),
220        )
221    }
222
223    fn test_blocks() -> Vec<ExecutedBlock<EthPrimitives>> {
224        TestBlockBuilder::eth()
225            .get_executed_blocks(1..4)
226            .collect::<Vec<_>>()
227            .into_iter()
228            .rev()
229            .enumerate()
230            .map(|(index, block)| with_unique_state(&block, index as u8 + 1))
231            .collect()
232    }
233
234    #[test]
235    fn single_block_uses_data_directly() {
236        let block = TestBlockBuilder::eth().get_executed_block_with_number(1, B256::random());
237        let anchor_hash = block.recovered_block().parent_hash();
238        let overlay = LazyOverlay::new(vec![block]);
239
240        assert!(!overlay.is_computed(anchor_hash));
241        let _ = overlay.get(anchor_hash);
242        assert!(overlay.is_computed(anchor_hash));
243    }
244
245    #[test]
246    fn caches_results_per_anchor() {
247        let blocks = test_blocks();
248        let prefix_anchor = blocks[2].recovered_block().hash();
249        let full_anchor = blocks[2].recovered_block().parent_hash();
250        let overlay = LazyOverlay::new(blocks);
251
252        let prefix = overlay.get(prefix_anchor);
253        let full = overlay.get(full_anchor);
254
255        assert!(overlay.is_computed(prefix_anchor));
256        assert!(overlay.is_computed(full_anchor));
257        assert!(!Arc::ptr_eq(&prefix, &full));
258        assert!(Arc::ptr_eq(&prefix, &overlay.get(prefix_anchor)));
259        assert!(Arc::ptr_eq(&full, &overlay.get(full_anchor)));
260    }
261
262    #[test]
263    fn requested_anchor_limits_the_merged_prefix() {
264        let blocks = test_blocks();
265        let prefix_anchor = blocks[2].recovered_block().hash();
266        let expected = LazyOverlay::merge_blocks(&blocks[..2]);
267        let overlay = LazyOverlay::new(blocks);
268        let actual = overlay.get(prefix_anchor);
269
270        assert_eq!(actual.nodes.as_ref(), expected.nodes.as_ref());
271        assert_eq!(actual.state.as_ref(), expected.state.as_ref());
272    }
273
274    #[test]
275    fn anchor_hash_returns_oldest_served_anchor() {
276        let blocks = test_blocks();
277        let expected_anchor = blocks.last().unwrap().recovered_block().parent_hash();
278        let overlay = LazyOverlay::new(blocks);
279
280        assert_eq!(overlay.anchor_hash(), Some(expected_anchor));
281    }
282
283    #[test]
284    fn reuses_tip_overlay_when_anchor_matches() {
285        let mut blocks = test_blocks();
286        let prefix_anchor = blocks[2].recovered_block().hash();
287        let tip_overlay = Arc::new(LazyOverlay::merge_blocks(&blocks[..2]));
288        let tip_data = blocks[0].trie_data();
289
290        blocks[0] = ExecutedBlock::new(
291            Arc::clone(&blocks[0].recovered_block),
292            Arc::clone(&blocks[0].execution_output),
293            ComputedTrieData::with_trie_input(
294                tip_data.hashed_state,
295                tip_data.trie_updates,
296                prefix_anchor,
297                Arc::clone(&tip_overlay),
298            ),
299        );
300
301        let overlay = LazyOverlay::new(blocks);
302        let actual = overlay.get(prefix_anchor);
303
304        assert!(Arc::ptr_eq(&actual, &tip_overlay));
305    }
306
307    #[test]
308    #[should_panic(
309        expected = "LazyOverlay does not contain a block whose parent hash matches requested anchor"
310    )]
311    fn missing_anchor_panics() {
312        let blocks = test_blocks();
313        let missing_anchor = blocks[0].recovered_block().hash();
314        let overlay = LazyOverlay::new(blocks);
315
316        let _ = overlay.get(missing_anchor);
317    }
318
319    #[test]
320    #[should_panic(
321        expected = "LazyOverlay blocks must be ordered newest to oldest along a single chain"
322    )]
323    fn misordered_blocks_panic() {
324        let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..3).collect();
325        let _ = LazyOverlay::new(blocks);
326    }
327}