From e9cbdf711bbf567cde8d0854548e212e103c3a5c Mon Sep 17 00:00:00 2001
From: carllin <wumu727@gmail.com>
Date: Tue, 14 Jul 2020 00:38:48 -0700
Subject: [PATCH] Add TreeDiff trait to reuse tree functions (#11046)

Co-authored-by: Carl <carl@solana.com>
---
 core/src/heaviest_subtree_fork_choice.rs | 47 +++++++-----------------
 core/src/lib.rs                          |  1 +
 core/src/repair_weight.rs                |  1 +
 core/src/repair_weighted_traversal.rs    |  2 +-
 core/src/tree_diff.rs                    | 32 ++++++++++++++++
 5 files changed, 48 insertions(+), 35 deletions(-)
 create mode 100644 core/src/tree_diff.rs

diff --git a/core/src/heaviest_subtree_fork_choice.rs b/core/src/heaviest_subtree_fork_choice.rs
index 25e39ac2a..ae2cd4dbd 100644
--- a/core/src/heaviest_subtree_fork_choice.rs
+++ b/core/src/heaviest_subtree_fork_choice.rs
@@ -2,6 +2,7 @@ use crate::{
     consensus::{ComputedBankState, Tower},
     fork_choice::ForkChoice,
     progress_map::ProgressMap,
+    tree_diff::TreeDiff,
 };
 use solana_runtime::{bank::Bank, bank_forks::BankForks, epoch_stakes::EpochStakes};
 use solana_sdk::{
@@ -141,10 +142,6 @@ impl HeaviestSubtreeForkChoice {
             .map(|fork_info| fork_info.stake_voted_subtree)
     }
 
-    pub fn contains_slot(&self, slot: Slot) -> bool {
-        self.fork_infos.contains_key(&slot)
-    }
-
     pub fn root(&self) -> Slot {
         self.root
     }
@@ -248,30 +245,6 @@ impl HeaviestSubtreeForkChoice {
         self.propagate_new_leaf(slot, parent)
     }
 
-    // Find all nodes reachable from `root1`, excluding subtree at `root2`
-    pub fn subtree_diff(&self, root1: Slot, root2: Slot) -> HashSet<Slot> {
-        if !self.contains_slot(root1) {
-            return HashSet::new();
-        }
-        let mut pending_slots = vec![root1];
-        let mut reachable_set = HashSet::new();
-        while !pending_slots.is_empty() {
-            let current_slot = pending_slots.pop().unwrap();
-            if current_slot == root2 {
-                continue;
-            }
-            reachable_set.insert(current_slot);
-            for child in self
-                .children(current_slot)
-                .expect("slot was discovered earlier, must exist")
-            {
-                pending_slots.push(*child);
-            }
-        }
-
-        reachable_set
-    }
-
     // Returns if the given `maybe_best_child` is the heaviest among the children
     // it's parent
     fn is_best_child(&self, maybe_best_child: Slot) -> bool {
@@ -305,12 +278,6 @@ impl HeaviestSubtreeForkChoice {
         AncestorIterator::new(start_slot, &self.fork_infos).collect()
     }
 
-    pub fn children(&self, slot: Slot) -> Option<&[Slot]> {
-        self.fork_infos
-            .get(&slot)
-            .map(|fork_info| &fork_info.children[..])
-    }
-
     pub fn merge(
         &mut self,
         other: HeaviestSubtreeForkChoice,
@@ -542,6 +509,18 @@ impl HeaviestSubtreeForkChoice {
     }
 }
 
+impl TreeDiff for HeaviestSubtreeForkChoice {
+    fn contains_slot(&self, slot: Slot) -> bool {
+        self.fork_infos.contains_key(&slot)
+    }
+
+    fn children(&self, slot: Slot) -> Option<&[Slot]> {
+        self.fork_infos
+            .get(&slot)
+            .map(|fork_info| &fork_info.children[..])
+    }
+}
+
 impl ForkChoice for HeaviestSubtreeForkChoice {
     fn compute_bank_stats(
         &mut self,
diff --git a/core/src/lib.rs b/core/src/lib.rs
index 758a5325b..bd3a21d8d 100644
--- a/core/src/lib.rs
+++ b/core/src/lib.rs
@@ -62,6 +62,7 @@ pub mod sigverify_stage;
 pub mod snapshot_packager_service;
 pub mod tpu;
 pub mod transaction_status_service;
+pub mod tree_diff;
 pub mod tvu;
 pub mod validator;
 pub mod verified_vote_packets;
diff --git a/core/src/repair_weight.rs b/core/src/repair_weight.rs
index 639a5b463..d230abc72 100644
--- a/core/src/repair_weight.rs
+++ b/core/src/repair_weight.rs
@@ -3,6 +3,7 @@ use crate::{
     repair_service::RepairTiming,
     repair_weighted_traversal::{self, Contains},
     serve_repair::RepairType,
+    tree_diff::TreeDiff,
 };
 use solana_ledger::{ancestor_iterator::AncestorIterator, blockstore::Blockstore};
 use solana_measure::measure::Measure;
diff --git a/core/src/repair_weighted_traversal.rs b/core/src/repair_weighted_traversal.rs
index 24b5cd8c8..667230100 100644
--- a/core/src/repair_weighted_traversal.rs
+++ b/core/src/repair_weighted_traversal.rs
@@ -1,6 +1,6 @@
 use crate::{
     heaviest_subtree_fork_choice::HeaviestSubtreeForkChoice, repair_service::RepairService,
-    serve_repair::RepairType,
+    serve_repair::RepairType, tree_diff::TreeDiff,
 };
 use solana_ledger::blockstore::Blockstore;
 use solana_sdk::clock::Slot;
diff --git a/core/src/tree_diff.rs b/core/src/tree_diff.rs
new file mode 100644
index 000000000..bad6d3143
--- /dev/null
+++ b/core/src/tree_diff.rs
@@ -0,0 +1,32 @@
+use solana_sdk::clock::Slot;
+use std::collections::HashSet;
+
+pub trait TreeDiff {
+    fn children(&self, slot: Slot) -> Option<&[Slot]>;
+
+    fn contains_slot(&self, slot: Slot) -> bool;
+
+    // Find all nodes reachable from `root1`, excluding subtree at `root2`
+    fn subtree_diff(&self, root1: Slot, root2: Slot) -> HashSet<Slot> {
+        if !self.contains_slot(root1) {
+            return HashSet::new();
+        }
+        let mut pending_slots = vec![root1];
+        let mut reachable_set = HashSet::new();
+        while !pending_slots.is_empty() {
+            let current_slot = pending_slots.pop().unwrap();
+            if current_slot == root2 {
+                continue;
+            }
+            reachable_set.insert(current_slot);
+            for child in self
+                .children(current_slot)
+                .expect("slot was discovered earlier, must exist")
+            {
+                pending_slots.push(*child);
+            }
+        }
+
+        reachable_set
+    }
+}