Track pruned subtrees in repair weight

core/src/heaviest_subtree_fork_choice.rs

@@ -16,8 +16,10 @@ use {
     },
     std::{
         borrow::Borrow,
+        cmp::Reverse,
         collections::{
-            btree_set::Iter, hash_map::Entry, BTreeMap, BTreeSet, HashMap, HashSet, VecDeque,
+            btree_set::Iter, hash_map::Entry, BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet,
+            VecDeque,
         },
         sync::{Arc, RwLock},
         time::Instant,
@@ -309,6 +311,40 @@ impl HeaviestSubtreeForkChoice {
         self.last_root_time = Instant::now();
     }
 
+    /// Prunes all slots not descending from `new_root`.
+    /// Returns slots that were removed (`< new_root`) and the newly pruned subtrees with only the slots > `new_root`
+    pub fn prune(&mut self, new_root: SlotHashKey) -> (Vec<SlotHashKey>, Vec<Self>) {
+        let mut pruned = vec![];
+        let mut removed = vec![];
+        // Find the pruned trees
+        let mut to_visit = BinaryHeap::from([Reverse(self.tree_root)]);
+        while !to_visit.is_empty() {
+            let cur_slot = to_visit.pop().unwrap().0;
+            if cur_slot == new_root {
+                continue;
+            }
+            if cur_slot < new_root {
+                for child in (&*self)
+                    .children(&cur_slot)
+                    .expect("slot was discovered earlier, must exist")
+                {
+                    to_visit.push(Reverse(*child))
+                }
+                removed.push(cur_slot);
+            } else {
+                // The start of a pruned subtree. Split it out and stop traversing this subtree.
+                pruned.push(self.split_off(&cur_slot));
+            }
+        }
+
+        for slot in removed.iter() {
+            self.fork_infos
+                .remove(slot)
+                .expect("Slots reachable from old root must exist in tree");
+        }
+        (removed, pruned)
+    }
+
     pub fn add_root_parent(&mut self, root_parent: SlotHashKey) {
         assert!(root_parent.0 < self.tree_root.0);
         assert!(self.fork_infos.get(&root_parent).is_none());