From 3e14af503304969209237c2af153ba673600db06 Mon Sep 17 00:00:00 2001
From: anatoly yakovenko <aeyakovenko@gmail.com>
Date: Sun, 28 Apr 2019 10:27:09 -0700
Subject: [PATCH] bast bank ancestor check (#4050)

---
 core/src/bank_forks.rs | 12 ++++++++----
 runtime/src/bank.rs    |  2 +-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/core/src/bank_forks.rs b/core/src/bank_forks.rs
index c15625b9d8..6a1114d92a 100644
--- a/core/src/bank_forks.rs
+++ b/core/src/bank_forks.rs
@@ -35,7 +35,8 @@ impl BankForks {
     pub fn ancestors(&self) -> HashMap<u64, HashSet<u64>> {
         let mut ancestors = HashMap::new();
         for bank in self.banks.values() {
-            let set = bank.parents().into_iter().map(|b| b.slot()).collect();
+            let mut set: HashSet<u64> = bank.ancestors.keys().cloned().collect();
+            set.remove(&bank.slot());
             ancestors.insert(bank.slot(), set);
         }
         ancestors
@@ -46,9 +47,11 @@ impl BankForks {
         let mut descendants = HashMap::new();
         for bank in self.banks.values() {
             let _ = descendants.entry(bank.slot()).or_insert(HashSet::new());
-            for parent in bank.parents() {
+            let mut set: HashSet<u64> = bank.ancestors.keys().cloned().collect();
+            set.remove(&bank.slot());
+            for parent in set {
                 descendants
-                    .entry(parent.slot())
+                    .entry(parent)
                     .or_insert(HashSet::new())
                     .insert(bank.slot());
             }
@@ -116,8 +119,9 @@ impl BankForks {
     }
 
     fn prune_non_root(&mut self, root: u64) {
+        let descendants = self.descendants();
         self.banks
-            .retain(|slot, bank| *slot >= root || bank.is_in_subtree_of(root))
+            .retain(|slot, _| descendants[&root].contains(slot))
     }
 }
 
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 4e770481d2..101355342a 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -124,7 +124,7 @@ pub struct Bank {
     parent: RwLock<Option<Arc<Bank>>>,
 
     /// The set of parents including this bank
-    ancestors: HashMap<u64, usize>,
+    pub ancestors: HashMap<u64, usize>,
 
     /// Hash of this Bank's state. Only meaningful after freezing.
     hash: RwLock<Hash>,