From 9774826b0803f733a6dbc01a7775ea5462cc2bcb Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Tue, 13 Jun 2023 16:45:19 -0500
Subject: [PATCH] add calculate_validator_rewards (#32107)

---
 runtime/src/bank.rs       | 31 ++++++++++++++++++++++++++++++
 runtime/src/bank/tests.rs | 40 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 7bcca39f98..1213b9066b 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -2896,6 +2896,37 @@ impl Bank {
         }
     }
 
+    #[allow(dead_code)]
+    /// Calculate epoch reward and return vote and stake rewards.
+    fn calculate_validator_rewards(
+        &self,
+        rewarded_epoch: Epoch,
+        rewards: u64,
+        reward_calc_tracer: Option<impl RewardCalcTracer>,
+        thread_pool: &ThreadPool,
+        metrics: &mut RewardsMetrics,
+    ) -> Option<(VoteRewardsAccounts, StakeRewardCalculation)> {
+        let stakes = self.stakes_cache.stakes();
+        let reward_calculate_param = self.get_epoch_reward_calculate_param_info(&stakes);
+
+        self.calculate_reward_points_partitioned(
+            &reward_calculate_param,
+            rewards,
+            thread_pool,
+            metrics,
+        )
+        .map(|point_value| {
+            self.calculate_stake_vote_rewards(
+                &reward_calculate_param,
+                rewarded_epoch,
+                point_value,
+                thread_pool,
+                reward_calc_tracer,
+                metrics,
+            )
+        })
+    }
+
     /// Load, calculate and payout epoch rewards for stake and vote accounts
     fn pay_validator_rewards_with_thread_pool(
         &mut self,
diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
index e946a234e7..d008ef24a1 100644
--- a/runtime/src/bank/tests.rs
+++ b/runtime/src/bank/tests.rs
@@ -12606,6 +12606,46 @@ fn test_rewards_point_calculation_empty() {
 
 /// Test reward computation at epoch boundary
 #[test]
+fn test_rewards_computation() {
+    solana_logger::setup();
+
+    let expected_num_delegations = 100;
+    let bank = create_reward_bank(expected_num_delegations).0;
+
+    // Calculate rewards
+    let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap();
+    let mut rewards_metrics = RewardsMetrics::default();
+    let expected_rewards = 100_000_000_000;
+
+    let calculated_rewards = bank.calculate_validator_rewards(
+        1,
+        expected_rewards,
+        null_tracer(),
+        &thread_pool,
+        &mut rewards_metrics,
+    );
+
+    let vote_rewards = &calculated_rewards.as_ref().unwrap().0;
+    let stake_rewards = &calculated_rewards.as_ref().unwrap().1;
+
+    let total_vote_rewards: u64 = vote_rewards
+        .rewards
+        .iter()
+        .map(|reward| reward.1.lamports)
+        .sum::<i64>() as u64;
+
+    // assert that total rewards matches the sum of vote rewards and stake rewards
+    assert_eq!(
+        stake_rewards.total_stake_rewards_lamports + total_vote_rewards,
+        expected_rewards
+    );
+
+    // assert that number of stake rewards matches
+    assert_eq!(stake_rewards.stake_rewards.len(), expected_num_delegations);
+}
+
+/// Test rewards compuation and partitioned rewards distribution at the epoch boundary
+#[test]
 fn test_store_stake_accounts_in_partition() {
     let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
     let bank = Bank::new_for_tests(&genesis_config);