reth_engine_tree/tree/
state.rs

1//! Functionality related to tree state.
2
3use crate::engine::EngineApiKind;
4use alloy_eips::{eip1898::BlockWithParent, merge::EPOCH_SLOTS, BlockNumHash};
5use alloy_primitives::{
6    map::{HashMap, HashSet},
7    BlockNumber, B256,
8};
9use reth_chain_state::{EthPrimitives, ExecutedBlockWithTrieUpdates};
10use reth_primitives_traits::{AlloyBlockHeader, NodePrimitives, SealedHeader};
11use reth_trie::updates::TrieUpdates;
12use std::{
13    collections::{btree_map, hash_map, BTreeMap, VecDeque},
14    ops::Bound,
15    sync::Arc,
16};
17use tracing::debug;
18
19/// Default number of blocks to retain persisted trie updates
20const DEFAULT_PERSISTED_TRIE_UPDATES_RETENTION: u64 = EPOCH_SLOTS * 2;
21
22/// Number of blocks to retain persisted trie updates for OP Stack chains
23/// OP Stack chains only need `EPOCH_BLOCKS` as reorgs are relevant only when
24/// op-node reorgs to the same chain twice
25const OPSTACK_PERSISTED_TRIE_UPDATES_RETENTION: u64 = EPOCH_SLOTS;
26
27/// Keeps track of the state of the tree.
28///
29/// ## Invariants
30///
31/// - This only stores blocks that are connected to the canonical chain.
32/// - All executed blocks are valid and have been executed.
33#[derive(Debug, Default)]
34pub struct TreeState<N: NodePrimitives = EthPrimitives> {
35    /// __All__ unique executed blocks by block hash that are connected to the canonical chain.
36    ///
37    /// This includes blocks of all forks.
38    pub(crate) blocks_by_hash: HashMap<B256, ExecutedBlockWithTrieUpdates<N>>,
39    /// Executed blocks grouped by their respective block number.
40    ///
41    /// This maps unique block number to all known blocks for that height.
42    ///
43    /// Note: there can be multiple blocks at the same height due to forks.
44    pub(crate) blocks_by_number: BTreeMap<BlockNumber, Vec<ExecutedBlockWithTrieUpdates<N>>>,
45    /// Map of any parent block hash to its children.
46    pub(crate) parent_to_child: HashMap<B256, HashSet<B256>>,
47    /// Map of hash to trie updates for canonical blocks that are persisted but not finalized.
48    ///
49    /// Contains the block number for easy removal.
50    pub(crate) persisted_trie_updates: HashMap<B256, (BlockNumber, Arc<TrieUpdates>)>,
51    /// Currently tracked canonical head of the chain.
52    pub(crate) current_canonical_head: BlockNumHash,
53    /// The engine API variant of this handler
54    pub(crate) engine_kind: EngineApiKind,
55}
56
57impl<N: NodePrimitives> TreeState<N> {
58    /// Returns a new, empty tree state that points to the given canonical head.
59    pub(crate) fn new(current_canonical_head: BlockNumHash, engine_kind: EngineApiKind) -> Self {
60        Self {
61            blocks_by_hash: HashMap::default(),
62            blocks_by_number: BTreeMap::new(),
63            current_canonical_head,
64            parent_to_child: HashMap::default(),
65            persisted_trie_updates: HashMap::default(),
66            engine_kind,
67        }
68    }
69
70    /// Resets the state and points to the given canonical head.
71    pub(crate) fn reset(&mut self, current_canonical_head: BlockNumHash) {
72        *self = Self::new(current_canonical_head, self.engine_kind);
73    }
74
75    /// Returns the number of executed blocks stored.
76    pub(crate) fn block_count(&self) -> usize {
77        self.blocks_by_hash.len()
78    }
79
80    /// Returns the [`ExecutedBlockWithTrieUpdates`] by hash.
81    pub(crate) fn executed_block_by_hash(
82        &self,
83        hash: B256,
84    ) -> Option<&ExecutedBlockWithTrieUpdates<N>> {
85        self.blocks_by_hash.get(&hash)
86    }
87
88    /// Returns the sealed block header by hash.
89    pub(crate) fn sealed_header_by_hash(
90        &self,
91        hash: &B256,
92    ) -> Option<SealedHeader<N::BlockHeader>> {
93        self.blocks_by_hash.get(hash).map(|b| b.sealed_block().sealed_header().clone())
94    }
95
96    /// Returns all available blocks for the given hash that lead back to the canonical chain, from
97    /// newest to oldest. And the parent hash of the oldest block that is missing from the buffer.
98    ///
99    /// Returns `None` if the block for the given hash is not found.
100    pub(crate) fn blocks_by_hash(
101        &self,
102        hash: B256,
103    ) -> Option<(B256, Vec<ExecutedBlockWithTrieUpdates<N>>)> {
104        let block = self.blocks_by_hash.get(&hash).cloned()?;
105        let mut parent_hash = block.recovered_block().parent_hash();
106        let mut blocks = vec![block];
107        while let Some(executed) = self.blocks_by_hash.get(&parent_hash) {
108            parent_hash = executed.recovered_block().parent_hash();
109            blocks.push(executed.clone());
110        }
111
112        Some((parent_hash, blocks))
113    }
114
115    /// Insert executed block into the state.
116    pub(crate) fn insert_executed(&mut self, executed: ExecutedBlockWithTrieUpdates<N>) {
117        let hash = executed.recovered_block().hash();
118        let parent_hash = executed.recovered_block().parent_hash();
119        let block_number = executed.recovered_block().number();
120
121        if self.blocks_by_hash.contains_key(&hash) {
122            return;
123        }
124
125        self.blocks_by_hash.insert(hash, executed.clone());
126
127        self.blocks_by_number.entry(block_number).or_default().push(executed);
128
129        self.parent_to_child.entry(parent_hash).or_default().insert(hash);
130
131        for children in self.parent_to_child.values_mut() {
132            children.retain(|child| self.blocks_by_hash.contains_key(child));
133        }
134    }
135
136    /// Remove single executed block by its hash.
137    ///
138    /// ## Returns
139    ///
140    /// The removed block and the block hashes of its children.
141    fn remove_by_hash(
142        &mut self,
143        hash: B256,
144    ) -> Option<(ExecutedBlockWithTrieUpdates<N>, HashSet<B256>)> {
145        let executed = self.blocks_by_hash.remove(&hash)?;
146
147        // Remove this block from collection of children of its parent block.
148        let parent_entry = self.parent_to_child.entry(executed.recovered_block().parent_hash());
149        if let hash_map::Entry::Occupied(mut entry) = parent_entry {
150            entry.get_mut().remove(&hash);
151
152            if entry.get().is_empty() {
153                entry.remove();
154            }
155        }
156
157        // Remove point to children of this block.
158        let children = self.parent_to_child.remove(&hash).unwrap_or_default();
159
160        // Remove this block from `blocks_by_number`.
161        let block_number_entry = self.blocks_by_number.entry(executed.recovered_block().number());
162        if let btree_map::Entry::Occupied(mut entry) = block_number_entry {
163            // We have to find the index of the block since it exists in a vec
164            if let Some(index) = entry.get().iter().position(|b| b.recovered_block().hash() == hash)
165            {
166                entry.get_mut().swap_remove(index);
167
168                // If there are no blocks left then remove the entry for this block
169                if entry.get().is_empty() {
170                    entry.remove();
171                }
172            }
173        }
174
175        Some((executed, children))
176    }
177
178    /// Returns whether or not the hash is part of the canonical chain.
179    pub(crate) fn is_canonical(&self, hash: B256) -> bool {
180        let mut current_block = self.current_canonical_head.hash;
181        if current_block == hash {
182            return true
183        }
184
185        while let Some(executed) = self.blocks_by_hash.get(&current_block) {
186            current_block = executed.recovered_block().parent_hash();
187            if current_block == hash {
188                return true
189            }
190        }
191
192        false
193    }
194
195    /// Removes canonical blocks below the upper bound, only if the last persisted hash is
196    /// part of the canonical chain.
197    pub(crate) fn remove_canonical_until(
198        &mut self,
199        upper_bound: BlockNumber,
200        last_persisted_hash: B256,
201    ) {
202        debug!(target: "engine::tree", ?upper_bound, ?last_persisted_hash, "Removing canonical blocks from the tree");
203
204        // If the last persisted hash is not canonical, then we don't want to remove any canonical
205        // blocks yet.
206        if !self.is_canonical(last_persisted_hash) {
207            return
208        }
209
210        // First, let's walk back the canonical chain and remove canonical blocks lower than the
211        // upper bound
212        let mut current_block = self.current_canonical_head.hash;
213        while let Some(executed) = self.blocks_by_hash.get(&current_block) {
214            current_block = executed.recovered_block().parent_hash();
215            if executed.recovered_block().number() <= upper_bound {
216                let num_hash = executed.recovered_block().num_hash();
217                debug!(target: "engine::tree", ?num_hash, "Attempting to remove block walking back from the head");
218                if let Some((mut removed, _)) =
219                    self.remove_by_hash(executed.recovered_block().hash())
220                {
221                    debug!(target: "engine::tree", ?num_hash, "Removed block walking back from the head");
222                    // finally, move the trie updates
223                    let Some(trie_updates) = removed.trie.take_present() else {
224                        debug!(target: "engine::tree", ?num_hash, "No trie updates found for persisted block");
225                        continue;
226                    };
227                    self.persisted_trie_updates.insert(
228                        removed.recovered_block().hash(),
229                        (removed.recovered_block().number(), trie_updates),
230                    );
231                }
232            }
233        }
234        debug!(target: "engine::tree", ?upper_bound, ?last_persisted_hash, "Removed canonical blocks from the tree");
235    }
236
237    /// Prunes old persisted trie updates based on the current block number
238    /// and chain type (OP Stack or regular)
239    pub(crate) fn prune_persisted_trie_updates(&mut self) {
240        let retention_blocks = if self.engine_kind.is_opstack() {
241            OPSTACK_PERSISTED_TRIE_UPDATES_RETENTION
242        } else {
243            DEFAULT_PERSISTED_TRIE_UPDATES_RETENTION
244        };
245
246        let earliest_block_to_retain =
247            self.current_canonical_head.number.saturating_sub(retention_blocks);
248
249        self.persisted_trie_updates
250            .retain(|_, (block_number, _)| *block_number > earliest_block_to_retain);
251    }
252
253    /// Removes all blocks that are below the finalized block, as well as removing non-canonical
254    /// sidechains that fork from below the finalized block.
255    pub(crate) fn prune_finalized_sidechains(&mut self, finalized_num_hash: BlockNumHash) {
256        let BlockNumHash { number: finalized_num, hash: finalized_hash } = finalized_num_hash;
257
258        // We remove disconnected sidechains in three steps:
259        // * first, remove everything with a block number __below__ the finalized block.
260        // * next, we populate a vec with parents __at__ the finalized block.
261        // * finally, we iterate through the vec, removing children until the vec is empty
262        // (BFS).
263
264        // We _exclude_ the finalized block because we will be dealing with the blocks __at__
265        // the finalized block later.
266        let blocks_to_remove = self
267            .blocks_by_number
268            .range((Bound::Unbounded, Bound::Excluded(finalized_num)))
269            .flat_map(|(_, blocks)| blocks.iter().map(|b| b.recovered_block().hash()))
270            .collect::<Vec<_>>();
271        for hash in blocks_to_remove {
272            if let Some((removed, _)) = self.remove_by_hash(hash) {
273                debug!(target: "engine::tree", num_hash=?removed.recovered_block().num_hash(), "Removed finalized sidechain block");
274            }
275        }
276
277        self.prune_persisted_trie_updates();
278
279        // The only block that should remain at the `finalized` number now, is the finalized
280        // block, if it exists.
281        //
282        // For all other blocks, we  first put their children into this vec.
283        // Then, we will iterate over them, removing them, adding their children, etc,
284        // until the vec is empty.
285        let mut blocks_to_remove = self.blocks_by_number.remove(&finalized_num).unwrap_or_default();
286
287        // re-insert the finalized hash if we removed it
288        if let Some(position) =
289            blocks_to_remove.iter().position(|b| b.recovered_block().hash() == finalized_hash)
290        {
291            let finalized_block = blocks_to_remove.swap_remove(position);
292            self.blocks_by_number.insert(finalized_num, vec![finalized_block]);
293        }
294
295        let mut blocks_to_remove = blocks_to_remove
296            .into_iter()
297            .map(|e| e.recovered_block().hash())
298            .collect::<VecDeque<_>>();
299        while let Some(block) = blocks_to_remove.pop_front() {
300            if let Some((removed, children)) = self.remove_by_hash(block) {
301                debug!(target: "engine::tree", num_hash=?removed.recovered_block().num_hash(), "Removed finalized sidechain child block");
302                blocks_to_remove.extend(children);
303            }
304        }
305    }
306
307    /// Remove all blocks up to __and including__ the given block number.
308    ///
309    /// If a finalized hash is provided, the only non-canonical blocks which will be removed are
310    /// those which have a fork point at or below the finalized hash.
311    ///
312    /// Canonical blocks below the upper bound will still be removed.
313    ///
314    /// NOTE: if the finalized block is greater than the upper bound, the only blocks that will be
315    /// removed are canonical blocks and sidechains that fork below the `upper_bound`. This is the
316    /// same behavior as if the `finalized_num` were `Some(upper_bound)`.
317    pub(crate) fn remove_until(
318        &mut self,
319        upper_bound: BlockNumHash,
320        last_persisted_hash: B256,
321        finalized_num_hash: Option<BlockNumHash>,
322    ) {
323        debug!(target: "engine::tree", ?upper_bound, ?finalized_num_hash, "Removing blocks from the tree");
324
325        // If the finalized num is ahead of the upper bound, and exists, we need to instead ensure
326        // that the only blocks removed, are canonical blocks less than the upper bound
327        let finalized_num_hash = finalized_num_hash.map(|mut finalized| {
328            if upper_bound.number < finalized.number {
329                finalized = upper_bound;
330                debug!(target: "engine::tree", ?finalized, "Adjusted upper bound");
331            }
332            finalized
333        });
334
335        // We want to do two things:
336        // * remove canonical blocks that are persisted
337        // * remove forks whose root are below the finalized block
338        // We can do this in 2 steps:
339        // * remove all canonical blocks below the upper bound
340        // * fetch the number of the finalized hash, removing any sidechains that are __below__ the
341        // finalized block
342        self.remove_canonical_until(upper_bound.number, last_persisted_hash);
343
344        // Now, we have removed canonical blocks (assuming the upper bound is above the finalized
345        // block) and only have sidechains below the finalized block.
346        if let Some(finalized_num_hash) = finalized_num_hash {
347            self.prune_finalized_sidechains(finalized_num_hash);
348        }
349    }
350
351    /// Determines if the second block is a direct descendant of the first block.
352    ///
353    /// If the two blocks are the same, this returns `false`.
354    pub(crate) fn is_descendant(&self, first: BlockNumHash, second: BlockWithParent) -> bool {
355        // If the second block's parent is the first block's hash, then it is a direct descendant
356        // and we can return early.
357        if second.parent == first.hash {
358            return true
359        }
360
361        // If the second block is lower than, or has the same block number, they are not
362        // descendants.
363        if second.block.number <= first.number {
364            return false
365        }
366
367        // iterate through parents of the second until we reach the number
368        let Some(mut current_block) = self.blocks_by_hash.get(&second.parent) else {
369            // If we can't find its parent in the tree, we can't continue, so return false
370            return false
371        };
372
373        while current_block.recovered_block().number() > first.number + 1 {
374            let Some(block) =
375                self.blocks_by_hash.get(&current_block.recovered_block().parent_hash())
376            else {
377                // If we can't find its parent in the tree, we can't continue, so return false
378                return false
379            };
380
381            current_block = block;
382        }
383
384        // Now the block numbers should be equal, so we compare hashes.
385        current_block.recovered_block().parent_hash() == first.hash
386    }
387
388    /// Updates the canonical head to the given block.
389    pub(crate) const fn set_canonical_head(&mut self, new_head: BlockNumHash) {
390        self.current_canonical_head = new_head;
391    }
392
393    /// Returns the tracked canonical head.
394    pub(crate) const fn canonical_head(&self) -> &BlockNumHash {
395        &self.current_canonical_head
396    }
397
398    /// Returns the block hash of the canonical head.
399    pub(crate) const fn canonical_block_hash(&self) -> B256 {
400        self.canonical_head().hash
401    }
402
403    /// Returns the block number of the canonical head.
404    pub(crate) const fn canonical_block_number(&self) -> BlockNumber {
405        self.canonical_head().number
406    }
407}
408
409#[cfg(test)]
410mod tests {
411    use super::*;
412    use reth_chain_state::test_utils::TestBlockBuilder;
413
414    #[test]
415    fn test_tree_state_normal_descendant() {
416        let mut tree_state = TreeState::new(BlockNumHash::default(), EngineApiKind::Ethereum);
417        let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..4).collect();
418
419        tree_state.insert_executed(blocks[0].clone());
420        assert!(tree_state.is_descendant(
421            blocks[0].recovered_block().num_hash(),
422            blocks[1].recovered_block().block_with_parent()
423        ));
424
425        tree_state.insert_executed(blocks[1].clone());
426
427        assert!(tree_state.is_descendant(
428            blocks[0].recovered_block().num_hash(),
429            blocks[2].recovered_block().block_with_parent()
430        ));
431        assert!(tree_state.is_descendant(
432            blocks[1].recovered_block().num_hash(),
433            blocks[2].recovered_block().block_with_parent()
434        ));
435    }
436
437    #[tokio::test]
438    async fn test_tree_state_insert_executed() {
439        let mut tree_state = TreeState::new(BlockNumHash::default(), EngineApiKind::Ethereum);
440        let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..4).collect();
441
442        tree_state.insert_executed(blocks[0].clone());
443        tree_state.insert_executed(blocks[1].clone());
444
445        assert_eq!(
446            tree_state.parent_to_child.get(&blocks[0].recovered_block().hash()),
447            Some(&HashSet::from_iter([blocks[1].recovered_block().hash()]))
448        );
449
450        assert!(!tree_state.parent_to_child.contains_key(&blocks[1].recovered_block().hash()));
451
452        tree_state.insert_executed(blocks[2].clone());
453
454        assert_eq!(
455            tree_state.parent_to_child.get(&blocks[1].recovered_block().hash()),
456            Some(&HashSet::from_iter([blocks[2].recovered_block().hash()]))
457        );
458        assert!(tree_state.parent_to_child.contains_key(&blocks[1].recovered_block().hash()));
459
460        assert!(!tree_state.parent_to_child.contains_key(&blocks[2].recovered_block().hash()));
461    }
462
463    #[tokio::test]
464    async fn test_tree_state_insert_executed_with_reorg() {
465        let mut tree_state = TreeState::new(BlockNumHash::default(), EngineApiKind::Ethereum);
466        let mut test_block_builder = TestBlockBuilder::eth();
467        let blocks: Vec<_> = test_block_builder.get_executed_blocks(1..6).collect();
468
469        for block in &blocks {
470            tree_state.insert_executed(block.clone());
471        }
472        assert_eq!(tree_state.blocks_by_hash.len(), 5);
473
474        let fork_block_3 = test_block_builder
475            .get_executed_block_with_number(3, blocks[1].recovered_block().hash());
476        let fork_block_4 = test_block_builder
477            .get_executed_block_with_number(4, fork_block_3.recovered_block().hash());
478        let fork_block_5 = test_block_builder
479            .get_executed_block_with_number(5, fork_block_4.recovered_block().hash());
480
481        tree_state.insert_executed(fork_block_3.clone());
482        tree_state.insert_executed(fork_block_4.clone());
483        tree_state.insert_executed(fork_block_5.clone());
484
485        assert_eq!(tree_state.blocks_by_hash.len(), 8);
486        assert_eq!(tree_state.blocks_by_number[&3].len(), 2); // two blocks at height 3 (original and fork)
487        assert_eq!(tree_state.parent_to_child[&blocks[1].recovered_block().hash()].len(), 2); // block 2 should have two children
488
489        // verify that we can insert the same block again without issues
490        tree_state.insert_executed(fork_block_4.clone());
491        assert_eq!(tree_state.blocks_by_hash.len(), 8);
492
493        assert!(tree_state.parent_to_child[&fork_block_3.recovered_block().hash()]
494            .contains(&fork_block_4.recovered_block().hash()));
495        assert!(tree_state.parent_to_child[&fork_block_4.recovered_block().hash()]
496            .contains(&fork_block_5.recovered_block().hash()));
497
498        assert_eq!(tree_state.blocks_by_number[&4].len(), 2);
499        assert_eq!(tree_state.blocks_by_number[&5].len(), 2);
500    }
501
502    #[tokio::test]
503    async fn test_tree_state_remove_before() {
504        let start_num_hash = BlockNumHash::default();
505        let mut tree_state = TreeState::new(start_num_hash, EngineApiKind::Ethereum);
506        let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..6).collect();
507
508        for block in &blocks {
509            tree_state.insert_executed(block.clone());
510        }
511
512        let last = blocks.last().unwrap();
513
514        // set the canonical head
515        tree_state.set_canonical_head(last.recovered_block().num_hash());
516
517        // inclusive bound, so we should remove anything up to and including 2
518        tree_state.remove_until(
519            BlockNumHash::new(2, blocks[1].recovered_block().hash()),
520            start_num_hash.hash,
521            Some(blocks[1].recovered_block().num_hash()),
522        );
523
524        assert!(!tree_state.blocks_by_hash.contains_key(&blocks[0].recovered_block().hash()));
525        assert!(!tree_state.blocks_by_hash.contains_key(&blocks[1].recovered_block().hash()));
526        assert!(!tree_state.blocks_by_number.contains_key(&1));
527        assert!(!tree_state.blocks_by_number.contains_key(&2));
528
529        assert!(tree_state.blocks_by_hash.contains_key(&blocks[2].recovered_block().hash()));
530        assert!(tree_state.blocks_by_hash.contains_key(&blocks[3].recovered_block().hash()));
531        assert!(tree_state.blocks_by_hash.contains_key(&blocks[4].recovered_block().hash()));
532        assert!(tree_state.blocks_by_number.contains_key(&3));
533        assert!(tree_state.blocks_by_number.contains_key(&4));
534        assert!(tree_state.blocks_by_number.contains_key(&5));
535
536        assert!(!tree_state.parent_to_child.contains_key(&blocks[0].recovered_block().hash()));
537        assert!(!tree_state.parent_to_child.contains_key(&blocks[1].recovered_block().hash()));
538        assert!(tree_state.parent_to_child.contains_key(&blocks[2].recovered_block().hash()));
539        assert!(tree_state.parent_to_child.contains_key(&blocks[3].recovered_block().hash()));
540        assert!(!tree_state.parent_to_child.contains_key(&blocks[4].recovered_block().hash()));
541
542        assert_eq!(
543            tree_state.parent_to_child.get(&blocks[2].recovered_block().hash()),
544            Some(&HashSet::from_iter([blocks[3].recovered_block().hash()]))
545        );
546        assert_eq!(
547            tree_state.parent_to_child.get(&blocks[3].recovered_block().hash()),
548            Some(&HashSet::from_iter([blocks[4].recovered_block().hash()]))
549        );
550    }
551
552    #[tokio::test]
553    async fn test_tree_state_remove_before_finalized() {
554        let start_num_hash = BlockNumHash::default();
555        let mut tree_state = TreeState::new(start_num_hash, EngineApiKind::Ethereum);
556        let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..6).collect();
557
558        for block in &blocks {
559            tree_state.insert_executed(block.clone());
560        }
561
562        let last = blocks.last().unwrap();
563
564        // set the canonical head
565        tree_state.set_canonical_head(last.recovered_block().num_hash());
566
567        // we should still remove everything up to and including 2
568        tree_state.remove_until(
569            BlockNumHash::new(2, blocks[1].recovered_block().hash()),
570            start_num_hash.hash,
571            None,
572        );
573
574        assert!(!tree_state.blocks_by_hash.contains_key(&blocks[0].recovered_block().hash()));
575        assert!(!tree_state.blocks_by_hash.contains_key(&blocks[1].recovered_block().hash()));
576        assert!(!tree_state.blocks_by_number.contains_key(&1));
577        assert!(!tree_state.blocks_by_number.contains_key(&2));
578
579        assert!(tree_state.blocks_by_hash.contains_key(&blocks[2].recovered_block().hash()));
580        assert!(tree_state.blocks_by_hash.contains_key(&blocks[3].recovered_block().hash()));
581        assert!(tree_state.blocks_by_hash.contains_key(&blocks[4].recovered_block().hash()));
582        assert!(tree_state.blocks_by_number.contains_key(&3));
583        assert!(tree_state.blocks_by_number.contains_key(&4));
584        assert!(tree_state.blocks_by_number.contains_key(&5));
585
586        assert!(!tree_state.parent_to_child.contains_key(&blocks[0].recovered_block().hash()));
587        assert!(!tree_state.parent_to_child.contains_key(&blocks[1].recovered_block().hash()));
588        assert!(tree_state.parent_to_child.contains_key(&blocks[2].recovered_block().hash()));
589        assert!(tree_state.parent_to_child.contains_key(&blocks[3].recovered_block().hash()));
590        assert!(!tree_state.parent_to_child.contains_key(&blocks[4].recovered_block().hash()));
591
592        assert_eq!(
593            tree_state.parent_to_child.get(&blocks[2].recovered_block().hash()),
594            Some(&HashSet::from_iter([blocks[3].recovered_block().hash()]))
595        );
596        assert_eq!(
597            tree_state.parent_to_child.get(&blocks[3].recovered_block().hash()),
598            Some(&HashSet::from_iter([blocks[4].recovered_block().hash()]))
599        );
600    }
601
602    #[tokio::test]
603    async fn test_tree_state_remove_before_lower_finalized() {
604        let start_num_hash = BlockNumHash::default();
605        let mut tree_state = TreeState::new(start_num_hash, EngineApiKind::Ethereum);
606        let blocks: Vec<_> = TestBlockBuilder::eth().get_executed_blocks(1..6).collect();
607
608        for block in &blocks {
609            tree_state.insert_executed(block.clone());
610        }
611
612        let last = blocks.last().unwrap();
613
614        // set the canonical head
615        tree_state.set_canonical_head(last.recovered_block().num_hash());
616
617        // we have no forks so we should still remove anything up to and including 2
618        tree_state.remove_until(
619            BlockNumHash::new(2, blocks[1].recovered_block().hash()),
620            start_num_hash.hash,
621            Some(blocks[0].recovered_block().num_hash()),
622        );
623
624        assert!(!tree_state.blocks_by_hash.contains_key(&blocks[0].recovered_block().hash()));
625        assert!(!tree_state.blocks_by_hash.contains_key(&blocks[1].recovered_block().hash()));
626        assert!(!tree_state.blocks_by_number.contains_key(&1));
627        assert!(!tree_state.blocks_by_number.contains_key(&2));
628
629        assert!(tree_state.blocks_by_hash.contains_key(&blocks[2].recovered_block().hash()));
630        assert!(tree_state.blocks_by_hash.contains_key(&blocks[3].recovered_block().hash()));
631        assert!(tree_state.blocks_by_hash.contains_key(&blocks[4].recovered_block().hash()));
632        assert!(tree_state.blocks_by_number.contains_key(&3));
633        assert!(tree_state.blocks_by_number.contains_key(&4));
634        assert!(tree_state.blocks_by_number.contains_key(&5));
635
636        assert!(!tree_state.parent_to_child.contains_key(&blocks[0].recovered_block().hash()));
637        assert!(!tree_state.parent_to_child.contains_key(&blocks[1].recovered_block().hash()));
638        assert!(tree_state.parent_to_child.contains_key(&blocks[2].recovered_block().hash()));
639        assert!(tree_state.parent_to_child.contains_key(&blocks[3].recovered_block().hash()));
640        assert!(!tree_state.parent_to_child.contains_key(&blocks[4].recovered_block().hash()));
641
642        assert_eq!(
643            tree_state.parent_to_child.get(&blocks[2].recovered_block().hash()),
644            Some(&HashSet::from_iter([blocks[3].recovered_block().hash()]))
645        );
646        assert_eq!(
647            tree_state.parent_to_child.get(&blocks[3].recovered_block().hash()),
648            Some(&HashSet::from_iter([blocks[4].recovered_block().hash()]))
649        );
650    }
651}