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 the trie updates resulting from the execution outcome.
880    ///
881    /// May wait for trie data if the deferred task hasn't completed.
882    #[inline]
883    pub fn trie_updates(&self) -> Arc<TrieUpdatesSorted> {
884        self.trie_data().sorted.trie_updates
885    }
886
887    /// Returns a [`BlockNumber`] of the block.
888    #[inline]
889    pub fn block_number(&self) -> BlockNumber {
890        self.recovered_block.header().number()
891    }
892}
893
894/// Non-empty chain of blocks.
895#[derive(Debug)]
896pub enum NewCanonicalChain<N: NodePrimitives = EthPrimitives> {
897    /// A simple append to the current canonical head
898    Commit {
899        /// all blocks that lead back to the canonical head
900        new: Vec<ExecutedBlock<N>>,
901    },
902    /// A reorged chain consists of two chains that trace back to a shared ancestor block at which
903    /// point they diverge.
904    Reorg {
905        /// All blocks of the _new_ chain
906        new: Vec<ExecutedBlock<N>>,
907        /// All blocks of the _old_ chain
908        old: Vec<ExecutedBlock<N>>,
909    },
910}
911
912impl<N: NodePrimitives<SignedTx: SignedTransaction>> NewCanonicalChain<N> {
913    /// Returns the length of the new chain.
914    pub const fn new_block_count(&self) -> usize {
915        match self {
916            Self::Commit { new } | Self::Reorg { new, .. } => new.len(),
917        }
918    }
919
920    /// Returns the length of the reorged chain.
921    pub const fn reorged_block_count(&self) -> usize {
922        match self {
923            Self::Commit { .. } => 0,
924            Self::Reorg { old, .. } => old.len(),
925        }
926    }
927
928    /// Converts the new chain into a notification that will be emitted to listeners
929    pub fn to_chain_notification(&self) -> CanonStateNotification<N> {
930        match self {
931            Self::Commit { new } => {
932                CanonStateNotification::Commit { new: Arc::new(Self::blocks_to_chain(new)) }
933            }
934            Self::Reorg { new, old } => CanonStateNotification::Reorg {
935                new: Arc::new(Self::blocks_to_chain(new)),
936                old: Arc::new(Self::blocks_to_chain(old)),
937            },
938        }
939    }
940
941    /// Converts a slice of executed blocks into a [`Chain`].
942    fn blocks_to_chain(blocks: &[ExecutedBlock<N>]) -> Chain<N> {
943        match blocks {
944            [] => Chain::default(),
945            [first, rest @ ..] => {
946                let mut chain = Chain::from_block(
947                    Arc::clone(&first.recovered_block),
948                    ExecutionOutcome::from((
949                        first.execution_outcome().clone(),
950                        first.block_number(),
951                    )),
952                    first.trie_data_handle(),
953                );
954                for exec in rest {
955                    chain.append_block(
956                        Arc::clone(&exec.recovered_block),
957                        ExecutionOutcome::from((
958                            exec.execution_outcome().clone(),
959                            exec.block_number(),
960                        )),
961                        exec.trie_data_handle(),
962                    );
963                }
964                chain
965            }
966        }
967    }
968
969    /// Returns the new tip of the chain.
970    ///
971    /// Returns the new tip for [`Self::Reorg`] and [`Self::Commit`] variants which commit at least
972    /// 1 new block.
973    pub fn tip(&self) -> &RecoveredBlock<N::Block> {
974        match self {
975            Self::Commit { new } | Self::Reorg { new, .. } => {
976                new.last().expect("non empty blocks").recovered_block()
977            }
978        }
979    }
980}
981
982#[cfg(test)]
983mod tests {
984    use super::*;
985    use crate::test_utils::TestBlockBuilder;
986    use alloy_eips::eip7685::Requests;
987    use alloy_primitives::{Address, BlockNumber, Bytes, StorageKey, StorageValue};
988    use rand::Rng;
989    use reth_errors::ProviderResult;
990    use reth_ethereum_primitives::{EthPrimitives, Receipt};
991    use reth_primitives_traits::{Account, Bytecode};
992    use reth_storage_api::{
993        AccountReader, BlockHashReader, BytecodeReader, HashedPostStateProvider,
994        StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
995    };
996    use reth_trie::{
997        updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
998        MultiProofTargets, StorageMultiProof, StorageProof, TrieInput,
999    };
1000
1001    fn create_mock_state(
1002        test_block_builder: &mut TestBlockBuilder<EthPrimitives>,
1003        block_number: u64,
1004        parent_hash: B256,
1005    ) -> BlockState {
1006        BlockState::new(
1007            test_block_builder.get_executed_block_with_number(block_number, parent_hash),
1008        )
1009    }
1010
1011    fn create_mock_state_chain(
1012        test_block_builder: &mut TestBlockBuilder<EthPrimitives>,
1013        num_blocks: u64,
1014    ) -> Vec<BlockState> {
1015        let mut chain = Vec::with_capacity(num_blocks as usize);
1016        let mut parent_hash = B256::random();
1017        let mut parent_state: Option<BlockState> = None;
1018
1019        for i in 1..=num_blocks {
1020            let mut state = create_mock_state(test_block_builder, i, parent_hash);
1021            if let Some(parent) = parent_state {
1022                state.parent = Some(Arc::new(parent));
1023            }
1024            parent_hash = state.hash();
1025            parent_state = Some(state.clone());
1026            chain.push(state);
1027        }
1028
1029        chain
1030    }
1031
1032    struct MockStateProvider;
1033
1034    impl StateProvider for MockStateProvider {
1035        fn storage(
1036            &self,
1037            _address: Address,
1038            _storage_key: StorageKey,
1039        ) -> ProviderResult<Option<StorageValue>> {
1040            Ok(None)
1041        }
1042    }
1043
1044    impl BytecodeReader for MockStateProvider {
1045        fn bytecode_by_hash(&self, _code_hash: &B256) -> ProviderResult<Option<Bytecode>> {
1046            Ok(None)
1047        }
1048    }
1049
1050    impl BlockHashReader for MockStateProvider {
1051        fn block_hash(&self, _number: BlockNumber) -> ProviderResult<Option<B256>> {
1052            Ok(None)
1053        }
1054
1055        fn canonical_hashes_range(
1056            &self,
1057            _start: BlockNumber,
1058            _end: BlockNumber,
1059        ) -> ProviderResult<Vec<B256>> {
1060            Ok(vec![])
1061        }
1062    }
1063
1064    impl AccountReader for MockStateProvider {
1065        fn basic_account(&self, _address: &Address) -> ProviderResult<Option<Account>> {
1066            Ok(None)
1067        }
1068    }
1069
1070    impl StateRootProvider for MockStateProvider {
1071        fn state_root(&self, _hashed_state: HashedPostState) -> ProviderResult<B256> {
1072            Ok(B256::random())
1073        }
1074
1075        fn state_root_from_nodes(&self, _input: TrieInput) -> ProviderResult<B256> {
1076            Ok(B256::random())
1077        }
1078
1079        fn state_root_with_updates(
1080            &self,
1081            _hashed_state: HashedPostState,
1082        ) -> ProviderResult<(B256, TrieUpdates)> {
1083            Ok((B256::random(), TrieUpdates::default()))
1084        }
1085
1086        fn state_root_from_nodes_with_updates(
1087            &self,
1088            _input: TrieInput,
1089        ) -> ProviderResult<(B256, TrieUpdates)> {
1090            Ok((B256::random(), TrieUpdates::default()))
1091        }
1092    }
1093
1094    impl HashedPostStateProvider for MockStateProvider {
1095        fn hashed_post_state(
1096            &self,
1097            _bundle_state: &revm::database::BundleState,
1098        ) -> HashedPostState {
1099            HashedPostState::default()
1100        }
1101    }
1102
1103    impl StorageRootProvider for MockStateProvider {
1104        fn storage_root(
1105            &self,
1106            _address: Address,
1107            _hashed_storage: HashedStorage,
1108        ) -> ProviderResult<B256> {
1109            Ok(B256::random())
1110        }
1111
1112        fn storage_proof(
1113            &self,
1114            _address: Address,
1115            slot: B256,
1116            _hashed_storage: HashedStorage,
1117        ) -> ProviderResult<StorageProof> {
1118            Ok(StorageProof::new(slot))
1119        }
1120
1121        fn storage_multiproof(
1122            &self,
1123            _address: Address,
1124            _slots: &[B256],
1125            _hashed_storage: HashedStorage,
1126        ) -> ProviderResult<StorageMultiProof> {
1127            Ok(StorageMultiProof::empty())
1128        }
1129    }
1130
1131    impl StateProofProvider for MockStateProvider {
1132        fn proof(
1133            &self,
1134            _input: TrieInput,
1135            _address: Address,
1136            _slots: &[B256],
1137        ) -> ProviderResult<AccountProof> {
1138            Ok(AccountProof::new(Address::random()))
1139        }
1140
1141        fn multiproof(
1142            &self,
1143            _input: TrieInput,
1144            _targets: MultiProofTargets,
1145        ) -> ProviderResult<MultiProof> {
1146            Ok(MultiProof::default())
1147        }
1148
1149        fn witness(
1150            &self,
1151            _input: TrieInput,
1152            _target: HashedPostState,
1153            _mode: reth_trie::ExecutionWitnessMode,
1154        ) -> ProviderResult<Vec<Bytes>> {
1155            Ok(Vec::default())
1156        }
1157    }
1158
1159    #[test]
1160    fn test_in_memory_state_impl_state_by_hash() {
1161        let mut state_by_hash = B256Map::default();
1162        let number = rand::rng().random::<u64>();
1163        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1164        let state = Arc::new(create_mock_state(&mut test_block_builder, number, B256::random()));
1165        state_by_hash.insert(state.hash(), state.clone());
1166
1167        let in_memory_state = InMemoryState::new(state_by_hash, BTreeMap::new(), None);
1168
1169        assert_eq!(in_memory_state.state_by_hash(state.hash()), Some(state));
1170        assert_eq!(in_memory_state.state_by_hash(B256::random()), None);
1171    }
1172
1173    #[test]
1174    fn test_in_memory_state_impl_state_by_number() {
1175        let mut state_by_hash = B256Map::default();
1176        let mut hash_by_number = BTreeMap::new();
1177
1178        let number = rand::rng().random::<u64>();
1179        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1180        let state = Arc::new(create_mock_state(&mut test_block_builder, number, B256::random()));
1181        let hash = state.hash();
1182
1183        state_by_hash.insert(hash, state.clone());
1184        hash_by_number.insert(number, hash);
1185
1186        let in_memory_state = InMemoryState::new(state_by_hash, hash_by_number, None);
1187
1188        assert_eq!(in_memory_state.state_by_number(number), Some(state));
1189        assert_eq!(in_memory_state.state_by_number(number + 1), None);
1190    }
1191
1192    #[test]
1193    fn test_in_memory_state_impl_head_state() {
1194        let mut state_by_hash = B256Map::default();
1195        let mut hash_by_number = BTreeMap::new();
1196        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1197        let state1 = Arc::new(create_mock_state(&mut test_block_builder, 1, B256::random()));
1198        let hash1 = state1.hash();
1199        let state2 = Arc::new(create_mock_state(&mut test_block_builder, 2, hash1));
1200        let hash2 = state2.hash();
1201        hash_by_number.insert(1, hash1);
1202        hash_by_number.insert(2, hash2);
1203        state_by_hash.insert(hash1, state1);
1204        state_by_hash.insert(hash2, state2);
1205
1206        let in_memory_state = InMemoryState::new(state_by_hash, hash_by_number, None);
1207        let head_state = in_memory_state.head_state().unwrap();
1208
1209        assert_eq!(head_state.hash(), hash2);
1210        assert_eq!(head_state.number(), 2);
1211    }
1212
1213    #[test]
1214    fn test_in_memory_state_impl_pending_state() {
1215        let pending_number = rand::rng().random::<u64>();
1216        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1217        let pending_state =
1218            create_mock_state(&mut test_block_builder, pending_number, B256::random());
1219        let pending_hash = pending_state.hash();
1220
1221        let in_memory_state =
1222            InMemoryState::new(B256Map::default(), BTreeMap::new(), Some(pending_state));
1223
1224        let result = in_memory_state.pending_state();
1225        assert!(result.is_some());
1226        let actual_pending_state = result.unwrap();
1227        assert_eq!(actual_pending_state.block.recovered_block().hash(), pending_hash);
1228        assert_eq!(actual_pending_state.block.recovered_block().number, pending_number);
1229    }
1230
1231    #[test]
1232    fn test_in_memory_state_impl_no_pending_state() {
1233        let in_memory_state: InMemoryState =
1234            InMemoryState::new(B256Map::default(), BTreeMap::new(), None);
1235
1236        assert_eq!(in_memory_state.pending_state(), None);
1237    }
1238
1239    #[test]
1240    fn test_state() {
1241        let number = rand::rng().random::<u64>();
1242        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1243        let block = test_block_builder.get_executed_block_with_number(number, B256::random());
1244
1245        let state = BlockState::new(block.clone());
1246
1247        assert_eq!(state.block(), block);
1248        assert_eq!(state.hash(), block.recovered_block().hash());
1249        assert_eq!(state.number(), number);
1250        assert_eq!(state.state_root(), block.recovered_block().state_root);
1251    }
1252
1253    #[test]
1254    fn test_state_receipts() {
1255        let receipts = vec![vec![Receipt::default()]];
1256        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1257        let block =
1258            test_block_builder.get_executed_block_with_receipts(receipts.clone(), B256::random());
1259
1260        let state = BlockState::new(block);
1261
1262        assert_eq!(state.receipts(), receipts.first().unwrap());
1263    }
1264
1265    #[test]
1266    fn test_in_memory_state_chain_update() {
1267        let state: CanonicalInMemoryState = CanonicalInMemoryState::empty();
1268        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1269        let block1 = test_block_builder.get_executed_block_with_number(0, B256::random());
1270        let block2 = test_block_builder.get_executed_block_with_number(0, B256::random());
1271        let chain = NewCanonicalChain::Commit { new: vec![block1.clone()] };
1272        state.update_chain(chain);
1273        assert_eq!(
1274            state.head_state().unwrap().block_ref().recovered_block().hash(),
1275            block1.recovered_block().hash()
1276        );
1277        assert_eq!(
1278            state.state_by_number(0).unwrap().block_ref().recovered_block().hash(),
1279            block1.recovered_block().hash()
1280        );
1281
1282        let chain = NewCanonicalChain::Reorg { new: vec![block2.clone()], old: vec![block1] };
1283        state.update_chain(chain);
1284        assert_eq!(
1285            state.head_state().unwrap().block_ref().recovered_block().hash(),
1286            block2.recovered_block().hash()
1287        );
1288        assert_eq!(
1289            state.state_by_number(0).unwrap().block_ref().recovered_block().hash(),
1290            block2.recovered_block().hash()
1291        );
1292
1293        assert_eq!(state.inner.in_memory_state.block_count(), 1);
1294    }
1295
1296    #[test]
1297    fn test_in_memory_state_set_pending_block() {
1298        let state: CanonicalInMemoryState = CanonicalInMemoryState::empty();
1299        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1300
1301        // First random block
1302        let block1 = test_block_builder.get_executed_block_with_number(0, B256::random());
1303
1304        // Second block with parent hash of the first block
1305        let block2 =
1306            test_block_builder.get_executed_block_with_number(1, block1.recovered_block().hash());
1307
1308        // Commit the two blocks
1309        let chain = NewCanonicalChain::Commit { new: vec![block1.clone(), block2.clone()] };
1310        state.update_chain(chain);
1311
1312        // Assert that the pending state is None before setting it
1313        assert!(state.pending_state().is_none());
1314
1315        // Set the pending block
1316        state.set_pending_block(block2.clone());
1317
1318        // Check the pending state
1319        assert_eq!(
1320            state.pending_state().unwrap(),
1321            BlockState::with_parent(block2.clone(), Some(Arc::new(BlockState::new(block1))))
1322        );
1323
1324        // Check the pending block
1325        assert_eq!(state.pending_block().unwrap(), block2.recovered_block().sealed_block().clone());
1326
1327        // Check the pending block number and hash
1328        assert_eq!(
1329            state.pending_block_num_hash().unwrap(),
1330            BlockNumHash { number: 1, hash: block2.recovered_block().hash() }
1331        );
1332
1333        // Check the pending header
1334        assert_eq!(state.pending_header().unwrap(), block2.recovered_block().header().clone());
1335
1336        // Check the pending sealed header
1337        assert_eq!(
1338            state.pending_sealed_header().unwrap(),
1339            block2.recovered_block().clone_sealed_header()
1340        );
1341
1342        // Check the pending block with senders
1343        assert_eq!(state.pending_recovered_block().unwrap(), block2.recovered_block().clone());
1344
1345        // Check the pending block and receipts
1346        assert_eq!(
1347            state.pending_block_and_receipts().unwrap(),
1348            (block2.recovered_block().clone(), vec![])
1349        );
1350    }
1351
1352    #[test]
1353    fn test_canonical_in_memory_state_state_provider() {
1354        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1355        let block1 = test_block_builder.get_executed_block_with_number(1, B256::random());
1356        let block2 =
1357            test_block_builder.get_executed_block_with_number(2, block1.recovered_block().hash());
1358        let block3 =
1359            test_block_builder.get_executed_block_with_number(3, block2.recovered_block().hash());
1360
1361        let state1 = Arc::new(BlockState::new(block1.clone()));
1362        let state2 = Arc::new(BlockState::with_parent(block2.clone(), Some(state1.clone())));
1363        let state3 = Arc::new(BlockState::with_parent(block3.clone(), Some(state2.clone())));
1364
1365        let mut blocks = B256Map::default();
1366        blocks.insert(block1.recovered_block().hash(), state1);
1367        blocks.insert(block2.recovered_block().hash(), state2);
1368        blocks.insert(block3.recovered_block().hash(), state3);
1369
1370        let mut numbers = BTreeMap::new();
1371        numbers.insert(1, block1.recovered_block().hash());
1372        numbers.insert(2, block2.recovered_block().hash());
1373        numbers.insert(3, block3.recovered_block().hash());
1374
1375        let canonical_state = CanonicalInMemoryState::new(blocks, numbers, None, None, None);
1376
1377        let historical: StateProviderBox = Box::new(MockStateProvider);
1378
1379        let overlay_provider =
1380            canonical_state.state_provider(block3.recovered_block().hash(), historical);
1381
1382        assert_eq!(overlay_provider.in_memory.len(), 3);
1383        assert_eq!(overlay_provider.in_memory[0].recovered_block().number, 3);
1384        assert_eq!(overlay_provider.in_memory[1].recovered_block().number, 2);
1385        assert_eq!(overlay_provider.in_memory[2].recovered_block().number, 1);
1386
1387        assert_eq!(
1388            overlay_provider.in_memory[0].recovered_block().parent_hash,
1389            overlay_provider.in_memory[1].recovered_block().hash()
1390        );
1391        assert_eq!(
1392            overlay_provider.in_memory[1].recovered_block().parent_hash,
1393            overlay_provider.in_memory[2].recovered_block().hash()
1394        );
1395
1396        let unknown_hash = B256::random();
1397        let empty_overlay_provider =
1398            canonical_state.state_provider(unknown_hash, Box::new(MockStateProvider));
1399        assert_eq!(empty_overlay_provider.in_memory.len(), 0);
1400    }
1401
1402    #[test]
1403    fn test_canonical_in_memory_state_canonical_chain_empty() {
1404        let state: CanonicalInMemoryState = CanonicalInMemoryState::empty();
1405        assert!(state.canonical_chain().next().is_none());
1406    }
1407
1408    #[test]
1409    fn test_canonical_in_memory_state_canonical_chain_single_block() {
1410        let block = TestBlockBuilder::eth().get_executed_block_with_number(1, B256::random());
1411        let hash = block.recovered_block().hash();
1412        let mut blocks = B256Map::default();
1413        blocks.insert(hash, Arc::new(BlockState::new(block)));
1414        let mut numbers = BTreeMap::new();
1415        numbers.insert(1, hash);
1416
1417        let state = CanonicalInMemoryState::new(blocks, numbers, None, None, None);
1418        let chain: Vec<_> = state.canonical_chain().collect();
1419
1420        assert_eq!(chain.len(), 1);
1421        assert_eq!(chain[0].number(), 1);
1422        assert_eq!(chain[0].hash(), hash);
1423    }
1424
1425    #[test]
1426    fn test_canonical_in_memory_state_canonical_chain_multiple_blocks() {
1427        let mut parent_hash = B256::random();
1428        let mut block_builder = TestBlockBuilder::eth();
1429        let state: CanonicalInMemoryState = CanonicalInMemoryState::empty();
1430
1431        for i in 1..=3 {
1432            let block = block_builder.get_executed_block_with_number(i, parent_hash);
1433            let hash = block.recovered_block().hash();
1434            state.update_blocks(Some(block), None);
1435            parent_hash = hash;
1436        }
1437
1438        let chain: Vec<_> = state.canonical_chain().collect();
1439
1440        assert_eq!(chain.len(), 3);
1441        assert_eq!(chain[0].number(), 3);
1442        assert_eq!(chain[1].number(), 2);
1443        assert_eq!(chain[2].number(), 1);
1444    }
1445
1446    // ensures the pending block is not part of the canonical chain
1447    #[test]
1448    fn test_canonical_in_memory_state_canonical_chain_with_pending_block() {
1449        let mut parent_hash = B256::random();
1450        let mut block_builder = TestBlockBuilder::<EthPrimitives>::eth();
1451        let state: CanonicalInMemoryState = CanonicalInMemoryState::empty();
1452
1453        for i in 1..=2 {
1454            let block = block_builder.get_executed_block_with_number(i, parent_hash);
1455            let hash = block.recovered_block().hash();
1456            state.update_blocks(Some(block), None);
1457            parent_hash = hash;
1458        }
1459
1460        let pending_block = block_builder.get_executed_block_with_number(3, parent_hash);
1461        state.set_pending_block(pending_block);
1462        let chain: Vec<_> = state.canonical_chain().collect();
1463
1464        assert_eq!(chain.len(), 2);
1465        assert_eq!(chain[0].number(), 2);
1466        assert_eq!(chain[1].number(), 1);
1467    }
1468
1469    #[test]
1470    fn test_block_state_parent_blocks() {
1471        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1472        let chain = create_mock_state_chain(&mut test_block_builder, 4);
1473
1474        let parents: Vec<_> = chain[3].parent_state_chain().collect();
1475        assert_eq!(parents.len(), 3);
1476        assert_eq!(parents[0].block().recovered_block().number, 3);
1477        assert_eq!(parents[1].block().recovered_block().number, 2);
1478        assert_eq!(parents[2].block().recovered_block().number, 1);
1479
1480        let parents: Vec<_> = chain[2].parent_state_chain().collect();
1481        assert_eq!(parents.len(), 2);
1482        assert_eq!(parents[0].block().recovered_block().number, 2);
1483        assert_eq!(parents[1].block().recovered_block().number, 1);
1484
1485        assert_eq!(chain[0].parent_state_chain().count(), 0);
1486    }
1487
1488    #[test]
1489    fn test_block_state_single_block_state_chain() {
1490        let single_block_number = 1;
1491        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1492        let single_block =
1493            create_mock_state(&mut test_block_builder, single_block_number, B256::random());
1494        let single_block_hash = single_block.block().recovered_block().hash();
1495
1496        assert_eq!(single_block.parent_state_chain().count(), 0);
1497
1498        let block_state_chain = single_block.chain().collect::<Vec<_>>();
1499        assert_eq!(block_state_chain.len(), 1);
1500        assert_eq!(block_state_chain[0].block().recovered_block().number, single_block_number);
1501        assert_eq!(block_state_chain[0].block().recovered_block().hash(), single_block_hash);
1502    }
1503
1504    #[test]
1505    fn test_block_state_chain() {
1506        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1507        let chain = create_mock_state_chain(&mut test_block_builder, 3);
1508
1509        let block_state_chain = chain[2].chain().collect::<Vec<_>>();
1510        assert_eq!(block_state_chain.len(), 3);
1511        assert_eq!(block_state_chain[0].block().recovered_block().number, 3);
1512        assert_eq!(block_state_chain[1].block().recovered_block().number, 2);
1513        assert_eq!(block_state_chain[2].block().recovered_block().number, 1);
1514
1515        let block_state_chain = chain[1].chain().collect::<Vec<_>>();
1516        assert_eq!(block_state_chain.len(), 2);
1517        assert_eq!(block_state_chain[0].block().recovered_block().number, 2);
1518        assert_eq!(block_state_chain[1].block().recovered_block().number, 1);
1519
1520        let block_state_chain = chain[0].chain().collect::<Vec<_>>();
1521        assert_eq!(block_state_chain.len(), 1);
1522        assert_eq!(block_state_chain[0].block().recovered_block().number, 1);
1523    }
1524
1525    #[test]
1526    fn test_to_chain_notification() {
1527        // Generate 4 blocks
1528        let mut test_block_builder: TestBlockBuilder = TestBlockBuilder::default();
1529        let block0 = test_block_builder.get_executed_block_with_number(0, B256::random());
1530        let block1 =
1531            test_block_builder.get_executed_block_with_number(1, block0.recovered_block.hash());
1532        let block1a =
1533            test_block_builder.get_executed_block_with_number(1, block0.recovered_block.hash());
1534        let block2 =
1535            test_block_builder.get_executed_block_with_number(2, block1.recovered_block.hash());
1536        let block2a =
1537            test_block_builder.get_executed_block_with_number(2, block1.recovered_block.hash());
1538
1539        // Test commit notification
1540        let chain_commit = NewCanonicalChain::Commit { new: vec![block0.clone(), block1.clone()] };
1541
1542        // Build expected trie data map
1543        let mut expected_trie_data = BTreeMap::new();
1544        expected_trie_data.insert(0, LazyTrieData::ready(block0.trie_data()));
1545        expected_trie_data.insert(1, LazyTrieData::ready(block1.trie_data()));
1546
1547        // Build expected execution outcome (first_block matches first block number)
1548        let commit_execution_outcome = ExecutionOutcome {
1549            receipts: vec![vec![], vec![]],
1550            requests: vec![Requests::default(), Requests::default()],
1551            first_block: 0,
1552            ..Default::default()
1553        };
1554
1555        assert_eq!(
1556            chain_commit.to_chain_notification(),
1557            CanonStateNotification::Commit {
1558                new: Arc::new(Chain::new(
1559                    vec![block0.recovered_block().clone(), block1.recovered_block().clone()],
1560                    commit_execution_outcome,
1561                    expected_trie_data,
1562                ))
1563            }
1564        );
1565
1566        // Test reorg notification
1567        let chain_reorg = NewCanonicalChain::Reorg {
1568            new: vec![block1a.clone(), block2a.clone()],
1569            old: vec![block1.clone(), block2.clone()],
1570        };
1571
1572        // Build expected trie data for old chain
1573        let mut old_trie_data = BTreeMap::new();
1574        old_trie_data.insert(1, LazyTrieData::ready(block1.trie_data()));
1575        old_trie_data.insert(2, LazyTrieData::ready(block2.trie_data()));
1576
1577        // Build expected trie data for new chain
1578        let mut new_trie_data = BTreeMap::new();
1579        new_trie_data.insert(1, LazyTrieData::ready(block1a.trie_data()));
1580        new_trie_data.insert(2, LazyTrieData::ready(block2a.trie_data()));
1581
1582        // Build expected execution outcome for reorg chains (first_block matches first block
1583        // number)
1584        let reorg_execution_outcome = ExecutionOutcome {
1585            receipts: vec![vec![], vec![]],
1586            requests: vec![Requests::default(), Requests::default()],
1587            first_block: 1,
1588            ..Default::default()
1589        };
1590
1591        assert_eq!(
1592            chain_reorg.to_chain_notification(),
1593            CanonStateNotification::Reorg {
1594                old: Arc::new(Chain::new(
1595                    vec![block1.recovered_block().clone(), block2.recovered_block().clone()],
1596                    reorg_execution_outcome.clone(),
1597                    old_trie_data,
1598                )),
1599                new: Arc::new(Chain::new(
1600                    vec![block1a.recovered_block().clone(), block2a.recovered_block().clone()],
1601                    reorg_execution_outcome,
1602                    new_trie_data,
1603                ))
1604            }
1605        );
1606    }
1607}