From b4236328a6b670b09160c1e6b242abb905bf0a67 Mon Sep 17 00:00:00 2001
From: rwwwx <vlad.matiukhin@gmail.com>
Date: Wed, 17 Jul 2024 12:09:39 +0300
Subject: [PATCH 1/2] Threw an error if `weighted_stake_diffs.entry()` is empty

---
 programs/rewards/src/error.rs             |  5 +++
 programs/rewards/src/state/reward_pool.rs | 38 +++++++++++++++++------
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/programs/rewards/src/error.rs b/programs/rewards/src/error.rs
index 1b79a049..cb215dc3 100644
--- a/programs/rewards/src/error.rs
+++ b/programs/rewards/src/error.rs
@@ -75,6 +75,11 @@ pub enum MplxRewardsError {
     /// No need to transfer zero amount of rewards.
     #[error("Rewards: rewards amount must be positive")]
     RewardsMustBeGreaterThanZero,
+
+    /// 13
+    /// No need to transfer zero amount of rewards.
+    #[error("'weighted_stake_diffs' entry cannot be empty")]
+    WeightedStakeDiffEntryIsEmpty,
 }
 
 impl PrintProgramError for MplxRewardsError {
diff --git a/programs/rewards/src/state/reward_pool.rs b/programs/rewards/src/state/reward_pool.rs
index d6a0f346..e882bdc8 100644
--- a/programs/rewards/src/state/reward_pool.rs
+++ b/programs/rewards/src/state/reward_pool.rs
@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::collections::{btree_map::Entry, BTreeMap};
 
 use crate::{
     error::MplxRewardsError,
@@ -216,16 +216,17 @@ impl RewardPool {
                 .checked_sub(curr_part_of_weighted_stake_for_flex)
                 .ok_or(MplxRewardsError::MathOverflow)?;
 
-            self.calculator
-                .weighted_stake_diffs
-                .entry(deposit_old_expiration_ts)
-                .and_modify(|modifier| *modifier -= weighted_stake_diff);
+            Self::modify_weighted_stake_diffs(
+                &mut self.calculator.weighted_stake_diffs,
+                deposit_old_expiration_ts,
+                weighted_stake_diff,
+            )?;
 
-            mining
-                .index
-                .weighted_stake_diffs
-                .entry(deposit_old_expiration_ts)
-                .and_modify(|modifier| *modifier -= weighted_stake_diff);
+            Self::modify_weighted_stake_diffs(
+                &mut mining.index.weighted_stake_diffs,
+                deposit_old_expiration_ts,
+                weighted_stake_diff,
+            )?;
 
             // also, we need to reduce staking power because we want to extend stake from "scratch"
             mining.share = mining
@@ -258,6 +259,23 @@ impl RewardPool {
 
         Ok(())
     }
+
+    fn modify_weighted_stake_diffs(
+        diffs: &mut BTreeMap<u64, u64>,
+        timestamp: u64,
+        weighted_stake_diff: u64,
+    ) -> Result<(), MplxRewardsError> {
+        match diffs.entry(timestamp) {
+            Entry::Vacant(_) => Err(MplxRewardsError::WeightedStakeDiffEntryIsEmpty),
+            Entry::Occupied(mut entry) => {
+                let modifier = entry.get_mut();
+                *modifier = modifier
+                    .checked_sub(weighted_stake_diff)
+                    .ok_or(MplxRewardsError::MathOverflow)?;
+                Ok(())
+            }
+        }
+    }
 }
 
 impl Sealed for RewardPool {}

From eab79337341315107c802f2a4d1e69d4cfd28232 Mon Sep 17 00:00:00 2001
From: rwwwx <vlad.matiukhin@gmail.com>
Date: Wed, 17 Jul 2024 12:39:37 +0300
Subject: [PATCH 2/2] renamed error, changed err message

---
 programs/rewards/src/error.rs             | 4 ++--
 programs/rewards/src/state/reward_pool.rs | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/programs/rewards/src/error.rs b/programs/rewards/src/error.rs
index cb215dc3..b1de3293 100644
--- a/programs/rewards/src/error.rs
+++ b/programs/rewards/src/error.rs
@@ -78,8 +78,8 @@ pub enum MplxRewardsError {
 
     /// 13
     /// No need to transfer zero amount of rewards.
-    #[error("'weighted_stake_diffs' entry cannot be empty")]
-    WeightedStakeDiffEntryIsEmpty,
+    #[error("No changes at the date in weighted stake modifiers while they're expected")]
+    NoWeightedStakeModifiersAtADate,
 }
 
 impl PrintProgramError for MplxRewardsError {
diff --git a/programs/rewards/src/state/reward_pool.rs b/programs/rewards/src/state/reward_pool.rs
index e882bdc8..0df6a816 100644
--- a/programs/rewards/src/state/reward_pool.rs
+++ b/programs/rewards/src/state/reward_pool.rs
@@ -266,7 +266,7 @@ impl RewardPool {
         weighted_stake_diff: u64,
     ) -> Result<(), MplxRewardsError> {
         match diffs.entry(timestamp) {
-            Entry::Vacant(_) => Err(MplxRewardsError::WeightedStakeDiffEntryIsEmpty),
+            Entry::Vacant(_) => Err(MplxRewardsError::NoWeightedStakeModifiersAtADate),
             Entry::Occupied(mut entry) => {
                 let modifier = entry.get_mut();
                 *modifier = modifier