From 63a4ed74a4b7b34f9d291e3e35d2b71ad5a85611 Mon Sep 17 00:00:00 2001
From: Rob Walker <rob@solana.com>
Date: Wed, 6 Mar 2019 14:44:21 -0800
Subject: [PATCH] consolidate logic for epoch and slot_index into Bank (#3144)

---
 core/src/leader_schedule_utils.rs |  3 +--
 core/src/staking_utils.rs         | 15 +++++++++------
 runtime/src/bank.rs               | 27 +++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/core/src/leader_schedule_utils.rs b/core/src/leader_schedule_utils.rs
index 977f1d0d10..43b87c7442 100644
--- a/core/src/leader_schedule_utils.rs
+++ b/core/src/leader_schedule_utils.rs
@@ -32,8 +32,7 @@ fn sort_stakes(stakes: &mut Vec<(Pubkey, u64)>) {
 
 /// Return the leader for the given slot.
 pub fn slot_leader_at(slot: u64, bank: &Bank) -> Pubkey {
-    let slot_index = slot % bank.slots_per_epoch();
-    let epoch = slot / bank.slots_per_epoch();
+    let (epoch, slot_index) = bank.get_epoch_and_slot_index(slot);
 
     let leader_schedule = leader_schedule(epoch, bank);
     leader_schedule[slot_index as usize]
diff --git a/core/src/staking_utils.rs b/core/src/staking_utils.rs
index f8ef25c62f..33732bb002 100644
--- a/core/src/staking_utils.rs
+++ b/core/src/staking_utils.rs
@@ -176,7 +176,12 @@ mod tests {
     fn test_epoch_stakes_and_lockouts() {
         let validator = Keypair::new();
 
-        let (genesis_block, mint_keypair) = GenesisBlock::new(500);
+        let (mut genesis_block, mint_keypair) = GenesisBlock::new(500);
+
+        // makes the bank generate epoch_vote_accounts right on epoch boundaries
+        //   handy for tests
+        genesis_block.stakers_slot_offset = 0;
+
         let bank = Bank::new(&genesis_block);
         let bank_voter = Keypair::new();
 
@@ -190,13 +195,11 @@ mod tests {
         // should show up in the active set
         voting_keypair_tests::new_vote_account_with_vote(&mint_keypair, &bank_voter, &bank, 499, 0);
 
-        // Have to wait until the epoch at (stakers_slot_offset / slots_per_epoch) + 1
-        // for the new votes to take effect. Earlier epochs were generated by genesis
-        let epoch = (bank.stakers_slot_offset() / bank.slots_per_epoch()) + 1;
+        // Build a bank in the next epoch, it will generate an epoch_vote_accounts() for epoch 1
+        let epoch = 1;
         let epoch_slot = epoch * bank.slots_per_epoch();
-        let epoch_slot_offset = epoch_slot - bank.stakers_slot_offset();
 
-        let bank = new_from_parent(&Arc::new(bank), epoch_slot_offset);
+        let bank = new_from_parent(&Arc::new(bank), epoch_slot);
 
         let result: Vec<_> = epoch_stakes_and_lockouts(&bank, 0);
         assert_eq!(result, vec![(1, None)]);
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 02754156b2..5b331ae716 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -794,6 +794,15 @@ impl Bank {
     pub fn epoch_vote_accounts(&self, epoch: u64) -> Option<&HashMap<Pubkey, Account>> {
         self.epoch_vote_accounts.get(&epoch)
     }
+
+    /// given a slot, return the epoch and offset into the epoch this slot falls
+    /// e.g. with a fixed number for slots_per_epoch, the calculation is simply:
+    ///
+    ///  ( slot/slots_per_epoch, slot % slots_per_epoch )
+    ///
+    pub fn get_epoch_and_slot_index(&self, slot: u64) -> (u64, u64) {
+        (slot / self.slots_per_epoch(), slot % self.slots_per_epoch())
+    }
 }
 
 #[cfg(test)]
@@ -1515,4 +1524,22 @@ mod tests {
         assert_eq!(bank.get_balance(&key.pubkey()), 0);
     }
 
+    #[test]
+    fn test_bank_get_epoch_and_slot_offset() {
+        let (mut genesis_block, _) = GenesisBlock::new(500);
+
+        // set this up weird, forces:
+        //  1. genesis bank to cover epochs 0, 1, *and* 2
+        //  2. child banks to cover epochs in their future
+        //
+        const SLOTS_PER_EPOCH: u64 = 8;
+        genesis_block.slots_per_epoch = SLOTS_PER_EPOCH;
+
+        let bank = Bank::new(&genesis_block);
+
+        assert_eq!(bank.get_epoch_and_slot_index(0), (0, 0));
+        assert_eq!(bank.get_epoch_and_slot_index(SLOTS_PER_EPOCH), (1, 0));
+        assert_eq!(bank.get_epoch_and_slot_index(SLOTS_PER_EPOCH + 1), (1, 1));
+    }
+
 }