Expose the new leader schedule functionality from the bank.

runtime/Cargo.toml

@@ -15,6 +15,7 @@ fnv = "1.0.6"
 hashbrown = "0.1.8"
 log = "0.4.2"
 rand = "0.6.5"
+rand_chacha = "0.1.1"
 serde = "1.0.88"
 serde_derive = "1.0.88"
 serde_json = "1.0.38"

runtime/src/bank.rs

@@ -5,6 +5,7 @@
 
 use crate::accounts::{Accounts, ErrorCounters, InstructionAccounts, InstructionLoaders};
 use crate::last_id_queue::LastIdQueue;
+use crate::leader_schedule::LeaderSchedule;
 use crate::runtime::{self, RuntimeError};
 use crate::status_cache::StatusCache;
 use bincode::serialize;
@@ -675,7 +676,7 @@ impl Bank {
 
     /// Return the checkpointed bank that should be used to generate a leader schedule.
     /// Return None if a sufficiently old bank checkpoint doesn't exist.
-    pub fn leader_schedule_bank(&self) -> Option<Arc<Bank>> {
+    fn leader_schedule_bank(&self) -> Option<Arc<Bank>> {
         let epoch_slot_height = self.slot_height() - self.slot_index();
         let expected = epoch_slot_height.saturating_sub(self.leader_schedule_slot_offset);
         self.parents()
@@ -683,6 +684,19 @@ impl Bank {
             .find(|bank| bank.slot_height() <= expected)
     }
 
+    /// Return the leader schedule for the current epoch.
+    fn leader_schedule(&self) -> LeaderSchedule {
+        match self.leader_schedule_bank() {
+            None => LeaderSchedule::new_with_bank(self),
+            Some(bank) => LeaderSchedule::new_with_bank(&bank),
+        }
+    }
+
+    /// Return the leader for the current slot.
+    pub fn slot_leader(&self) -> Pubkey {
+        self.leader_schedule()[self.slot_index() as usize]
+    }
+
     /// Return the number of ticks since genesis.
     pub fn tick_height(&self) -> u64 {
         self.last_id_queue.read().unwrap().tick_height
@@ -1365,4 +1379,11 @@ mod tests {
             bank.merge_parents();
         }
     }
+
+    #[test]
+    fn test_bank_slot_leader_basic() {
+        let pubkey = Keypair::new().pubkey();
+        let bank = Bank::new(&GenesisBlock::new_with_leader(2, pubkey, 2).0);
+        assert_eq!(bank.slot_leader(), pubkey);
+    }
 }

runtime/src/leader_schedule.rs

@@ -0,0 +1,93 @@
+use crate::bank::Bank;
+use rand::distributions::{Distribution, WeightedIndex};
+use rand::SeedableRng;
+use rand_chacha::ChaChaRng;
+use solana_sdk::pubkey::Pubkey;
+use std::ops::Index;
+
+/// Stake-weighted leader schedule for one epoch.
+#[derive(Debug, PartialEq)]
+pub struct LeaderSchedule {
+    slot_leaders: Vec<Pubkey>,
+}
+
+impl LeaderSchedule {
+    pub fn new(ids_and_stakes: &[(Pubkey, u64)], seed: &[u8; 32], len: u64) -> Self {
+        let (pubkeys, stakes): (Vec<Pubkey>, Vec<u64>) = ids_and_stakes
+            .iter()
+            .map(|&(ref id, ref stake)| (id, stake))
+            .unzip();
+
+        // Should have no zero weighted stakes
+        let mut rng = ChaChaRng::from_seed(*seed);
+        let weighted_index = WeightedIndex::new(stakes).unwrap();
+        let slot_leaders = (0..len)
+            .map(|_| pubkeys[weighted_index.sample(&mut rng)])
+            .collect();
+
+        Self { slot_leaders }
+    }
+
+    pub fn new_with_bank(bank: &Bank) -> Self {
+        let mut seed = [0u8; 32];
+        seed.copy_from_slice(bank.last_id().as_ref());
+        let stakes: Vec<_> = bank.staked_nodes().into_iter().collect();
+        Self::new(&stakes, &seed, bank.slots_per_epoch())
+    }
+}
+
+impl Index<usize> for LeaderSchedule {
+    type Output = Pubkey;
+    fn index(&self, index: usize) -> &Pubkey {
+        &self.slot_leaders[index % self.slot_leaders.len()]
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use solana_sdk::genesis_block::GenesisBlock;
+    use solana_sdk::signature::{Keypair, KeypairUtil};
+    use std::iter;
+
+    #[test]
+    fn test_leader_schedule_index() {
+        let pubkey0 = Keypair::new().pubkey();
+        let pubkey1 = Keypair::new().pubkey();
+        let leader_schedule = LeaderSchedule {
+            slot_leaders: vec![pubkey0, pubkey1],
+        };
+        assert_eq!(leader_schedule[0], pubkey0);
+        assert_eq!(leader_schedule[1], pubkey1);
+        assert_eq!(leader_schedule[2], pubkey0);
+    }
+
+    #[test]
+    fn test_leader_schedule_basic() {
+        let num_keys = 10;
+        let stakes: Vec<_> = (0..num_keys)
+            .map(|i| (Keypair::new().pubkey(), i))
+            .collect();
+
+        let seed = Keypair::new().pubkey();
+        let mut seed_bytes = [0u8; 32];
+        seed_bytes.copy_from_slice(seed.as_ref());
+        let len = num_keys * 10;
+        let leader_schedule = LeaderSchedule::new(&stakes, &seed_bytes, len);
+        let leader_schedule2 = LeaderSchedule::new(&stakes, &seed_bytes, len);
+        assert_eq!(leader_schedule.slot_leaders.len() as u64, len);
+        // Check that the same schedule is reproducibly generated
+        assert_eq!(leader_schedule, leader_schedule2);
+    }
+
+    #[test]
+    fn test_leader_schedule_via_bank() {
+        let pubkey = Keypair::new().pubkey();
+        let (genesis_block, _mint_keypair) = GenesisBlock::new_with_leader(2, pubkey, 2);
+        let bank = Bank::new(&genesis_block);
+        let leader_schedule = LeaderSchedule::new_with_bank(&bank);
+        let len = bank.slots_per_epoch() as usize;
+        let expected: Vec<_> = iter::repeat(pubkey).take(len).collect();
+        assert_eq!(leader_schedule.slot_leaders, expected);
+    }
+}

runtime/src/lib.rs

@@ -2,6 +2,7 @@ mod accounts;
 pub mod bank;
 pub mod bloom;
 mod last_id_queue;
+mod leader_schedule;
 mod runtime;
 mod status_cache;