PR comments

core/src/bank_forks.rs

@@ -84,11 +84,9 @@ impl BankForks {
         }
     }
 
-    // TODO: use the bank's own ID instead of receiving a parameter?
+    pub fn insert(&mut self, bank: Bank) {
-    pub fn insert(&mut self, bank_slot: u64, bank: Bank) {
         let bank = Arc::new(bank);
-        assert_eq!(bank_slot, bank.slot());
+        let prev = self.banks.insert(bank.slot(), bank.clone());
-        let prev = self.banks.insert(bank_slot, bank.clone());
         assert!(prev.is_none());
 
         self.working_bank = bank.clone();
@@ -109,13 +107,8 @@ impl BankForks {
     }
 
     fn prune_non_root(&mut self, root: u64) {
-        self.banks.retain(|slot, bank| {
+        self.banks
-            if *slot < root {
+            .retain(|slot, bank| *slot >= root || bank.is_in_subtree_of(root))
-                false
-            } else {
-                bank.is_descendant_of(root)
-            }
-        })
     }
 }
 
@@ -133,7 +126,7 @@ mod tests {
         let mut bank_forks = BankForks::new(0, bank);
         let child_bank = Bank::new_from_parent(&bank_forks[0u64], &Pubkey::default(), 1);
         child_bank.register_tick(&Hash::default());
-        bank_forks.insert(1, child_bank);
+        bank_forks.insert(child_bank);
         assert_eq!(bank_forks[1u64].tick_height(), 1);
         assert_eq!(bank_forks.working_bank().tick_height(), 1);
     }
@@ -145,9 +138,9 @@ mod tests {
         let mut bank_forks = BankForks::new(0, bank);
         let bank0 = bank_forks[0].clone();
         let bank = Bank::new_from_parent(&bank0, &Pubkey::default(), 1);
-        bank_forks.insert(1, bank);
+        bank_forks.insert(bank);
         let bank = Bank::new_from_parent(&bank0, &Pubkey::default(), 2);
-        bank_forks.insert(2, bank);
+        bank_forks.insert(bank);
         let descendants = bank_forks.descendants();
         let children: Vec<u64> = descendants[&0].iter().cloned().collect();
         assert_eq!(children, vec![1, 2]);
@@ -162,9 +155,9 @@ mod tests {
         let mut bank_forks = BankForks::new(0, bank);
         let bank0 = bank_forks[0].clone();
         let bank = Bank::new_from_parent(&bank0, &Pubkey::default(), 1);
-        bank_forks.insert(1, bank);
+        bank_forks.insert(bank);
         let bank = Bank::new_from_parent(&bank0, &Pubkey::default(), 2);
-        bank_forks.insert(2, bank);
+        bank_forks.insert(bank);
         let ancestors = bank_forks.ancestors();
         assert!(ancestors[&0].is_empty());
         let parents: Vec<u64> = ancestors[&1].iter().cloned().collect();
@@ -179,7 +172,7 @@ mod tests {
         let bank = Bank::new(&genesis_block);
         let mut bank_forks = BankForks::new(0, bank);
         let child_bank = Bank::new_from_parent(&bank_forks[0u64], &Pubkey::default(), 1);
-        bank_forks.insert(1, child_bank);
+        bank_forks.insert(child_bank);
         assert!(bank_forks.frozen_banks().get(&0).is_some());
         assert!(bank_forks.frozen_banks().get(&1).is_none());
     }
@@ -190,7 +183,7 @@ mod tests {
         let bank = Bank::new(&genesis_block);
         let mut bank_forks = BankForks::new(0, bank);
         let child_bank = Bank::new_from_parent(&bank_forks[0u64], &Pubkey::default(), 1);
-        bank_forks.insert(1, child_bank);
+        bank_forks.insert(child_bank);
         assert_eq!(bank_forks.active_banks(), vec![1]);
     }
 

core/src/replay_stage.rs

@@ -283,7 +283,7 @@ impl ReplayStage {
                                 )
                                 .to_owned(),);
                         let tpu_bank = Bank::new_from_parent(parent, my_id, poh_slot);
-                        bank_forks.write().unwrap().insert(poh_slot, tpu_bank);
+                        bank_forks.write().unwrap().insert(tpu_bank);
                         if let Some(tpu_bank) = bank_forks.read().unwrap().get(poh_slot).cloned() {
                             assert_eq!(
                                 bank_forks.read().unwrap().working_bank().slot(),
@@ -415,10 +415,7 @@ impl ReplayStage {
                 }
                 let leader = leader_schedule_utils::slot_leader_at(child_id, &parent_bank).unwrap();
                 info!("new fork:{} parent:{}", child_id, parent_id);
-                forks.insert(
+                forks.insert(Bank::new_from_parent(&parent_bank, &leader, child_id));
-                    child_id,
-                    Bank::new_from_parent(&parent_bank, &leader, child_id),
-                );
             }
         }
     }
@@ -599,7 +596,7 @@ mod test {
             ReplayStage::generate_new_bank_forks(&blocktree, &mut bank_forks);
             assert!(bank_forks.get(1).is_some());
 
-            // Inset blob for slot 3, generate new forks, check result
+            // Insert blob for slot 3, generate new forks, check result
             let mut blob_slot_2 = Blob::default();
             blob_slot_2.set_slot(2);
             blob_slot_2.set_parent(0);

runtime/src/bank.rs

@@ -860,7 +860,7 @@ impl Bank {
             .store_slow(self.accounts_id, &program_id, &bogus_account);
     }
 
-    pub fn is_descendant_of(&self, parent: u64) -> bool {
+    pub fn is_in_subtree_of(&self, parent: u64) -> bool {
         if self.slot() == parent {
             return true;
         }
@@ -1644,7 +1644,7 @@ mod tests {
     }
 
     #[test]
-    fn test_is_descendant_of() {
+    fn test_is_in_subtree_of() {
         let (genesis_block, _) = GenesisBlock::new(1);
         let parent = Arc::new(Bank::new(&genesis_block));
         // Bank 1
@@ -1655,15 +1655,15 @@ mod tests {
         let bank5 = Bank::new_from_parent(&bank, &Pubkey::default(), 5);
 
         // Parents of bank 2: 0 -> 1 -> 2
-        assert!(bank2.is_descendant_of(0));
+        assert!(bank2.is_in_subtree_of(0));
-        assert!(bank2.is_descendant_of(1));
+        assert!(bank2.is_in_subtree_of(1));
-        assert!(bank2.is_descendant_of(2));
+        assert!(bank2.is_in_subtree_of(2));
-        assert!(!bank2.is_descendant_of(3));
+        assert!(!bank2.is_in_subtree_of(3));
 
-        // Parents of bank 3: 0 -> 1 -> 5
+        // Parents of bank 5: 0 -> 1 -> 5
-        assert!(bank5.is_descendant_of(0));
+        assert!(bank5.is_in_subtree_of(0));
-        assert!(bank5.is_descendant_of(1));
+        assert!(bank5.is_in_subtree_of(1));
-        assert!(!bank5.is_descendant_of(2));
+        assert!(!bank5.is_in_subtree_of(2));
-        assert!(!bank5.is_descendant_of(4));
+        assert!(!bank5.is_in_subtree_of(4));
     }
 }