Merge heaviest bank modules (#10672)

* Merge heaviest bank modules

* Merge heaviest fork choice modules

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

core/src/bank_weight_fork_choice.rs

@@ -20,7 +20,7 @@ impl ForkChoice for BankWeightForkChoice {
         bank: &Bank,
         _tower: &Tower,
         progress: &mut ProgressMap,
-        computed_bank_stats: &ComputedBankState,
+        computed_bank_state: &ComputedBankState,
     ) {
         let bank_slot = bank.slot();
         // Only time progress map should be missing a bank slot
@@ -36,14 +36,14 @@ impl ForkChoice for BankWeightForkChoice {
             .get_fork_stats_mut(bank_slot)
             .expect("All frozen banks must exist in the Progress map");
 
-        let ComputedBankState { bank_weight, .. } = computed_bank_stats;
+        let ComputedBankState { bank_weight, .. } = computed_bank_state;
         stats.weight = *bank_weight;
         stats.fork_weight = stats.weight + parent_weight;
     }
 
     // Returns:
     // 1) The heaviest overall bank
-    // 2) The heavest bank on the same fork as the last vote (doesn't require a
+    // 2) The heaviest bank on the same fork as the last vote (doesn't require a
     // switching proof to vote for)
     fn select_forks(
         &self,

core/src/fork_choice.rs

@@ -21,12 +21,12 @@ pub(crate) trait ForkChoice {
         bank: &Bank,
         tower: &Tower,
         progress: &mut ProgressMap,
-        computed_bank_stats: &ComputedBankState,
+        computed_bank_state: &ComputedBankState,
     );
 
     // Returns:
-    // 1) The heaviest overall bbank
-    // 2) The heavest bank on the same fork as the last vote (doesn't require a
+    // 1) The heaviest overall bank
+    // 2) The heaviest bank on the same fork as the last vote (doesn't require a
     // switching proof to vote for)
     fn select_forks(
         &self,

core/src/heaviest_subtree_fork_choice.rs

@@ -1,19 +1,21 @@
-#[cfg(test)]
-use solana_runtime::bank_forks::BankForks;
-use solana_runtime::{bank::Bank, epoch_stakes::EpochStakes};
+use crate::{
+    consensus::{ComputedBankState, Tower},
+    fork_choice::ForkChoice,
+    progress_map::ProgressMap,
+};
+use solana_runtime::{bank::Bank, bank_forks::BankForks, epoch_stakes::EpochStakes};
 use solana_sdk::{
     clock::{Epoch, Slot},
     epoch_schedule::EpochSchedule,
     pubkey::Pubkey,
 };
 use std::{
-    collections::{BTreeMap, HashMap},
-    sync::Arc,
+    collections::{BTreeMap, HashMap, HashSet},
+    sync::{Arc, RwLock},
 };
 #[cfg(test)]
 use trees::{Tree, TreeWalk};
 
-mod fork_choice;
 pub type ForkWeight = u64;
 
 #[derive(PartialEq, Eq, Clone, Debug, PartialOrd, Ord)]
@@ -401,6 +403,66 @@ impl HeaviestSubtreeForkChoice {
     }
 }
 
+impl ForkChoice for HeaviestSubtreeForkChoice {
+    fn compute_bank_stats(
+        &mut self,
+        bank: &Bank,
+        _tower: &Tower,
+        _progress: &mut ProgressMap,
+        computed_bank_state: &ComputedBankState,
+    ) {
+        let ComputedBankState { pubkey_votes, .. } = computed_bank_state;
+
+        // Update `heaviest_subtree_fork_choice` to find the best fork to build on
+        let best_overall_slot = self.add_votes(
+            &pubkey_votes,
+            bank.epoch_stakes_map(),
+            bank.epoch_schedule(),
+        );
+
+        datapoint_info!(
+            "best_slot",
+            ("slot", bank.slot(), i64),
+            ("best_slot", best_overall_slot, i64),
+        );
+    }
+
+    // Returns:
+    // 1) The heaviest overall bank
+    // 2) The heaviest bank on the same fork as the last vote (doesn't require a
+    // switching proof to vote for)
+    fn select_forks(
+        &self,
+        _frozen_banks: &[Arc<Bank>],
+        tower: &Tower,
+        _progress: &ProgressMap,
+        _ancestors: &HashMap<u64, HashSet<u64>>,
+        bank_forks: &RwLock<BankForks>,
+    ) -> (Arc<Bank>, Option<Arc<Bank>>) {
+        let last_vote = tower.last_vote().slots.last().cloned();
+        let heaviest_slot_on_same_voted_fork = last_vote.map(|last_vote| {
+            let heaviest_slot_on_same_voted_fork =
+                self.best_slot(last_vote).expect("last_vote is a frozen bank so must have been added to heaviest_subtree_fork_choice at time of freezing");
+            if heaviest_slot_on_same_voted_fork == last_vote {
+                None
+            } else {
+                Some(heaviest_slot_on_same_voted_fork)
+            }
+        }).unwrap_or(None);
+        let heaviest_slot = self.best_overall_slot();
+        let r_bank_forks = bank_forks.read().unwrap();
+        (
+            r_bank_forks.get(heaviest_slot).unwrap().clone(),
+            heaviest_slot_on_same_voted_fork.map(|heaviest_slot_on_same_voted_fork| {
+                r_bank_forks
+                    .get(heaviest_slot_on_same_voted_fork)
+                    .unwrap()
+                    .clone()
+            }),
+        )
+    }
+}
+
 struct AncestorIterator<'a> {
     current_slot: Slot,
     fork_infos: &'a HashMap<Slot, ForkInfo>,

core/src/heaviest_subtree_fork_choice/fork_choice.rs

@@ -1,71 +0,0 @@
-use crate::{
-    consensus::{ComputedBankState, Tower},
-    fork_choice::ForkChoice,
-    heaviest_subtree_fork_choice::HeaviestSubtreeForkChoice,
-    progress_map::ProgressMap,
-};
-use solana_runtime::{bank::Bank, bank_forks::BankForks};
-use std::{
-    collections::{HashMap, HashSet},
-    sync::{Arc, RwLock},
-};
-
-impl ForkChoice for HeaviestSubtreeForkChoice {
-    fn compute_bank_stats(
-        &mut self,
-        bank: &Bank,
-        _tower: &Tower,
-        _progress: &mut ProgressMap,
-        computed_bank_stats: &ComputedBankState,
-    ) {
-        let ComputedBankState { pubkey_votes, .. } = computed_bank_stats;
-
-        // Update `heaviest_subtree_fork_choice` to find the best fork to build on
-        let best_overall_slot = self.add_votes(
-            &pubkey_votes,
-            bank.epoch_stakes_map(),
-            bank.epoch_schedule(),
-        );
-
-        datapoint_info!(
-            "best_slot",
-            ("slot", bank.slot(), i64),
-            ("best_slot", best_overall_slot, i64),
-        );
-    }
-
-    // Returns:
-    // 1) The heaviest overall bbank
-    // 2) The heavest bank on the same fork as the last vote (doesn't require a
-    // switching proof to vote for)
-    fn select_forks(
-        &self,
-        _frozen_banks: &[Arc<Bank>],
-        tower: &Tower,
-        _progress: &ProgressMap,
-        _ancestors: &HashMap<u64, HashSet<u64>>,
-        bank_forks: &RwLock<BankForks>,
-    ) -> (Arc<Bank>, Option<Arc<Bank>>) {
-        let last_vote = tower.last_vote().slots.last().cloned();
-        let heaviest_slot_on_same_voted_fork = last_vote.map(|last_vote| {
-            let heaviest_slot_on_same_voted_fork =
-                self.best_slot(last_vote).expect("last_vote is a frozen bank so must have been added to heaviest_subtree_fork_choice at time of freezing");
-            if heaviest_slot_on_same_voted_fork == last_vote {
-                None
-            } else {
-                Some(heaviest_slot_on_same_voted_fork)
-            }
-        }).unwrap_or(None);
-        let heaviest_slot = self.best_overall_slot();
-        let r_bank_forks = bank_forks.read().unwrap();
-        (
-            r_bank_forks.get(heaviest_slot).unwrap().clone(),
-            heaviest_slot_on_same_voted_fork.map(|heaviest_slot_on_same_voted_fork| {
-                r_bank_forks
-                    .get(heaviest_slot_on_same_voted_fork)
-                    .unwrap()
-                    .clone()
-            }),
-        )
-    }
-}