From 5741400713971d94a2905bc1bd07c7a41737d6a0 Mon Sep 17 00:00:00 2001
From: Rob Walker <rob@solana.com>
Date: Thu, 14 Mar 2019 16:06:56 -0700
Subject: [PATCH] add support for finding the next slot a node will be leader
 (#3298)

---
 core/src/leader_schedule.rs       |  5 +--
 core/src/leader_schedule_utils.rs | 59 ++++++++++++++++++++++++++++++-
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/core/src/leader_schedule.rs b/core/src/leader_schedule.rs
index e5f6b9a8de..0a433df18d 100644
--- a/core/src/leader_schedule.rs
+++ b/core/src/leader_schedule.rs
@@ -21,9 +21,10 @@ impl LeaderSchedule {
     }
 }
 
-impl Index<usize> for LeaderSchedule {
+impl Index<u64> for LeaderSchedule {
     type Output = Pubkey;
-    fn index(&self, index: usize) -> &Pubkey {
+    fn index(&self, index: u64) -> &Pubkey {
+        let index = index as usize;
         &self.slot_leaders[index % self.slot_leaders.len()]
     }
 }
diff --git a/core/src/leader_schedule_utils.rs b/core/src/leader_schedule_utils.rs
index 4e0d6bf2e6..a91a4d82d8 100644
--- a/core/src/leader_schedule_utils.rs
+++ b/core/src/leader_schedule_utils.rs
@@ -34,7 +34,30 @@ fn sort_stakes(stakes: &mut Vec<(Pubkey, u64)>) {
 pub fn slot_leader_at(slot: u64, bank: &Bank) -> Option<Pubkey> {
     let (epoch, slot_index) = bank.get_epoch_and_slot_index(slot);
 
-    leader_schedule(epoch, bank).map(|leader_schedule| leader_schedule[slot_index as usize])
+    leader_schedule(epoch, bank).map(|leader_schedule| leader_schedule[slot_index])
+}
+
+/// Return the next slot after the given current_slot that the given node will be leader
+pub fn next_leader_slot(pubkey: &Pubkey, current_slot: u64, bank: &Bank) -> Option<u64> {
+    let (epoch, slot_index) = bank.get_epoch_and_slot_index(current_slot + 1);
+
+    if let Some(leader_schedule) = leader_schedule(epoch, bank) {
+        // clippy thinks I should do this:
+        //  for (i, <item>) in leader_schedule
+        //                           .iter()
+        //                           .enumerate()
+        //                           .take(bank.get_slots_in_epoch(epoch))
+        //                           .skip(from_slot_index + 1) {
+        //
+        //  but leader_schedule doesn't implement Iter...
+        #[allow(clippy::needless_range_loop)]
+        for i in slot_index..bank.get_slots_in_epoch(epoch) {
+            if *pubkey == leader_schedule[i] {
+                return Some(current_slot + 1 + (i - slot_index) as u64);
+            }
+        }
+    }
+    None
 }
 
 // Returns the number of ticks remaining from the specified tick_height to the end of the
@@ -54,6 +77,40 @@ mod tests {
     use solana_sdk::genesis_block::{GenesisBlock, BOOTSTRAP_LEADER_LAMPORTS};
     use solana_sdk::signature::{Keypair, KeypairUtil};
 
+    #[test]
+    fn test_next_leader_slot() {
+        let pubkey = Keypair::new().pubkey();
+        let mut genesis_block = GenesisBlock::new_with_leader(
+            BOOTSTRAP_LEADER_LAMPORTS,
+            &pubkey,
+            BOOTSTRAP_LEADER_LAMPORTS,
+        )
+        .0;
+        genesis_block.epoch_warmup = false;
+
+        let bank = Bank::new(&genesis_block);
+        assert_eq!(slot_leader_at(bank.slot(), &bank).unwrap(), pubkey);
+        assert_eq!(next_leader_slot(&pubkey, 0, &bank), Some(1));
+        assert_eq!(next_leader_slot(&pubkey, 1, &bank), Some(2));
+        assert_eq!(
+            next_leader_slot(
+                &pubkey,
+                2 * genesis_block.slots_per_epoch - 1, // no schedule generated for epoch 2
+                &bank
+            ),
+            None
+        );
+
+        assert_eq!(
+            next_leader_slot(
+                &Keypair::new().pubkey(), // not in leader_schedule
+                0,
+                &bank
+            ),
+            None
+        );
+    }
+
     #[test]
     fn test_leader_schedule_via_bank() {
         let pubkey = Keypair::new().pubkey();