cleanup math in distribute_partitioned_epoch_rewards (#32164)

runtime/src/bank.rs

@@ -1934,15 +1934,15 @@ impl Bank {
                 return;
             };
 
-        assert!(
-            self.epoch_schedule.get_slots_in_epoch(self.epoch)
-                > self.get_reward_total_num_blocks(status.calculated_epoch_stake_rewards.len())
-        );
         let height = self.block_height();
         let start_block_height = status.start_block_height;
         let credit_start = start_block_height + self.get_reward_calculation_num_blocks();
-        let credit_end_exclusive = credit_start
-            + self.get_reward_distribution_num_blocks(status.calculated_epoch_stake_rewards.len());
+        let credit_end_exclusive =
+            credit_start + status.calculated_epoch_stake_rewards.len() as u64;
+        assert!(
+            self.epoch_schedule.get_slots_in_epoch(self.epoch)
+                > credit_end_exclusive.saturating_sub(credit_start)
+        );
 
         if height >= credit_start && height < credit_end_exclusive {
             let partition_index = height - credit_start;

runtime/src/bank/tests.rs

@@ -12633,7 +12633,30 @@ fn test_distribute_partitioned_epoch_rewards() {
         .map(|_| StakeReward::new_random())
         .collect::<Vec<_>>();
 
-    let stake_rewards = hash_rewards_into_partitions(stake_rewards, &Hash::new(&[1; 32]), 100);
+    let stake_rewards = hash_rewards_into_partitions(stake_rewards, &Hash::new(&[1; 32]), 2);
+
+    bank.set_epoch_reward_status_active(stake_rewards);
+
+    bank.distribute_partitioned_epoch_rewards();
+}
+
+#[test]
+#[should_panic(expected = "assertion failed: self.epoch_schedule.get_slots_in_epoch")]
+fn test_distribute_partitioned_epoch_rewards_too_many_partitions() {
+    let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+    let mut bank = Bank::new_for_tests(&genesis_config);
+
+    let expected_num = 1;
+
+    let stake_rewards = (0..expected_num)
+        .map(|_| StakeReward::new_random())
+        .collect::<Vec<_>>();
+
+    let stake_rewards = hash_rewards_into_partitions(
+        stake_rewards,
+        &Hash::new(&[1; 32]),
+        bank.epoch_schedule().slots_per_epoch as usize + 1,
+    );
 
     bank.set_epoch_reward_status_active(stake_rewards);