From 88d6db85379b40c718fade3bafd51120d6535853 Mon Sep 17 00:00:00 2001
From: Greg Fitzgerald <greg@solana.com>
Date: Wed, 20 Feb 2019 09:21:29 -0700
Subject: [PATCH] And ranking and simplify

---
 src/active_stakers.rs   | 35 ++++++++++++++++++++++-------------
 src/leader_scheduler.rs |  3 ++-
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/active_stakers.rs b/src/active_stakers.rs
index b45f1168b1..72a257efb7 100644
--- a/src/active_stakers.rs
+++ b/src/active_stakers.rs
@@ -1,5 +1,4 @@
 use crate::leader_schedule::LeaderSchedule;
-use hashbrown::{HashMap, HashSet};
 use solana_runtime::bank::Bank;
 use solana_sdk::pubkey::Pubkey;
 use solana_sdk::vote_program::VoteState;
@@ -13,14 +12,30 @@ fn is_active_staker(vote_state: &VoteState, lower_bound: u64, upper_bound: u64)
         .is_some()
 }
 
+fn rank_stakes(stakes: &mut Vec<(Pubkey, u64)>) {
+    // Rank first by stake. If stakes are the same we rank by pubkey to ensure a
+    // deterministic result.
+    // Note: Use unstable sort, because we dedup right after to remove the equal elements.
+    stakes.sort_unstable_by(|(pubkey0, stake0), (pubkey1, stake1)| {
+        if stake0 == stake1 {
+            pubkey0.cmp(&pubkey1)
+        } else {
+            stake0.cmp(&stake1)
+        }
+    });
+
+    // Now that it's sorted, we can do an O(n) dedup.
+    stakes.dedup();
+}
+
 /// The set of stakers that have voted near the time of construction
 pub struct ActiveStakers {
-    stakes: HashMap<Pubkey, u64>,
+    stakes: Vec<(Pubkey, u64)>,
 }
 
 impl ActiveStakers {
     pub fn new_with_upper_bound(bank: &Bank, lower_bound: u64, upper_bound: u64) -> Self {
-        let stakes = bank
+        let mut stakes: Vec<_> = bank
             .vote_states(|vote_state| is_active_staker(vote_state, lower_bound, upper_bound))
             .iter()
             .filter_map(|vote_state| {
@@ -33,6 +48,7 @@ impl ActiveStakers {
                 }
             })
             .collect();
+        rank_stakes(&mut stakes);
         Self { stakes }
     }
 
@@ -40,19 +56,12 @@ impl ActiveStakers {
         Self::new_with_upper_bound(bank, lower_bound, bank.tick_height())
     }
 
-    /// Return a map from staker pubkeys to their respective stakes.
-    pub fn stakes(&self) -> HashMap<Pubkey, u64> {
-        self.stakes.clone()
-    }
-
     /// Return the pubkeys of each staker.
-    pub fn stakers(&self) -> HashSet<Pubkey> {
-        self.stakes.keys().cloned().collect()
+    pub fn pubkeys(&self) -> Vec<Pubkey> {
+        self.stakes.iter().map(|(pubkey, _stake)| *pubkey).collect()
     }
 
     pub fn leader_schedule(&self) -> LeaderSchedule {
-        let mut stakers: Vec<_> = self.stakes.keys().cloned().collect();
-        stakers.sort();
-        LeaderSchedule::new(stakers)
+        LeaderSchedule::new(self.pubkeys())
     }
 }
diff --git a/src/leader_scheduler.rs b/src/leader_scheduler.rs
index 7519576e52..092658045e 100644
--- a/src/leader_scheduler.rs
+++ b/src/leader_scheduler.rs
@@ -214,7 +214,8 @@ impl LeaderScheduler {
             upper_bound
         );
 
-        ActiveStakers::new_with_upper_bound(&bank, lower_bound, upper_bound).stakers()
+        let active_stakers = ActiveStakers::new_with_upper_bound(&bank, lower_bound, upper_bound);
+        active_stakers.pubkeys().into_iter().collect()
     }
 
     // Updates the leader schedule to include ticks from tick_height to the first tick of the next epoch