small improvements to the deposit_mining() logic

programs/rewards/src/instruction.rs

@@ -169,7 +169,7 @@ pub enum RewardsInstruction {
     },
 
     #[account(0, signer, name = "deposit_authority", desc = "The address of the Staking program's Registrar, which is PDA and is responsible for signing CPIs")]
-    #[account(1, name = "reward_pool", desc = "The address of the reward pool")]
+    #[account(1, writable, name = "reward_pool", desc = "The address of the reward pool")]
     #[account(2, writable, name = "mining", desc = "The address of the mining account which belongs to the user and stores info about user's rewards")]
     Slash {
         mining_owner: Pubkey,

programs/rewards/src/state/reward_pool.rs

@@ -218,27 +218,34 @@ impl<'a> WrappedRewardPool<'a> {
         self.pool.total_share = self.pool.total_share.safe_add(weighted_stake)?;
         mining.mining.share = mining.mining.share.safe_add(weighted_stake)?;
 
-        let stake_expiraion_date = lockup_period.end_timestamp(get_curr_unix_ts())?;
-        let modifier = if let Some(modifier) = self.weighted_stake_diffs.get(&stake_expiraion_date)
+        let stake_expiration_date = lockup_period.end_timestamp(get_curr_unix_ts())?;
+
+        let modifier = if let Some(modifier) = self.weighted_stake_diffs.get(&stake_expiration_date)
         {
             *modifier
         } else {
             0
         };
 
         self.weighted_stake_diffs.insert(
-            lockup_period.end_timestamp(get_curr_unix_ts())?,
+            stake_expiration_date,
             modifier.safe_add(weighted_stake_diff)?,
         );
 
-        let date_to_insert = &lockup_period.end_timestamp(get_curr_unix_ts())?;
-        if mining.weighted_stake_diffs.get(date_to_insert).is_some() {
-            let modifier = mining.weighted_stake_diffs.get_mut(date_to_insert).unwrap();
+        if mining
+            .weighted_stake_diffs
+            .get(&stake_expiration_date)
+            .is_some()
+        {
+            let modifier = mining
+                .weighted_stake_diffs
+                .get_mut(&stake_expiration_date)
+                .unwrap();
             *modifier = modifier.safe_add(weighted_stake_diff)?;
         } else {
             mining
                 .weighted_stake_diffs
-                .insert(*date_to_insert, weighted_stake_diff);
+                .insert(stake_expiration_date, weighted_stake_diff);
         }
 
         if let Some(delegate_mining_acc) = delegate_mining {
@@ -296,12 +303,13 @@ impl<'a> WrappedRewardPool<'a> {
     ) -> ProgramResult {
         self.withdraw(mining, slash_amount_multiplied_by_period, None)?;
 
-        if stake_expiration_date.is_some() {
-            let stake_expiration_date = stake_expiration_date.unwrap();
+        if let Some(stake_expiration_date) = stake_expiration_date {
             let beginning_of_the_stake_expiration_date =
                 stake_expiration_date - (stake_expiration_date % SECONDS_PER_DAY);
+
             let diff_by_expiration_date =
                 slash_amount_multiplied_by_period.safe_sub(slash_amount_in_native)?;
+
             let diff_record = mining
                 .weighted_stake_diffs
                 .get_mut(&beginning_of_the_stake_expiration_date)
@@ -314,6 +322,7 @@ impl<'a> WrappedRewardPool<'a> {
                 .ok_or(MplxRewardsError::NoWeightedStakeModifiersAtADate)?;
             *diff_record = diff_record.safe_sub(diff_by_expiration_date)?;
         }
+
         Ok(())
     }