reth_engine_tree/tree/
state.rs

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

reth_engine_tree/tree/
state.rs
