Refactor Vote Program Account setup (#2992)

src/leader_schedule_utils.rs

@@ -5,13 +5,30 @@ use solana_sdk::pubkey::Pubkey;
 
 /// Return the leader schedule for the given epoch.
 fn leader_schedule(epoch_height: u64, bank: &Bank) -> LeaderSchedule {
-    let stakes = staking_utils::staked_nodes_at_epoch(bank, epoch_height);
+    let stakes = staking_utils::node_stakes_at_epoch(bank, epoch_height);
     let mut seed = [0u8; 32];
     seed[0..8].copy_from_slice(&epoch_height.to_le_bytes());
-    let stakes: Vec<_> = stakes.into_iter().collect();
+    let mut stakes: Vec<_> = stakes.into_iter().collect();
+    sort_stakes(&mut stakes);
     LeaderSchedule::new(&stakes, seed, bank.slots_per_epoch())
 }
 
+fn sort_stakes(stakes: &mut Vec<(Pubkey, u64)>) {
+    // Sort first by stake. If stakes are the same, sort 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(|(l_id, l_stake), (r_id, r_stake)| {
+        if r_stake == l_stake {
+            r_id.cmp(&l_id)
+        } else {
+            r_stake.cmp(&l_stake)
+        }
+    });
+
+    // Now that it's sorted, we can do an O(n) dedup.
+    stakes.dedup();
+}
+
 /// Return the leader for the slot at the slot_index and epoch_height returned
 /// by the given function.
 fn slot_leader_by<F>(bank: &Bank, get_slot_index: F) -> Pubkey
@@ -92,7 +109,7 @@ mod tests {
         let (genesis_block, _mint_keypair) = GenesisBlock::new_with_leader(2, pubkey, 2);
         let bank = Bank::new(&genesis_block);
 
-        let ids_and_stakes: Vec<_> = staking_utils::staked_nodes(&bank).into_iter().collect();
+        let ids_and_stakes: Vec<_> = staking_utils::node_stakes(&bank).into_iter().collect();
         let seed = [0u8; 32];
         let leader_schedule =
             LeaderSchedule::new(&ids_and_stakes, seed, genesis_block.slots_per_epoch);
@@ -122,4 +139,31 @@ mod tests {
         assert_eq!(next_slot_leader_index(0, 0, 2), (1, 0));
         assert_eq!(next_slot_leader_index(1, 0, 2), (0, 1));
     }
+
+    #[test]
+    fn test_sort_stakes_basic() {
+        let pubkey0 = Keypair::new().pubkey();
+        let pubkey1 = Keypair::new().pubkey();
+        let mut stakes = vec![(pubkey0, 1), (pubkey1, 2)];
+        sort_stakes(&mut stakes);
+        assert_eq!(stakes, vec![(pubkey1, 2), (pubkey0, 1)]);
+    }
+
+    #[test]
+    fn test_sort_stakes_with_dup() {
+        let pubkey0 = Keypair::new().pubkey();
+        let pubkey1 = Keypair::new().pubkey();
+        let mut stakes = vec![(pubkey0, 1), (pubkey1, 2), (pubkey0, 1)];
+        sort_stakes(&mut stakes);
+        assert_eq!(stakes, vec![(pubkey1, 2), (pubkey0, 1)]);
+    }
+
+    #[test]
+    fn test_sort_stakes_with_equal_stakes() {
+        let pubkey0 = Pubkey::default();
+        let pubkey1 = Keypair::new().pubkey();
+        let mut stakes = vec![(pubkey0, 1), (pubkey1, 1)];
+        sort_stakes(&mut stakes);
+        assert_eq!(stakes, vec![(pubkey1, 1), (pubkey0, 1)]);
+    }
 }