From dcd66534dda718160c80752711e6c3a8c6686992 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Tue, 13 Jun 2023 15:57:52 -0500
Subject: [PATCH] add partitioned rewards helpers (#32085)

* add partitioned rewards helpers

* remove pub(crate)

* credit -> distribution
---
 runtime/src/bank.rs       | 32 ++++++++++++++++++
 runtime/src/bank/tests.rs | 70 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+)

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index cb4331967e..7bcca39f98 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -1511,6 +1511,38 @@ impl Bank {
             .reward_calculation_num_blocks
     }
 
+    #[allow(dead_code)]
+    /// Calculate the number of blocks required to distribute rewards to all stake accounts.
+    fn get_reward_distribution_num_blocks(&self, total_stake_accounts: usize) -> u64 {
+        if self.epoch_schedule.warmup && self.epoch < self.first_normal_epoch() {
+            1
+        } else {
+            let num_chunks = crate::accounts_hash::AccountsHasher::div_ceil(
+                total_stake_accounts,
+                self.partitioned_rewards_stake_account_stores_per_block() as usize,
+            ) as u64;
+
+            // Limit the reward credit interval to 5% of the total number of slots in a epoch
+            num_chunks.clamp(1, (self.epoch_schedule.slots_per_epoch / 20).max(1))
+        }
+    }
+
+    #[allow(dead_code)]
+    /// Return the total number of blocks in reward interval (including both calculation and crediting).
+    fn get_reward_total_num_blocks(&self, total_stake_accounts: usize) -> u64 {
+        self.get_reward_calculation_num_blocks()
+            + self.get_reward_distribution_num_blocks(total_stake_accounts)
+    }
+
+    #[allow(dead_code)]
+    /// Return `RewardInterval` enum for current bank
+    fn get_reward_interval(&self) -> RewardInterval {
+        if matches!(self.epoch_reward_status, EpochRewardStatus::Active(_)) {
+            RewardInterval::InsideInterval
+        } else {
+            RewardInterval::OutsideInterval
+        }
+    }
     fn _new_from_parent(
         parent: &Arc<Bank>,
         collector_id: &Pubkey,
diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
index bdc5ce44c8..e946a234e7 100644
--- a/runtime/src/bank/tests.rs
+++ b/runtime/src/bank/tests.rs
@@ -13072,6 +13072,76 @@ fn test_calc_vote_accounts_to_store_normal() {
     }
 }
 
+/// Test get_reward_distribution_num_blocks, get_reward_calculation_num_blocks, get_reward_total_num_blocks during normal epoch gives the expected result
+#[test]
+fn test_get_reward_distribution_num_blocks_normal() {
+    solana_logger::setup();
+    let (mut genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+    genesis_config.epoch_schedule = EpochSchedule::custom(432000, 432000, false);
+
+    let bank = Bank::new_for_tests(&genesis_config);
+
+    // Given 8k rewards, it will take 2 blocks to credit all the rewards
+    let expected_num = 8192;
+    let stake_rewards = (0..expected_num)
+        .map(|_| StakeReward::new_random())
+        .collect::<Vec<_>>();
+
+    assert_eq!(
+        bank.get_reward_distribution_num_blocks(stake_rewards.len()),
+        2
+    );
+    assert_eq!(bank.get_reward_calculation_num_blocks(), 1);
+    assert_eq!(
+        bank.get_reward_total_num_blocks(stake_rewards.len()),
+        bank.get_reward_distribution_num_blocks(stake_rewards.len())
+            + bank.get_reward_calculation_num_blocks(),
+    );
+}
+
+/// Test get_reward_distribution_num_blocks, get_reward_calculation_num_blocks, get_reward_total_num_blocks during small epoch
+/// The num_credit_blocks should be cap to 5% of the total number of blocks in the epoch.
+#[test]
+fn test_get_reward_distribution_num_blocks_cap() {
+    let (mut genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+    genesis_config.epoch_schedule = EpochSchedule::custom(32, 32, false);
+
+    let bank = Bank::new_for_tests(&genesis_config);
+
+    // Given 8k rewards, normally it will take 2 blocks to credit all the rewards. However, because of
+    // the short epoch, i.e. 32 slots, we should cap the number of credit blocks to 32/20 = 1.
+    let expected_num = 8192;
+    let stake_rewards = (0..expected_num)
+        .map(|_| StakeReward::new_random())
+        .collect::<Vec<_>>();
+
+    assert_eq!(
+        bank.get_reward_distribution_num_blocks(stake_rewards.len()),
+        1
+    );
+    assert_eq!(bank.get_reward_calculation_num_blocks(), 1);
+    assert_eq!(
+        bank.get_reward_total_num_blocks(stake_rewards.len()),
+        bank.get_reward_distribution_num_blocks(stake_rewards.len())
+            + bank.get_reward_calculation_num_blocks(),
+    );
+}
+
+/// Test get_reward_distribution_num_blocks, get_reward_calculation_num_blocks, get_reward_total_num_blocks during warm up epoch gives the expected result.
+/// The num_credit_blocks should be 1 during warm up epoch.
+#[test]
+fn test_get_reward_distribution_num_blocks_warmup() {
+    let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+
+    let bank = Bank::new_for_tests(&genesis_config);
+    assert_eq!(bank.get_reward_distribution_num_blocks(0), 1);
+    assert_eq!(bank.get_reward_calculation_num_blocks(), 1);
+    assert_eq!(
+        bank.get_reward_total_num_blocks(0),
+        bank.get_reward_distribution_num_blocks(0) + bank.get_reward_calculation_num_blocks(),
+    );
+}
+
 #[test]
 fn test_calculate_stake_vote_rewards() {
     solana_logger::setup();