Skip to main content

reth_chain_state/
in_memory.rs

1//! Types for tracking the canonical chain state in memory.
2
3use crate::{
4    CanonStateNotification, CanonStateNotificationSender, CanonStateNotifications,
5    ChainInfoTracker, MemoryOverlayStateProvider,
6};
7use alloy_consensus::{transaction::TransactionMeta, BlockHeader};
8use alloy_eips::{BlockHashOrNumber, BlockNumHash};
9use alloy_primitives::{map::B256Map, BlockNumber, TxHash, B256};
10use parking_lot::RwLock;
11use reth_chainspec::ChainInfo;
12use reth_ethereum_primitives::EthPrimitives;
13use reth_execution_types::{BlockExecutionOutput, BlockExecutionResult, Chain, ExecutionOutcome};
14use reth_metrics::{metrics::Gauge, Metrics};
15use reth_primitives_traits::{
16    BlockBody as _, IndexedTx, NodePrimitives, RecoveredBlock, SealedBlock, SealedHeader,
17    SignedTransaction,
18};
19use reth_storage_api::StateProviderBox;
20use reth_trie::{
21    updates::TrieUpdatesSorted, ComputedTrieData, HashedPostStateSorted, LazyTrieData,
22};
23use std::{collections::BTreeMap, sync::Arc, time::Instant};
24use tokio::sync::{broadcast, watch};
25
26/// Size of the broadcast channel used to notify canonical state events.
27const CANON_STATE_NOTIFICATION_CHANNEL_SIZE: usize = 256;
28
29/// Metrics for the in-memory state.
30#[derive(Metrics)]
31#[metrics(scope = "blockchain_tree.in_mem_state")]
32pub(crate) struct InMemoryStateMetrics {
33    /// The block number of the earliest block in the in-memory state.
34    pub(crate) earliest_block: Gauge,
35    /// The block number of the latest block in the in-memory state.
36    pub(crate) latest_block: Gauge,
37    /// The number of blocks in the in-memory state.
38    pub(crate) num_blocks: Gauge,
39}
40
41/// Container type for in memory state data of the canonical chain.
42///
43/// This tracks blocks and their state that haven't been persisted to disk yet but are part of the
44/// canonical chain that can be traced back to a canonical block on disk.
45///
46/// # Locking behavior on state updates
47///
48/// All update calls must acquire all locks at once before modifying state to ensure the internal
49/// state remains consistent. This prevents readers from observing partially updated state where
50/// the numbers and blocks maps are out of sync.
51/// Update functions ensure that the numbers write lock is always acquired first, because lookup by
52/// numbers first read the numbers map and then the blocks map.
53/// By acquiring the numbers lock first, we ensure that read-only lookups don't deadlock updates.
54/// This holds, because only lookup by number functions need to acquire the numbers lock first to
55/// get the block hash.
56#[derive(Debug, Default)]
57pub(crate) struct InMemoryState<N: NodePrimitives = EthPrimitives> {
58    /// All canonical blocks that are not on disk yet.
59    blocks: RwLock<B256Map<Arc<BlockState<N>>>>,
60    /// Mapping of block numbers to block hashes.
61    numbers: RwLock<BTreeMap<u64, B256>>,
62    /// The pending block that has not yet been made canonical.
63    pending: watch::Sender<Option<BlockState<N>>>,
64    /// Metrics for the in-memory state.
65    metrics: InMemoryStateMetrics,
66}
67
68impl<N: NodePrimitives> InMemoryState<N> {
69    pub(crate) fn new(
70        blocks: B256Map<Arc<BlockState<N>>>,
71        numbers: BTreeMap<u64, B256>,
72        pending: Option<BlockState<N>>,
73    ) -> Self {
74        let (pending, _) = watch::channel(pending);
75        let this = Self {
76            blocks: RwLock::new(blocks),
77            numbers: RwLock::new(numbers),
78            pending,
79            metrics: Default::default(),
80        };
81        this.update_metrics();
82        this
83    }
84
85    /// Update the metrics for the in-memory state.
86    ///
87    /// # Locking behavior
88    ///
89    /// This tries to acquire a read lock. Drop any write locks before calling this.
90    pub(crate) fn update_metrics(&self) {
91        let (count, earliest, latest) = {
92            let numbers = self.numbers.read();
93            let count = numbers.len();
94            let earliest = numbers.first_key_value().map(|(number, _)| *number);
95            let latest = numbers.last_key_value().map(|(number, _)| *number);
96            (count, earliest, latest)
97        };
98        if let Some(earliest_block_number) = earliest {
99            self.metrics.earliest_block.set(earliest_block_number as f64);
100        }
101        if let Some(latest_block_number) = latest {
102            self.metrics.latest_block.set(latest_block_number as f64);
103        }
104        self.metrics.num_blocks.set(count as f64);
105    }
106
107    /// Returns the state for a given block hash.
108    pub(crate) fn state_by_hash(&self, hash: B256) -> Option<Arc<BlockState<N>>> {
109        self.blocks.read().get(&hash).cloned()
110    }
111
112    /// Returns the state for a given block number.
113    pub(crate) fn state_by_number(&self, number: u64) -> Option<Arc<BlockState<N>>> {
114        let hash = self.hash_by_number(number)?;
115        self.state_by_hash(hash)
116    }
117
118    /// Returns the hash for a specific block number
119    pub(crate) fn hash_by_number(&self, number: u64) -> Option<B256> {
120        self.numbers.read().get(&number).copied()
121    }
122
123    /// Returns the current chain head state.
124    pub(crate) fn head_state(&self) -> Option<Arc<BlockState<N>>> {
125        let hash = *self.numbers.read().last_key_value()?.1;
126        self.state_by_hash(hash)
127    }
128
129    /// Returns the pending state corresponding to the current head plus one,
130    /// from the payload received in newPayload that does not have a FCU yet.
131    pub(crate) fn pending_state(&self) -> Option<BlockState<N>> {
132        self.pending.borrow().clone()
133    }
134
135    #[cfg(test)]
136    fn block_count(&self) -> usize {
137        self.blocks.read().len()
138    }
139}
140
141/// Inner type to provide in memory state. It includes a chain tracker to be
142/// advanced internally by the tree.
143#[derive(Debug)]
144pub(crate) struct CanonicalInMemoryStateInner<N: NodePrimitives> {
145    /// Tracks certain chain information, such as the canonical head, safe head, and finalized
146    /// head.
147    pub(crate) chain_info_tracker: ChainInfoTracker<N>,
148    /// Tracks blocks at the tip of the chain that have not been persisted to disk yet.
149    pub(crate) in_memory_state: InMemoryState<N>,
150    /// A broadcast stream that emits events when the canonical chain is updated.
151    pub(crate) canon_state_notification_sender: CanonStateNotificationSender<N>,
152}
153
154impl<N: NodePrimitives> CanonicalInMemoryStateInner<N> {
155    /// Clears all entries in the in memory state.
156    fn clear(&self) {
157        {
158            // acquire locks, starting with the numbers lock
159            let mut numbers = self.in_memory_state.numbers.write();
160            let mut blocks = self.in_memory_state.blocks.write();
161            numbers.clear();
162            blocks.clear();
163            self.in_memory_state.pending.send_modify(|p| {
164                p.take();
165            });
166        }
167        self.in_memory_state.update_metrics();
168    }
169}
170
171type PendingBlockAndReceipts<N> =
172    (RecoveredBlock<<N as NodePrimitives>::Block>, Vec<reth_primitives_traits::ReceiptTy<N>>);
173
174/// This type is responsible for providing the blocks, receipts, and state for
175/// all canonical blocks not on disk yet and keeps track of the block range that
176/// is in memory.
177#[derive(Debug, Clone)]
178pub struct CanonicalInMemoryState<N: NodePrimitives = EthPrimitives> {
179    pub(crate) inner: Arc<CanonicalInMemoryStateInner<N>>,
180}
181
182impl<N: NodePrimitives> CanonicalInMemoryState<N> {
183    /// Create a new in-memory state with the given blocks, numbers, pending state, and optional
184    /// finalized header.
185    pub fn new(
186        blocks: B256Map<Arc<BlockState<N>>>,
187        numbers: BTreeMap<u64, B256>,
188        pending: Option<BlockState<N>>,
189        finalized: Option<SealedHeader<N::BlockHeader>>,
190        safe: Option<SealedHeader<N::BlockHeader>>,
191    ) -> Self {
192        let in_memory_state = InMemoryState::new(blocks, numbers, pending);
193        let header = in_memory_state.head_state().map_or_else(SealedHeader::default, |state| {
194            state.block_ref().recovered_block().clone_sealed_header()
195        });
196        let chain_info_tracker = ChainInfoTracker::new(header, finalized, safe);
197        let (canon_state_notification_sender, _) =
198            broadcast::channel(CANON_STATE_NOTIFICATION_CHANNEL_SIZE);
199
200        Self {
201            inner: Arc::new(CanonicalInMemoryStateInner {
202                chain_info_tracker,
203                in_memory_state,
204                canon_state_notification_sender,
205            }),
206        }
207    }
208
209    /// Create an empty state.
210    pub fn empty() -> Self {
211        Self::new(B256Map::default(), BTreeMap::new(), None, None, None)
212    }
213
214    /// Create a new in memory state with the given local head and finalized header
215    /// if it exists.
216    pub fn with_head(
217        head: SealedHeader<N::BlockHeader>,
218        finalized: Option<SealedHeader<N::BlockHeader>>,
219        safe: Option<SealedHeader<N::BlockHeader>>,
220    ) -> Self {
221        let chain_info_tracker = ChainInfoTracker::new(head, finalized, safe);
222        let in_memory_state = InMemoryState::default();
223        let (canon_state_notification_sender, _) =
224            broadcast::channel(CANON_STATE_NOTIFICATION_CHANNEL_SIZE);
225        let inner = CanonicalInMemoryStateInner {
226            chain_info_tracker,
227            in_memory_state,
228            canon_state_notification_sender,
229        };
230
231        Self { inner: Arc::new(inner) }
232    }
233
234    /// Returns the block hash corresponding to the given number.
235    pub fn hash_by_number(&self, number: u64) -> Option<B256> {
236        self.inner.in_memory_state.hash_by_number(number)
237    }
238
239    /// Returns the header corresponding to the given hash.
240    pub fn header_by_hash(&self, hash: B256) -> Option<SealedHeader<N::BlockHeader>> {
241        self.state_by_hash(hash)
242            .map(|block| block.block_ref().recovered_block().clone_sealed_header())
243    }
244
245    /// Clears all entries in the in memory state.
246    pub fn clear_state(&self) {
247        self.inner.clear()
248    }
249
250    /// Updates the pending block with the given block.
251    ///
252    /// Note: This assumes that the parent block of the pending block is canonical.
253    pub fn set_pending_block(&self, pending: ExecutedBlock<N>) {
254        // fetch the state of the pending block's parent block
255        let parent = self.state_by_hash(pending.recovered_block().parent_hash());
256        let pending = BlockState::with_parent(pending, parent);
257        self.inner.in_memory_state.pending.send_modify(|p| {
258            p.replace(pending);
259        });
260        self.inner.in_memory_state.update_metrics();
261    }
262
263    /// Append new blocks to the in memory state.
264    ///
265    /// This removes all reorged blocks and appends the new blocks to the tracked chain and connects
266    /// them to their parent blocks.
267    fn update_blocks<I, R>(&self, new_blocks: I, reorged: R)
268    where
269        I: IntoIterator<Item = ExecutedBlock<N>>,
270        R: IntoIterator<Item = ExecutedBlock<N>>,
271    {
272        {
273            // acquire locks, starting with the numbers lock
274            let mut numbers = self.inner.in_memory_state.numbers.write();
275            let mut blocks = self.inner.in_memory_state.blocks.write();
276
277            // we first remove the blocks from the reorged chain
278            for block in reorged {
279                let hash = block.recovered_block().hash();
280                let number = block.recovered_block().number();
281                blocks.remove(&hash);
282                numbers.remove(&number);
283            }
284
285            // insert the new blocks
286            for block in new_blocks {
287                let parent = blocks.get(&block.recovered_block().parent_hash()).cloned();
288                let block_state = BlockState::with_parent(block, parent);
289                let hash = block_state.hash();
290                let number = block_state.number();
291
292                // append new blocks
293                blocks.insert(hash, Arc::new(block_state));
294                numbers.insert(number, hash);
295            }
296
297            // remove the pending state
298            self.inner.in_memory_state.pending.send_modify(|p| {
299                p.take();
300            });
301        }
302        self.inner.in_memory_state.update_metrics();
303    }
304
305    /// Update the in memory state with the given chain update.
306    pub fn update_chain(&self, new_chain: NewCanonicalChain<N>) {
307        match new_chain {
308            NewCanonicalChain::Commit { new } => {
309                self.update_blocks(new, vec![]);
310            }
311            NewCanonicalChain::Reorg { new, old } => {
312                self.update_blocks(new, old);
313            }
314        }
315    }
316
317    /// Removes blocks from the in memory state that are persisted to the given height.
318    ///
319    /// This will update the links between blocks and remove all blocks that are [..
320    /// `persisted_height`].
321    pub fn remove_persisted_blocks(&self, persisted_num_hash: BlockNumHash) {
322        self.remove_persisted_blocks_until(persisted_num_hash, persisted_num_hash.number);
323    }
324
325    /// Removes blocks from the in-memory state through `remove_until` while still reporting the
326    /// provided block as the persisted tip.
327    pub fn remove_persisted_blocks_until(
328        &self,
329        persisted_num_hash: BlockNumHash,
330        remove_until: BlockNumber,
331    ) {
332        self.set_persisted(persisted_num_hash);
333        // if the persisted hash is not in the canonical in memory state, do nothing, because it
334        // means canonical blocks were not actually persisted.
335        //
336        // This can happen if the persistence task takes a long time, while a reorg is happening.
337        {
338            if self.inner.in_memory_state.blocks.read().get(&persisted_num_hash.hash).is_none() {
339                // do nothing
340                return
341            }
342        }
343
344        {
345            // acquire locks, starting with the numbers lock
346            let mut numbers = self.inner.in_memory_state.numbers.write();
347            let mut blocks = self.inner.in_memory_state.blocks.write();
348
349            let remove_until = remove_until.min(persisted_num_hash.number);
350
351            // clear all numbers
352            numbers.clear();
353
354            // Drain all blocks and keep only the suffix that still has to stay in memory.
355            let mut old_blocks = blocks
356                .drain()
357                .filter(|(_, b)| b.block_ref().recovered_block().number() > remove_until)
358                .map(|(_, b)| b.block.clone())
359                .collect::<Vec<_>>();
360
361            // sort the blocks by number so we can insert them back in natural order (low -> high)
362            old_blocks.sort_unstable_by_key(|block| block.recovered_block().number());
363
364            // re-insert the blocks in natural order and connect them to their parent blocks
365            for block in old_blocks {
366                let parent = blocks.get(&block.recovered_block().parent_hash()).cloned();
367                let block_state = BlockState::with_parent(block, parent);
368                let hash = block_state.hash();
369                let number = block_state.number();
370
371                // append new blocks
372                blocks.insert(hash, Arc::new(block_state));
373                numbers.insert(number, hash);
374            }
375
376            // also shift the pending state if it exists
377            self.inner.in_memory_state.pending.send_modify(|p| {
378                if let Some(p) = p.as_mut() {
379                    p.parent = blocks.get(&p.block_ref().recovered_block().parent_hash()).cloned();
380                }
381            });
382        }
383        self.inner.in_memory_state.update_metrics();
384    }
385
386    /// Returns in memory state corresponding the given hash.
387    pub fn state_by_hash(&self, hash: B256) -> Option<Arc<BlockState<N>>> {
388        self.inner.in_memory_state.state_by_hash(hash)
389    }
390
391    /// Returns in memory state corresponding the block number.
392    pub fn state_by_number(&self, number: u64) -> Option<Arc<BlockState<N>>> {
393        self.inner.in_memory_state.state_by_number(number)
394    }
395
396    /// Returns the in memory head state.
397    pub fn head_state(&self) -> Option<Arc<BlockState<N>>> {
398        self.inner.in_memory_state.head_state()
399    }
400
401    /// Returns the in memory pending state.
402    pub fn pending_state(&self) -> Option<BlockState<N>> {
403        self.inner.in_memory_state.pending_state()
404    }
405
406    /// Returns the in memory pending `BlockNumHash`.
407    pub fn pending_block_num_hash(&self) -> Option<BlockNumHash> {
408        self.inner
409            .in_memory_state
410            .pending_state()
411            .map(|state| BlockNumHash { number: state.number(), hash: state.hash() })
412    }
413
414    /// Returns the current `ChainInfo`.
415    pub fn chain_info(&self) -> ChainInfo {
416        self.inner.chain_info_tracker.chain_info()
417    }
418
419    /// Returns the latest canonical block number.
420    pub fn get_canonical_block_number(&self) -> u64 {
421        self.inner.chain_info_tracker.get_canonical_block_number()
422    }
423
424    /// Returns the `BlockNumHash` of the safe head.
425    pub fn get_safe_num_hash(&self) -> Option<BlockNumHash> {
426        self.inner.chain_info_tracker.get_safe_num_hash()
427    }
428
429    /// Returns the `BlockNumHash` of the finalized head.
430    pub fn get_finalized_num_hash(&self) -> Option<BlockNumHash> {
431        self.inner.chain_info_tracker.get_finalized_num_hash()
432    }
433
434    /// Hook for new fork choice update.
435    pub fn on_forkchoice_update_received(&self) {
436        self.inner.chain_info_tracker.on_forkchoice_update_received();
437    }
438
439    /// Returns the timestamp of the last received update.
440    pub fn last_received_update_timestamp(&self) -> Option<Instant> {
441        self.inner.chain_info_tracker.last_forkchoice_update_received_at()
442    }
443
444    /// Canonical head setter.
445    pub fn set_canonical_head(&self, header: SealedHeader<N::BlockHeader>) {
446        self.inner.chain_info_tracker.set_canonical_head(header);
447    }
448
449    /// Safe head setter.
450    pub fn set_safe(&self, header: SealedHeader<N::BlockHeader>) {
451        self.inner.chain_info_tracker.set_safe(header);
452    }
453
454    /// Finalized head setter.
455    pub fn set_finalized(&self, header: SealedHeader<N::BlockHeader>) {
456        self.inner.chain_info_tracker.set_finalized(header);
457    }
458
459    /// Persisted block setter.
460    pub fn set_persisted(&self, num_hash: BlockNumHash) {
461        self.inner.chain_info_tracker.set_persisted(num_hash);
462    }
463
464    /// Canonical head getter.
465    pub fn get_canonical_head(&self) -> SealedHeader<N::BlockHeader> {
466        self.inner.chain_info_tracker.get_canonical_head()
467    }
468
469    /// Finalized header getter.
470    pub fn get_finalized_header(&self) -> Option<SealedHeader<N::BlockHeader>> {
471        self.inner.chain_info_tracker.get_finalized_header()
472    }
473
474    /// Safe header getter.
475    pub fn get_safe_header(&self) -> Option<SealedHeader<N::BlockHeader>> {
476        self.inner.chain_info_tracker.get_safe_header()
477    }
478
479    /// Persisted block `BlockNumHash` getter.
480    pub fn get_persisted_num_hash(&self) -> Option<BlockNumHash> {
481        self.inner.chain_info_tracker.get_persisted_num_hash()
482    }
483
484    /// Returns the `SealedHeader` corresponding to the pending state.
485    pub fn pending_sealed_header(&self) -> Option<SealedHeader<N::BlockHeader>> {
486        self.pending_state().map(|h| h.block_ref().recovered_block().clone_sealed_header())
487    }
488
489    /// Returns the `Header` corresponding to the pending state.
490    pub fn pending_header(&self) -> Option<N::BlockHeader> {
491        self.pending_sealed_header().map(|sealed_header| sealed_header.unseal())
492    }
493
494    /// Returns the `SealedBlock` corresponding to the pending state.
495    pub fn pending_block(&self) -> Option<SealedBlock<N::Block>> {
496        self.pending_state()
497            .map(|block_state| block_state.block_ref().recovered_block().sealed_block().clone())
498    }
499
500    /// Returns the `RecoveredBlock` corresponding to the pending state.
501    pub fn pending_recovered_block(&self) -> Option<RecoveredBlock<N::Block>>
502    where
503        N::SignedTx: SignedTransaction,
504    {
505        self.pending_state().map(|block_state| block_state.block_ref().recovered_block().clone())
506    }
507
508    /// Returns a tuple with the `SealedBlock` corresponding to the pending
509    /// state and a vector of its `Receipt`s.
510    pub fn pending_block_and_receipts(&self) -> Option<PendingBlockAndReceipts<N>> {
511        self.pending_state().map(|block_state| {
512            (
513                block_state.block_ref().recovered_block().clone(),
514                block_state.executed_block_receipts(),
515            )
516        })
517    }
518
519    /// Subscribe to new blocks events.
520    pub fn subscribe_canon_state(&self) -> CanonStateNotifications<N> {
521        self.inner.canon_state_notification_sender.subscribe()
522    }
523
524    /// Subscribe to new safe block events.
525    pub fn subscribe_safe_block(&self) -> watch::Receiver<Option<SealedHeader<N::BlockHeader>>> {
526        self.inner.chain_info_tracker.subscribe_safe_block()
527    }
528
529    /// Subscribe to new finalized block events.
530    pub fn subscribe_finalized_block(
531        &self,
532    ) -> watch::Receiver<Option<SealedHeader<N::BlockHeader>>> {
533        self.inner.chain_info_tracker.subscribe_finalized_block()
534    }
535
536    /// Subscribe to new persisted block events.
537    pub fn subscribe_persisted_block(&self) -> watch::Receiver<Option<BlockNumHash>> {
538        self.inner.chain_info_tracker.subscribe_persisted_block()
539    }
540
541    /// Attempts to send a new [`CanonStateNotification`] to all active Receiver handles.
542    pub fn notify_canon_state(&self, event: CanonStateNotification<N>) {
543        self.inner.canon_state_notification_sender.send(event).ok();
544    }
545
546    /// Return state provider with reference to in-memory blocks that overlay database state.
547    ///
548    /// This merges the state of all blocks that are part of the chain that the requested block is
549    /// the head of. This includes all blocks that connect back to the canonical block on disk.
550    pub fn state_provider(
551        &self,
552        hash: B256,
553        historical: StateProviderBox,
554    ) -> MemoryOverlayStateProvider<N> {
555        let in_memory = if let Some(state) = self.state_by_hash(hash) {
556            state.chain().map(|block_state| block_state.block()).collect()
557        } else {
558            Vec::new()
559        };
560
561        MemoryOverlayStateProvider::new(historical, in_memory)
562    }
563
564    /// Returns an iterator over all __canonical blocks__ in the in-memory state, from newest to
565    /// oldest (highest to lowest).
566    ///
567    /// This iterator contains a snapshot of the in-memory state at the time of the call.
568    pub fn canonical_chain(&self) -> impl Iterator<Item = Arc<BlockState<N>>> {
569        self.inner.in_memory_state.head_state().into_iter().flat_map(|head| head.iter())
570    }
571
572    /// Returns [`SignedTransaction`] type for the given `TxHash` if found.
573    pub fn transaction_by_hash(&self, hash: TxHash) -> Option<N::SignedTx> {
574        for block_state in self.canonical_chain() {
575            if let Some(tx) =
576                block_state.block_ref().recovered_block().body().transaction_by_hash(&hash)
577            {
578                return Some(tx.clone())
579            }
580        }
581        None
582    }
583
584    /// Returns a tuple with [`SignedTransaction`] type and [`TransactionMeta`] for the
585    /// given [`TxHash`] if found.
586    pub fn transaction_by_hash_with_meta(
587        &self,
588        tx_hash: TxHash,
589    ) -> Option<(N::SignedTx, TransactionMeta)> {
590        for block_state in self.canonical_chain() {
591            if let Some(indexed) = block_state.find_indexed(tx_hash) {
592                return Some((indexed.tx().clone(), indexed.meta()));
593            }
594        }
595        None
596    }
597}
598
599/// State after applying the given block, this block is part of the canonical chain that partially
600/// stored in memory and can be traced back to a canonical block on disk.
601#[derive(Debug, Clone)]
602pub struct BlockState<N: NodePrimitives = EthPrimitives> {
603    /// The executed block that determines the state after this block has been executed.
604    block: ExecutedBlock<N>,
605    /// The block's parent block if it exists.
606    parent: Option<Arc<Self>>,
607}
608
609impl<N: NodePrimitives> PartialEq for BlockState<N> {
610    fn eq(&self, other: &Self) -> bool {
611        self.block == other.block && self.parent == other.parent
612    }
613}
614
615impl<N: NodePrimitives> BlockState<N> {
616    /// [`BlockState`] constructor.
617    pub const fn new(block: ExecutedBlock<N>) -> Self {
618        Self { block, parent: None }
619    }
620
621    /// [`BlockState`] constructor with parent.
622    pub const fn with_parent(block: ExecutedBlock<N>, parent: Option<Arc<Self>>) -> Self {
623        Self { block, parent }
624    }
625
626    /// Returns the hash and block of the on disk block this state can be traced back to.
627    pub fn anchor(&self) -> BlockNumHash {
628        let mut current = self;
629        while let Some(parent) = &current.parent {
630            current = parent;
631        }
632        current.block.recovered_block().parent_num_hash()
633    }
634
635    /// Returns the executed block that determines the state.
636    pub fn block(&self) -> ExecutedBlock<N> {
637        self.block.clone()
638    }
639
640    /// Returns a reference to the executed block that determines the state.
641    pub const fn block_ref(&self) -> &ExecutedBlock<N> {
642        &self.block
643    }
644
645    /// Returns the hash of executed block that determines the state.
646    pub fn hash(&self) -> B256 {
647        self.block.recovered_block().hash()
648    }
649
650    /// Returns the block number of executed block that determines the state.
651    pub fn number(&self) -> u64 {
652        self.block.recovered_block().number()
653    }
654
655    /// Returns the state root after applying the executed block that determines
656    /// the state.
657    pub fn state_root(&self) -> B256 {
658        self.block.recovered_block().state_root()
659    }
660
661    /// Returns the `Receipts` of executed block that determines the state.
662    pub fn receipts(&self) -> &Vec<N::Receipt> {
663        &self.block.execution_outcome().receipts
664    }
665
666    /// Returns a vector of `Receipt` of executed block that determines the state.
667    /// We assume that the `Receipts` in the executed block `ExecutionOutcome`
668    /// has only one element corresponding to the executed block associated to
669    /// the state.
670    ///
671    /// This clones the vector of receipts. To avoid it, use [`Self::executed_block_receipts_ref`].
672    pub fn executed_block_receipts(&self) -> Vec<N::Receipt> {
673        self.receipts().clone()
674    }
675
676    /// Returns a slice of `Receipt` of executed block that determines the state.
677    /// We assume that the `Receipts` in the executed block `ExecutionOutcome`
678    /// has only one element corresponding to the executed block associated to
679    /// the state.
680    pub fn executed_block_receipts_ref(&self) -> &[N::Receipt] {
681        self.receipts()
682    }
683
684    /// Returns an iterator over __parent__ `BlockStates`.
685    ///
686    /// The block state order is newest to oldest (highest to lowest):
687    /// `[5,4,3,2,1]`
688    ///
689    /// Note: This does not include self.
690    pub fn parent_state_chain(&self) -> impl Iterator<Item = &Self> + '_ {
691        std::iter::successors(self.parent.as_deref(), |state| state.parent.as_deref())
692    }
693
694    /// Returns a vector of `BlockStates` representing the entire in memory chain.
695    /// The block state order in the output vector is newest to oldest (highest to lowest),
696    /// including self as the first element.
697    pub fn chain(&self) -> impl Iterator<Item = &Self> {
698        std::iter::successors(Some(self), |state| state.parent.as_deref())
699    }
700
701    /// Appends the parent chain of this [`BlockState`] to the given vector.
702    ///
703    /// Parents are appended in order from newest to oldest (highest to lowest).
704    /// This does not include self, only the parent states.
705    ///
706    /// This is a convenience method equivalent to `chain.extend(self.parent_state_chain())`.
707    pub fn append_parent_chain<'a>(&'a self, chain: &mut Vec<&'a Self>) {
708        chain.extend(self.parent_state_chain());
709    }
710
711    /// Returns an iterator over the atomically captured chain of in memory blocks.
712    ///
713    /// This yields the blocks from newest to oldest (highest to lowest).
714    pub fn iter(self: Arc<Self>) -> impl Iterator<Item = Arc<Self>> {
715        std::iter::successors(Some(self), |state| state.parent.clone())
716    }
717
718    /// Return state provider with reference to in-memory blocks that overlay database state.
719    ///
720    /// This merges the state of all blocks that are part of the chain that the this block is
721    /// the head of. This includes all blocks that connect back to the canonical block on disk.
722    pub fn state_provider(&self, historical: StateProviderBox) -> MemoryOverlayStateProvider<N> {
723        let in_memory = self.chain().map(|block_state| block_state.block()).collect();
724
725        MemoryOverlayStateProvider::new(historical, in_memory)
726    }
727
728    /// Tries to find a block by [`BlockHashOrNumber`] in the chain ending at this block.
729    pub fn block_on_chain(&self, hash_or_num: BlockHashOrNumber) -> Option<&Self> {
730        self.chain().find(|block| match hash_or_num {
731            BlockHashOrNumber::Hash(hash) => block.hash() == hash,
732            BlockHashOrNumber::Number(number) => block.number() == number,
733        })
734    }
735
736    /// Tries to find a transaction by [`TxHash`] in the chain ending at this block.
737    pub fn transaction_on_chain(&self, hash: TxHash) -> Option<N::SignedTx> {
738        self.chain().find_map(|block_state| {
739            block_state.block_ref().recovered_block().body().transaction_by_hash(&hash).cloned()
740        })
741    }
742
743    /// Tries to find a transaction with meta by [`TxHash`] in the chain ending at this block.
744    pub fn transaction_meta_on_chain(
745        &self,
746        tx_hash: TxHash,
747    ) -> Option<(N::SignedTx, TransactionMeta)> {
748        self.chain().find_map(|block_state| {
749            block_state.find_indexed(tx_hash).map(|indexed| (indexed.tx().clone(), indexed.meta()))
750        })
751    }
752
753    /// Finds a transaction by hash and returns it with its index and block context.
754    pub fn find_indexed(&self, tx_hash: TxHash) -> Option<IndexedTx<'_, N::Block>> {
755        self.block_ref().recovered_block().find_indexed(tx_hash)
756    }
757}
758
759/// Represents an executed block stored in-memory.
760#[derive(Clone, Debug)]
761pub struct ExecutedBlock<N: NodePrimitives = EthPrimitives> {
762    /// Recovered Block
763    pub recovered_block: Arc<RecoveredBlock<N::Block>>,
764    /// Block's execution outcome.
765    pub execution_output: Arc<BlockExecutionOutput<N::Receipt>>,
766    /// Deferred trie data produced by execution.
767    ///
768    /// This allows deferring the computation of the trie data which can be expensive.
769    /// The data can be populated asynchronously after the block was validated.
770    pub trie_data: LazyTrieData,
771}
772
773impl<N: NodePrimitives> Default for ExecutedBlock<N> {
774    fn default() -> Self {
775        Self {
776            recovered_block: Default::default(),
777            execution_output: Arc::new(BlockExecutionOutput {
778                result: BlockExecutionResult {
779                    receipts: Default::default(),
780                    requests: Default::default(),
781                    gas_used: 0,
782                    blob_gas_used: 0,
783                },
784                state: Default::default(),
785            }),
786            trie_data: LazyTrieData::ready(ComputedTrieData::default()),
787        }
788    }
789}
790
791impl<N: NodePrimitives> PartialEq for ExecutedBlock<N> {
792    fn eq(&self, other: &Self) -> bool {
793        // Trie data is computed asynchronously and doesn't define block identity.
794        self.recovered_block == other.recovered_block &&
795            self.execution_output == other.execution_output
796    }
797}
798
799impl<N: NodePrimitives> ExecutedBlock<N> {
800    /// Create a new [`ExecutedBlock`] with already-computed trie data.
801    ///
802    /// Use this constructor when trie data is available immediately (e.g., sequencers,
803    /// payload builders). This is the safe default path.
804    pub fn new(
805        recovered_block: Arc<RecoveredBlock<N::Block>>,
806        execution_output: Arc<BlockExecutionOutput<N::Receipt>>,
807        trie_data: ComputedTrieData,
808    ) -> Self {
809        Self { recovered_block, execution_output, trie_data: LazyTrieData::ready(trie_data) }
810    }
811
812    /// Create a new [`ExecutedBlock`] with deferred trie data.
813    ///
814    /// This is useful if the trie data is populated somewhere else, e.g. asynchronously
815    /// after the block was validated.
816    ///
817    /// The [`LazyTrieData`] handle allows expensive trie operations (sorting hashed state and
818    /// trie updates) to be performed outside the critical validation path by a background task.
819    /// This can improve latency for time-sensitive operations like block validation.
820    ///
821    /// If the data hasn't been populated when [`Self::trie_data()`] is called, the caller waits
822    /// for the background task to publish it.
823    ///
824    /// Use [`Self::new()`] instead when trie data is already computed and available immediately.
825    pub const fn with_deferred_trie_data(
826        recovered_block: Arc<RecoveredBlock<N::Block>>,
827        execution_output: Arc<BlockExecutionOutput<N::Receipt>>,
828        trie_data: LazyTrieData,
829    ) -> Self {
830        Self { recovered_block, execution_output, trie_data }
831    }
832
833    /// Returns a reference to an inner [`SealedBlock`]
834    #[inline]
835    pub fn sealed_block(&self) -> &SealedBlock<N::Block> {
836        self.recovered_block.sealed_block()
837    }
838
839    /// Returns a reference to [`RecoveredBlock`]
840    #[inline]
841    pub fn recovered_block(&self) -> &RecoveredBlock<N::Block> {
842        &self.recovered_block
843    }
844
845    /// Returns a reference to the block's execution outcome
846    #[inline]
847    pub fn execution_outcome(&self) -> &BlockExecutionOutput<N::Receipt> {
848        &self.execution_output
849    }
850
851    /// Returns the trie data, waiting for the background task if not already cached.
852    ///
853    /// Uses `OnceLock::get_or_init` internally:
854    /// - If already computed: returns cached result immediately
855    /// - If not computed: first caller waits for the publishing task, others wait for that result
856    #[inline]
857    #[tracing::instrument(level = "debug", target = "engine::tree", name = "trie_data", skip_all)]
858    pub fn trie_data(&self) -> ComputedTrieData {
859        self.trie_data.get().clone()
860    }
861
862    /// Returns a clone of the deferred trie data handle.
863    ///
864    /// A handle is a lightweight reference that can be passed to descendants without
865    /// forcing trie data to be observed immediately. The actual work runs in the background task.
866    #[inline]
867    pub fn trie_data_handle(&self) -> LazyTrieData {
868        self.trie_data.clone()
869    }
870
871    /// Returns the hashed state result of the execution outcome.
872    ///
873    /// May wait for trie data if the deferred task hasn't completed.
874    #[inline]
875    pub fn hashed_state(&self) -> Arc<HashedPostStateSorted> {
876        self.trie_data().sorted.hashed_state
877    }
878
879    /// Returns a reference to the hashed state result of the execution outcome.
880    ///
881    /// May wait for trie data if the deferred task hasn't completed.
882    #[inline]
883    pub fn hashed_state_ref(&self) -> &HashedPostStateSorted {
884        &self.trie_data.get().sorted.hashed_state
885    }
886
887    /// Returns references to the hashed state results of the executed blocks.
888    ///
889    /// May wait for trie data if any deferred task hasn't completed.
890    pub fn hashed_state_refs(blocks: &[Self]) -> Vec<&HashedPostStateSorted> {
891        blocks.iter().map(Self::hashed_state_ref).collect()
892    }
893
894    /// Returns the trie updates resulting from the execution outcome.
895    ///
896    /// May wait for trie data if the deferred task hasn't completed.
897    #[inline]
898    pub fn trie_updates(&self) -> Arc<TrieUpdatesSorted> {
899        self.trie_data().sorted.trie_updates
900    }
901
902    /// Returns a reference to the trie updates resulting from the execution outcome.
903    ///
904    /// May wait for trie data if the deferred task hasn't completed.
905    #[inline]
906    pub fn trie_updates_ref(&self) -> &TrieUpdatesSorted {
907        &self.trie_data.get().sorted.trie_updates
908    }
909
910    /// Returns references to the trie updates of the executed blocks.
911    ///
912    /// May wait for trie data if any deferred task hasn't completed.
913    pub fn trie_updates_refs(blocks: &[Self]) -> Vec<&TrieUpdatesSorted> {
914        blocks.iter().map(Self::trie_updates_ref).collect()
915    }
916
917    /// Returns a [`BlockNumber`] of the block.
918    #[inline]
919    pub fn block_number(&self) -> BlockNumber {
920        self.recovered_block.header().number()
921    }
922}
923
924/// Non-empty chain of blocks.
925#[derive(Debug)]
926pub enum NewCanonicalChain<N: NodePrimitives = EthPrimitives> {
927    /// A simple append to the current canonical head
928    Commit {
929        /// all blocks that lead back to the canonical head
930        new: Vec<ExecutedBlock<N>>,
931    },
932    /// A reorged chain consists of two chains that trace back to a shared ancestor block at which
933    /// point they diverge.
934    Reorg {
935        /// All blocks of the _new_ chain
936        new: Vec<ExecutedBlock<N>>,
937        /// All blocks of the _old_ chain
938        old: Vec<ExecutedBlock<N>>,
939    },
940}
941
942impl<N: NodePrimitives<SignedTx: SignedTransaction>> NewCanonicalChain<N> {
943    /// Returns the length of the new chain.
944    pub const fn new_block_count(&self) -> usize {
945        match self {
946            Self::Commit { new } | Self::Reorg { new, .. } => new.len(),
947        }
948    }
949
950    /// Returns the length of the reorged chain.
951    pub const fn reorged_block_count(&self) -> usize {
952        match self {
953            Self::Commit { .. } => 0,
954            Self::Reorg { old, .. } => old.len(),
955        }
956    }
957
958    /// Converts the new chain into a notification that will be emitted to listeners
959    pub fn to_chain_notification(&self) -> CanonStateNotification<N> {
960        match self {
961            Self::Commit { new } => {
962                CanonStateNotification::Commit { new: Arc::new(Self::blocks_to_chain(new)) }
963            }
964            Self::Reorg { new, old } => CanonStateNotification::Reorg {
965                new: Arc::new(Self::blocks_to_chain(new)),
966                old: Arc::new(Self::blocks_to_chain(old)),
967            },
968        }
969    }
970
971    /// Converts a slice of executed blocks into a [`Chain`].
972    fn blocks_to_chain(blocks: &[ExecutedBlock<N>]) -> Chain<N> {
973        match blocks {
974            [] => Chain::default(),
975            [first, rest @ ..] => {
976                let mut chain = Chain::from_block(
977                    Arc::clone(&first.recovered_block),
978                    ExecutionOutcome::from((
979                        first.execution_outcome().clone(),
980                        first.block_number(),
981                    )),
982                    first.trie_data_handle(),
983                );
984                for exec in rest {
985                    chain.append_block(
986                        Arc::clone(&exec.recovered_block),
987                        ExecutionOutcome::from((
988                            exec.execution_outcome().clone(),
989                            exec.block_number(),
990                        )),
991                        exec.trie_data_handle(),
992                    );
993                }
994                chain
995            }
996        }
997    }
998
999    /// Returns the new tip of the chain.
1000    ///
1001    /// Returns the new tip for [`Self::Reorg`] and [`Self::Commit`] variants which commit at least
1002    /// 1 new block.
1003    pub fn tip(&self) -> &RecoveredBlock<N::Block> {
1004        match self {
1005            Self::Commit { new } | Self::Reorg { new, .. } => {
1006                new.last().expect("non empty blocks").recovered_block()
1007            }
1008        }
1009    }
1010}
1011
1012#[cfg(test)]
1013mod tests {
1014    use super::*;
1015    use crate::test_utils::TestBlockBuilder;
1016    use alloy_eips::eip7685::Requests;
1017    use alloy_primitives::{Address, BlockNumber, Bytes, StorageKey, StorageValue};
1018    use rand::Rng;
1019    use reth_errors::ProviderResult;
1020    use reth_ethereum_primitives::{EthPrimitives, Receipt};
1021    use reth_primitives_traits::{Account, Bytecode};
1022    use reth_storage_api::{
1023        AccountReader, BlockHashReader, BytecodeReader, HashedPostStateProvider,
1024        StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
1025    };
1026    use reth_trie::{
1027        updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
1028        MultiProofTargets, StorageMultiProof, StorageProof, TrieInput,
1029    };
1030
1031    fn create_mock_state(
1032        test_block_builder: &mut TestBlockBuilder<EthPrimitives>,
1033        block_number: u64,
1034        parent_hash: B256,
1035    ) -> BlockState {
1036        BlockState::new(
1037            test_block_builder.get_executed_block_with_number(block_number, parent_hash),
1038        )
1039    }
1040
1041    fn create_mock_state_chain(
1042        test_block_builder: &mut TestBlockBuilder<EthPrimitives>,
1043        num_blocks: u64,
1044    ) -> Vec<BlockState> {
1045        let mut chain = Vec::with_capacity(num_blocks as usize);
1046        let mut parent_hash = B256::random();
1047        let mut parent_state: Option<BlockState> = None;
1048
1049        for i in 1..=num_blocks {
1050            let mut state = create_mock_state(test_block_builder, i, parent_hash);
1051            if let Some(parent) = parent_state {
1052                state.parent = Some(Arc::new(parent));
1053            }
1054            parent_hash = state.hash();
1055            parent_state = Some(state.clone());
1056            chain.push(state);
1057        }
1058
1059        chain
1060    }
1061
1062    struct MockStateProvider;
1063
1064    impl StateProvider for MockStateProvider {
1065        fn storage(
1066            &self,
1067            _address: Address,
1068            _storage_key: StorageKey,
1069        ) -> ProviderResult<Option<StorageValue>> {
1070            Ok(None)
1071        }
1072    }
1073
1074    impl BytecodeReader for MockStateProvider {
1075        fn bytecode_by_hash(&self, _code_hash: &B256) -> ProviderResult<Option<Bytecode>> {
1076            Ok(None)
1077        }
1078    }
1079
1080    impl BlockHashReader for MockStateProvider {
1081        fn block_hash(&self, _number: BlockNumber) -> ProviderResult<Option<B256>> {
1082            Ok(None)
1083        }
1084
1085        fn canonical_hashes_range(
1086            &self,
1087            _start: BlockNumber,
1088            _end: BlockNumber,
1089        ) -> ProviderResult<Vec<B256>> {
1090            Ok(vec![])
1091        }
1092    }
1093
1094    impl AccountReader for MockStateProvider {
1095        fn basic_account(&self, _address: &Address) -> ProviderResult<Option<Account>> {
1096            Ok(None)
1097        }
1098    }
1099
1100    impl StateRootProvider for MockStateProvider {
1101        fn state_root(&self, _hashed_state: HashedPostState) -> ProviderResult<B256> {
1102            Ok(B256::random())
1103        }
1104
1105        fn state_root_from_nodes(&self, _input: TrieInput) -> ProviderResult<B256> {
1106            Ok(B256::random())
1107        }
1108
1109        fn state_root_with_updates(
1110            &self,
1111            _hashed_state: HashedPostState,
1112        ) -> ProviderResult<(B256, TrieUpdates)> {
1113            Ok((B256::random(), TrieUpdates::default()))
1114        }
1115
1116        fn state_root_from_nodes_with_updates(
1117            &self,
1118            _input: TrieInput,
1119        ) -> ProviderResult<(B256, TrieUpdates)> {
1120            Ok((B256::random(), TrieUpdates::default()))
1121        }
1122    }
1123
1124    impl HashedPostStateProvider for MockStateProvider {
1125        fn hashed_post_state(
1126            &self,
1127            _bundle_state: &revm::database::BundleState,
1128        ) -> ProviderResult<HashedPostState> {
1129            Ok(HashedPostState::default())
1130        }
1131    }
1132
1133    impl StorageRootProvider for MockStateProvider {
1134        fn storage_root(
1135            &self,
1136            _address: Address,
1137            _hashed_storage: HashedStorage,
1138        ) -> ProviderResult<B256> {
1139            Ok(B256::random())
1140        }
1141
1142        fn storage_proof(
1143            &self,
1144            _address: Address,
1145            slot: B256,
1146            _hashed_storage: HashedStorage,
1147        ) -> ProviderResult<StorageProof> {
1148            Ok(StorageProof::new(slot))
1149        }
1150
1151        fn storage_multiproof(
1152            &self,
1153            _address: Address,
1154            _slots: &[B256],
1155            _hashed_storage: HashedStorage,
1156        ) -> ProviderResult<StorageMultiProof> {
1157            Ok(StorageMultiProof::empty())
1158        }
1159    }
1160
1161    impl StateProofProvider for MockStateProvider {
1162        fn proof(
1163            &self,
1164            _input: TrieInput,
1165            _address: Address,
1166            _slots: &[B256],
1167        ) -> ProviderResult<AccountProof> {
1168            Ok(AccountProof::new(Address::random()))
1169        }
1170
1171        fn multiproof(
1172            &self,
1173            _input: TrieInput,
1174            _targets: MultiProofTargets,
1175        ) -> ProviderResult<MultiProof> {
1176            Ok(MultiProof::default())
1177        }
1178
1179        fn witness(
1180            &self,
1181            _input: TrieInput,
1182            _target: HashedPostState,
1183            _mode: reth_trie::ExecutionWitnessMode,
1184        ) -> ProviderResult<Vec<Bytes>> {
1185            Ok(Vec::default())
1186        }
1187    }
1188
1189    #[test]
1190    fn test_in_memory_state_impl_state_by_hash() {
1191        let mut state_by_hash = B256Map::default();
1192        let number = rand::rng().random::<u64>();
1193        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1194        let state = Arc::new(create_mock_state(&mut test_block_builder, number, B256::random()));
1195        state_by_hash.insert(state.hash(), state.clone());
1196
1197        let in_memory_state = InMemoryState::new(state_by_hash, BTreeMap::new(), None);
1198
1199        assert_eq!(in_memory_state.state_by_hash(state.hash()), Some(state));
1200        assert_eq!(in_memory_state.state_by_hash(B256::random()), None);
1201    }
1202
1203    #[test]
1204    fn test_in_memory_state_impl_state_by_number() {
1205        let mut state_by_hash = B256Map::default();
1206        let mut hash_by_number = BTreeMap::new();
1207
1208        let number = rand::rng().random::<u64>();
1209        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1210        let state = Arc::new(create_mock_state(&mut test_block_builder, number, B256::random()));
1211        let hash = state.hash();
1212
1213        state_by_hash.insert(hash, state.clone());
1214        hash_by_number.insert(number, hash);
1215
1216        let in_memory_state = InMemoryState::new(state_by_hash, hash_by_number, None);
1217
1218        assert_eq!(in_memory_state.state_by_number(number), Some(state));
1219        assert_eq!(in_memory_state.state_by_number(number + 1), None);
1220    }
1221
1222    #[test]
1223    fn test_in_memory_state_impl_head_state() {
1224        let mut state_by_hash = B256Map::default();
1225        let mut hash_by_number = BTreeMap::new();
1226        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1227        let state1 = Arc::new(create_mock_state(&mut test_block_builder, 1, B256::random()));
1228        let hash1 = state1.hash();
1229        let state2 = Arc::new(create_mock_state(&mut test_block_builder, 2, hash1));
1230        let hash2 = state2.hash();
1231        hash_by_number.insert(1, hash1);
1232        hash_by_number.insert(2, hash2);
1233        state_by_hash.insert(hash1, state1);
1234        state_by_hash.insert(hash2, state2);
1235
1236        let in_memory_state = InMemoryState::new(state_by_hash, hash_by_number, None);
1237        let head_state = in_memory_state.head_state().unwrap();
1238
1239        assert_eq!(head_state.hash(), hash2);
1240        assert_eq!(head_state.number(), 2);
1241    }
1242
1243    #[test]
1244    fn test_in_memory_state_impl_pending_state() {
1245        let pending_number = rand::rng().random::<u64>();
1246        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1247        let pending_state =
1248            create_mock_state(&mut test_block_builder, pending_number, B256::random());
1249        let pending_hash = pending_state.hash();
1250
1251        let in_memory_state =
1252            InMemoryState::new(B256Map::default(), BTreeMap::new(), Some(pending_state));
1253
1254        let result = in_memory_state.pending_state();
1255        assert!(result.is_some());
1256        let actual_pending_state = result.unwrap();
1257        assert_eq!(actual_pending_state.block.recovered_block().hash(), pending_hash);
1258        assert_eq!(actual_pending_state.block.recovered_block().number, pending_number);
1259    }
1260
1261    #[test]
1262    fn test_in_memory_state_impl_no_pending_state() {
1263        let in_memory_state: InMemoryState =
1264            InMemoryState::new(B256Map::default(), BTreeMap::new(), None);
1265
1266        assert_eq!(in_memory_state.pending_state(), None);
1267    }
1268
1269    #[test]
1270    fn test_state() {
1271        let number = rand::rng().random::<u64>();
1272        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1273        let block = test_block_builder.get_executed_block_with_number(number, B256::random());
1274
1275        let state = BlockState::new(block.clone());
1276
1277        assert_eq!(state.block(), block);
1278        assert_eq!(state.hash(), block.recovered_block().hash());
1279        assert_eq!(state.number(), number);
1280        assert_eq!(state.state_root(), block.recovered_block().state_root);
1281    }
1282
1283    #[test]
1284    fn test_state_receipts() {
1285        let receipts = vec![vec![Receipt::default()]];
1286        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1287        let block =
1288            test_block_builder.get_executed_block_with_receipts(receipts.clone(), B256::random());
1289
1290        let state = BlockState::new(block);
1291
1292        assert_eq!(state.receipts(), receipts.first().unwrap());
1293    }
1294
1295    #[test]
1296    fn test_in_memory_state_chain_update() {
1297        let state: CanonicalInMemoryState = CanonicalInMemoryState::empty();
1298        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1299        let block1 = test_block_builder.get_executed_block_with_number(0, B256::random());
1300        let block2 = test_block_builder.get_executed_block_with_number(0, B256::random());
1301        let chain = NewCanonicalChain::Commit { new: vec![block1.clone()] };
1302        state.update_chain(chain);
1303        assert_eq!(
1304            state.head_state().unwrap().block_ref().recovered_block().hash(),
1305            block1.recovered_block().hash()
1306        );
1307        assert_eq!(
1308            state.state_by_number(0).unwrap().block_ref().recovered_block().hash(),
1309            block1.recovered_block().hash()
1310        );
1311
1312        let chain = NewCanonicalChain::Reorg { new: vec![block2.clone()], old: vec![block1] };
1313        state.update_chain(chain);
1314        assert_eq!(
1315            state.head_state().unwrap().block_ref().recovered_block().hash(),
1316            block2.recovered_block().hash()
1317        );
1318        assert_eq!(
1319            state.state_by_number(0).unwrap().block_ref().recovered_block().hash(),
1320            block2.recovered_block().hash()
1321        );
1322
1323        assert_eq!(state.inner.in_memory_state.block_count(), 1);
1324    }
1325
1326    #[test]
1327    fn test_in_memory_state_set_pending_block() {
1328        let state: CanonicalInMemoryState = CanonicalInMemoryState::empty();
1329        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1330
1331        // First random block
1332        let block1 = test_block_builder.get_executed_block_with_number(0, B256::random());
1333
1334        // Second block with parent hash of the first block
1335        let block2 =
1336            test_block_builder.get_executed_block_with_number(1, block1.recovered_block().hash());
1337
1338        // Commit the two blocks
1339        let chain = NewCanonicalChain::Commit { new: vec![block1.clone(), block2.clone()] };
1340        state.update_chain(chain);
1341
1342        // Assert that the pending state is None before setting it
1343        assert!(state.pending_state().is_none());
1344
1345        // Set the pending block
1346        state.set_pending_block(block2.clone());
1347
1348        // Check the pending state
1349        assert_eq!(
1350            state.pending_state().unwrap(),
1351            BlockState::with_parent(block2.clone(), Some(Arc::new(BlockState::new(block1))))
1352        );
1353
1354        // Check the pending block
1355        assert_eq!(state.pending_block().unwrap(), block2.recovered_block().sealed_block().clone());
1356
1357        // Check the pending block number and hash
1358        assert_eq!(
1359            state.pending_block_num_hash().unwrap(),
1360            BlockNumHash { number: 1, hash: block2.recovered_block().hash() }
1361        );
1362
1363        // Check the pending header
1364        assert_eq!(state.pending_header().unwrap(), block2.recovered_block().header().clone());
1365
1366        // Check the pending sealed header
1367        assert_eq!(
1368            state.pending_sealed_header().unwrap(),
1369            block2.recovered_block().clone_sealed_header()
1370        );
1371
1372        // Check the pending block with senders
1373        assert_eq!(state.pending_recovered_block().unwrap(), block2.recovered_block().clone());
1374
1375        // Check the pending block and receipts
1376        assert_eq!(
1377            state.pending_block_and_receipts().unwrap(),
1378            (block2.recovered_block().clone(), vec![])
1379        );
1380    }
1381
1382    #[test]
1383    fn test_canonical_in_memory_state_state_provider() {
1384        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1385        let block1 = test_block_builder.get_executed_block_with_number(1, B256::random());
1386        let block2 =
1387            test_block_builder.get_executed_block_with_number(2, block1.recovered_block().hash());
1388        let block3 =
1389            test_block_builder.get_executed_block_with_number(3, block2.recovered_block().hash());
1390
1391        let state1 = Arc::new(BlockState::new(block1.clone()));
1392        let state2 = Arc::new(BlockState::with_parent(block2.clone(), Some(state1.clone())));
1393        let state3 = Arc::new(BlockState::with_parent(block3.clone(), Some(state2.clone())));
1394
1395        let mut blocks = B256Map::default();
1396        blocks.insert(block1.recovered_block().hash(), state1);
1397        blocks.insert(block2.recovered_block().hash(), state2);
1398        blocks.insert(block3.recovered_block().hash(), state3);
1399
1400        let mut numbers = BTreeMap::new();
1401        numbers.insert(1, block1.recovered_block().hash());
1402        numbers.insert(2, block2.recovered_block().hash());
1403        numbers.insert(3, block3.recovered_block().hash());
1404
1405        let canonical_state = CanonicalInMemoryState::new(blocks, numbers, None, None, None);
1406
1407        let historical: StateProviderBox = Box::new(MockStateProvider);
1408
1409        let overlay_provider =
1410            canonical_state.state_provider(block3.recovered_block().hash(), historical);
1411
1412        assert_eq!(overlay_provider.in_memory.len(), 3);
1413        assert_eq!(overlay_provider.in_memory[0].recovered_block().number, 3);
1414        assert_eq!(overlay_provider.in_memory[1].recovered_block().number, 2);
1415        assert_eq!(overlay_provider.in_memory[2].recovered_block().number, 1);
1416
1417        assert_eq!(
1418            overlay_provider.in_memory[0].recovered_block().parent_hash,
1419            overlay_provider.in_memory[1].recovered_block().hash()
1420        );
1421        assert_eq!(
1422            overlay_provider.in_memory[1].recovered_block().parent_hash,
1423            overlay_provider.in_memory[2].recovered_block().hash()
1424        );
1425
1426        let unknown_hash = B256::random();
1427        let empty_overlay_provider =
1428            canonical_state.state_provider(unknown_hash, Box::new(MockStateProvider));
1429        assert_eq!(empty_overlay_provider.in_memory.len(), 0);
1430    }
1431
1432    #[test]
1433    fn test_canonical_in_memory_state_canonical_chain_empty() {
1434        let state: CanonicalInMemoryState = CanonicalInMemoryState::empty();
1435        assert!(state.canonical_chain().next().is_none());
1436    }
1437
1438    #[test]
1439    fn test_canonical_in_memory_state_canonical_chain_single_block() {
1440        let block = TestBlockBuilder::eth().get_executed_block_with_number(1, B256::random());
1441        let hash = block.recovered_block().hash();
1442        let mut blocks = B256Map::default();
1443        blocks.insert(hash, Arc::new(BlockState::new(block)));
1444        let mut numbers = BTreeMap::new();
1445        numbers.insert(1, hash);
1446
1447        let state = CanonicalInMemoryState::new(blocks, numbers, None, None, None);
1448        let chain: Vec<_> = state.canonical_chain().collect();
1449
1450        assert_eq!(chain.len(), 1);
1451        assert_eq!(chain[0].number(), 1);
1452        assert_eq!(chain[0].hash(), hash);
1453    }
1454
1455    #[test]
1456    fn test_canonical_in_memory_state_canonical_chain_multiple_blocks() {
1457        let mut parent_hash = B256::random();
1458        let mut block_builder = TestBlockBuilder::eth();
1459        let state: CanonicalInMemoryState = CanonicalInMemoryState::empty();
1460
1461        for i in 1..=3 {
1462            let block = block_builder.get_executed_block_with_number(i, parent_hash);
1463            let hash = block.recovered_block().hash();
1464            state.update_blocks(Some(block), None);
1465            parent_hash = hash;
1466        }
1467
1468        let chain: Vec<_> = state.canonical_chain().collect();
1469
1470        assert_eq!(chain.len(), 3);
1471        assert_eq!(chain[0].number(), 3);
1472        assert_eq!(chain[1].number(), 2);
1473        assert_eq!(chain[2].number(), 1);
1474    }
1475
1476    // ensures the pending block is not part of the canonical chain
1477    #[test]
1478    fn test_canonical_in_memory_state_canonical_chain_with_pending_block() {
1479        let mut parent_hash = B256::random();
1480        let mut block_builder = TestBlockBuilder::<EthPrimitives>::eth();
1481        let state: CanonicalInMemoryState = CanonicalInMemoryState::empty();
1482
1483        for i in 1..=2 {
1484            let block = block_builder.get_executed_block_with_number(i, parent_hash);
1485            let hash = block.recovered_block().hash();
1486            state.update_blocks(Some(block), None);
1487            parent_hash = hash;
1488        }
1489
1490        let pending_block = block_builder.get_executed_block_with_number(3, parent_hash);
1491        state.set_pending_block(pending_block);
1492        let chain: Vec<_> = state.canonical_chain().collect();
1493
1494        assert_eq!(chain.len(), 2);
1495        assert_eq!(chain[0].number(), 2);
1496        assert_eq!(chain[1].number(), 1);
1497    }
1498
1499    #[test]
1500    fn test_block_state_parent_blocks() {
1501        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1502        let chain = create_mock_state_chain(&mut test_block_builder, 4);
1503
1504        let parents: Vec<_> = chain[3].parent_state_chain().collect();
1505        assert_eq!(parents.len(), 3);
1506        assert_eq!(parents[0].block().recovered_block().number, 3);
1507        assert_eq!(parents[1].block().recovered_block().number, 2);
1508        assert_eq!(parents[2].block().recovered_block().number, 1);
1509
1510        let parents: Vec<_> = chain[2].parent_state_chain().collect();
1511        assert_eq!(parents.len(), 2);
1512        assert_eq!(parents[0].block().recovered_block().number, 2);
1513        assert_eq!(parents[1].block().recovered_block().number, 1);
1514
1515        assert_eq!(chain[0].parent_state_chain().count(), 0);
1516    }
1517
1518    #[test]
1519    fn test_block_state_single_block_state_chain() {
1520        let single_block_number = 1;
1521        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1522        let single_block =
1523            create_mock_state(&mut test_block_builder, single_block_number, B256::random());
1524        let single_block_hash = single_block.block().recovered_block().hash();
1525
1526        assert_eq!(single_block.parent_state_chain().count(), 0);
1527
1528        let block_state_chain = single_block.chain().collect::<Vec<_>>();
1529        assert_eq!(block_state_chain.len(), 1);
1530        assert_eq!(block_state_chain[0].block().recovered_block().number, single_block_number);
1531        assert_eq!(block_state_chain[0].block().recovered_block().hash(), single_block_hash);
1532    }
1533
1534    #[test]
1535    fn test_block_state_chain() {
1536        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1537        let chain = create_mock_state_chain(&mut test_block_builder, 3);
1538
1539        let block_state_chain = chain[2].chain().collect::<Vec<_>>();
1540        assert_eq!(block_state_chain.len(), 3);
1541        assert_eq!(block_state_chain[0].block().recovered_block().number, 3);
1542        assert_eq!(block_state_chain[1].block().recovered_block().number, 2);
1543        assert_eq!(block_state_chain[2].block().recovered_block().number, 1);
1544
1545        let block_state_chain = chain[1].chain().collect::<Vec<_>>();
1546        assert_eq!(block_state_chain.len(), 2);
1547        assert_eq!(block_state_chain[0].block().recovered_block().number, 2);
1548        assert_eq!(block_state_chain[1].block().recovered_block().number, 1);
1549
1550        let block_state_chain = chain[0].chain().collect::<Vec<_>>();
1551        assert_eq!(block_state_chain.len(), 1);
1552        assert_eq!(block_state_chain[0].block().recovered_block().number, 1);
1553    }
1554
1555    #[test]
1556    fn test_to_chain_notification() {
1557        // Generate 4 blocks
1558        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1559        let block0 = test_block_builder.get_executed_block_with_number(0, B256::random());
1560        let block1 =
1561            test_block_builder.get_executed_block_with_number(1, block0.recovered_block.hash());
1562        let block1a =
1563            test_block_builder.get_executed_block_with_number(1, block0.recovered_block.hash());
1564        let block2 =
1565            test_block_builder.get_executed_block_with_number(2, block1.recovered_block.hash());
1566        let block2a =
1567            test_block_builder.get_executed_block_with_number(2, block1.recovered_block.hash());
1568
1569        // Test commit notification
1570        let chain_commit = NewCanonicalChain::Commit { new: vec![block0.clone(), block1.clone()] };
1571
1572        // Build expected trie data map
1573        let mut expected_trie_data = BTreeMap::new();
1574        expected_trie_data.insert(0, LazyTrieData::ready(block0.trie_data()));
1575        expected_trie_data.insert(1, LazyTrieData::ready(block1.trie_data()));
1576
1577        // Build expected execution outcome (first_block matches first block number)
1578        let commit_execution_outcome = ExecutionOutcome {
1579            receipts: vec![vec![], vec![]],
1580            requests: vec![Requests::default(), Requests::default()],
1581            first_block: 0,
1582            ..Default::default()
1583        };
1584
1585        assert_eq!(
1586            chain_commit.to_chain_notification(),
1587            CanonStateNotification::Commit {
1588                new: Arc::new(Chain::new(
1589                    vec![block0.recovered_block().clone(), block1.recovered_block().clone()],
1590                    commit_execution_outcome,
1591                    expected_trie_data,
1592                ))
1593            }
1594        );
1595
1596        // Test reorg notification
1597        let chain_reorg = NewCanonicalChain::Reorg {
1598            new: vec![block1a.clone(), block2a.clone()],
1599            old: vec![block1.clone(), block2.clone()],
1600        };
1601
1602        // Build expected trie data for old chain
1603        let mut old_trie_data = BTreeMap::new();
1604        old_trie_data.insert(1, LazyTrieData::ready(block1.trie_data()));
1605        old_trie_data.insert(2, LazyTrieData::ready(block2.trie_data()));
1606
1607        // Build expected trie data for new chain
1608        let mut new_trie_data = BTreeMap::new();
1609        new_trie_data.insert(1, LazyTrieData::ready(block1a.trie_data()));
1610        new_trie_data.insert(2, LazyTrieData::ready(block2a.trie_data()));
1611
1612        // Build expected execution outcome for reorg chains (first_block matches first block
1613        // number)
1614        let reorg_execution_outcome = ExecutionOutcome {
1615            receipts: vec![vec![], vec![]],
1616            requests: vec![Requests::default(), Requests::default()],
1617            first_block: 1,
1618            ..Default::default()
1619        };
1620
1621        assert_eq!(
1622            chain_reorg.to_chain_notification(),
1623            CanonStateNotification::Reorg {
1624                old: Arc::new(Chain::new(
1625                    vec![block1.recovered_block().clone(), block2.recovered_block().clone()],
1626                    reorg_execution_outcome.clone(),
1627                    old_trie_data,
1628                )),
1629                new: Arc::new(Chain::new(
1630                    vec![block1a.recovered_block().clone(), block2a.recovered_block().clone()],
1631                    reorg_execution_outcome,
1632                    new_trie_data,
1633                ))
1634            }
1635        );
1636    }
1637}