fix: remove persisted blocks properly (#9969)

crates/engine/tree/src/tree/mod.rs

@@ -42,6 +42,7 @@ use reth_stages_api::ControlFlow;
 use reth_trie::HashedPostState;
 use std::{
     collections::{BTreeMap, HashMap, HashSet},
+    ops::Bound,
     sync::{mpsc::Receiver, Arc},
     time::Instant,
 };
@@ -63,18 +64,20 @@ pub use config::TreeConfig;
 /// - All executed blocks are valid and have been executed.
 #[derive(Debug, Default)]
 pub struct TreeState {
-    /// __All__ executed blocks by block hash.
+    /// __All__ unique executed blocks by block hash that are connected to the canonical chain.
     ///
     /// This includes blocks of all forks.
     blocks_by_hash: HashMap<B256, ExecutedBlock>,
     /// Executed blocks grouped by their respective block number.
     ///
     /// This maps unique block number to all known blocks for that height.
+    ///
+    /// Note: there can be multiple blocks at the same height due to forks.
     blocks_by_number: BTreeMap<BlockNumber, Vec<ExecutedBlock>>,
-    /// Currently tracked canonical head of the chain.
-    current_canonical_head: BlockNumHash,
     /// Map of any parent block hash to its children.
     parent_to_child: HashMap<B256, HashSet<B256>>,
+    /// Currently tracked canonical head of the chain.
+    current_canonical_head: BlockNumHash,
 }
 
 impl TreeState {
@@ -127,10 +130,10 @@ impl TreeState {
         }
     }
 
-    /// Remove blocks before specified block number (exclusive).
-    pub(crate) fn remove_before(&mut self, block_number: BlockNumber) {
+    /// Remove all blocks up to the given block number.
+    pub(crate) fn remove_before(&mut self, upper_bound: Bound<BlockNumber>) {
         let mut numbers_to_remove = Vec::new();
-        for (&number, _) in self.blocks_by_number.range(..block_number) {
+        for (&number, _) in self.blocks_by_number.range((Bound::Unbounded, upper_bound)) {
             numbers_to_remove.push(number);
         }
 
@@ -629,7 +632,7 @@ where
 
         // state house keeping after backfill sync
         // remove all executed blocks below the backfill height
-        self.state.tree_state.remove_before(backfill_height + 1);
+        self.state.tree_state.remove_before(Bound::Included(backfill_height));
         // remove all buffered blocks below the backfill height
         self.state.buffer.remove_old_blocks(backfill_height);
         // we remove all entries because now we're synced to the backfill target and consider this
@@ -799,33 +802,13 @@ where
     ///
     /// Assumes that `finish` has been called on the `persistence_state` at least once
     fn on_new_persisted_block(&mut self) {
-        self.remove_persisted_blocks_from_tree_state();
+        self.state
+            .tree_state
+            .remove_before(Bound::Included(self.persistence_state.last_persisted_block_number));
         self.canonical_in_memory_state
             .remove_persisted_blocks(self.persistence_state.last_persisted_block_number);
     }
 
-    /// Clears persisted blocks from the in-memory tree state.
-    ///
-    /// Assumes that `finish` has been called on the `persistence_state` at least once
-    fn remove_persisted_blocks_from_tree_state(&mut self) {
-        let keys_to_remove: Vec<BlockNumber> = self
-            .state
-            .tree_state
-            .blocks_by_number
-            .range(..=self.persistence_state.last_persisted_block_number)
-            .map(|(&k, _)| k)
-            .collect();
-
-        for key in keys_to_remove {
-            if let Some(blocks) = self.state.tree_state.blocks_by_number.remove(&key) {
-                // Remove corresponding blocks from blocks_by_hash
-                for block in blocks {
-                    self.state.tree_state.blocks_by_hash.remove(&block.block().hash());
-                }
-            }
-        }
-    }
-
     /// Return block from database or in-memory state by hash.
     fn block_by_hash(&self, hash: B256) -> ProviderResult<Option<Block>> {
         // check database first
@@ -2096,7 +2079,7 @@ mod tests {
             tree_state.insert_executed(block.clone());
         }
 
-        tree_state.remove_before(3);
+        tree_state.remove_before(Bound::Excluded(3));
 
         assert!(!tree_state.blocks_by_hash.contains_key(&blocks[0].block.hash()));
         assert!(!tree_state.blocks_by_hash.contains_key(&blocks[1].block.hash()));