improve update_stake instruction to keep track of partially unstake

Author: RustNinja

solana/merkle-distributor/src/lib.rs

@@ -131,32 +131,74 @@ pub mod merkle_distributor {
         amount: u64,
         is_staked: bool, //true if staked, false if unstaked
     ) -> Result<()> {
+        // Check program ID
         let program_id = get_staking_caller(&ctx.accounts.instruction)?;
         let staking_contract = ctx.accounts.distributor.staking_contract_address;
         if program_id != staking_contract {
-            panic!("Invalid program id");
+            panic!("Invalid program ID");
         }
+
+        // Get the current timestamp
         let current_ts = Clock::get()?.unix_timestamp;
         let claimer_status = &mut ctx.accounts.claim_status;
+
         if is_staked {
+            // Handle staking: Add a new item to `staked_list`
             claimer_status.staked_list.push(StakedItem {
                 stake_amount: amount,
                 stake_timestamp_start: Some(current_ts),
                 already_staked_duration_time: 0,
             });
         } else {
-            let curr_ts = Clock::get()?.unix_timestamp;
-            let mut _unstake_amount = 0;
-            for i in claimer_status.staked_list.iter_mut(){
-                let Some(stake_timestamp_start) = i.stake_timestamp_start else{
+            // Handle unstaking: Process the required unstake amount
+            let mut remaining_unstake_amount = amount;
+
+            // We need a new list to store updated staked items after partial unstaking
+            let mut new_staked_item = None;
+
+            for staked_item in claimer_status.staked_list.iter_mut() {
+                if remaining_unstake_amount == 0 {
+                    // If no more amount needs to be unstaked, keep the rest as-is
+                    break;
+                }
+
+                // Only process items that have a start timestamp (still staked)
+                let Some(stake_timestamp_start) = staked_item.stake_timestamp_start else {
                     continue;
                 };
-                i.stake_timestamp_start = None;
-                let diff = curr_ts.checked_sub(stake_timestamp_start).ok_or(crate::error::ErrorCode::ArithmeticError)?;
-                i.already_staked_duration_time = i.already_staked_duration_time.checked_add(diff).ok_or(crate::error::ErrorCode::ArithmeticError)?;
-                _unstake_amount += i.stake_amount;
+
+                // Calculate time staked and update the duration
+                let time_staked = current_ts.checked_sub(stake_timestamp_start)
+                    .ok_or(crate::error::ErrorCode::ArithmeticError)?;
+                staked_item.already_staked_duration_time = staked_item
+                    .already_staked_duration_time
+                    .checked_add(time_staked)
+                    .ok_or(crate::error::ErrorCode::ArithmeticError)?;
+
+                if staked_item.stake_amount <= remaining_unstake_amount {
+                    // Fully unstake this item and update the remaining amount
+                    remaining_unstake_amount -= staked_item.stake_amount;
+                } else {
+                    // Partially unstake: Keep original as-is, and create a new item for unstaked portion
+                    new_staked_item = Some(StakedItem {
+                        stake_amount: staked_item.stake_amount - remaining_unstake_amount,
+                        stake_timestamp_start: Some(current_ts), // New start time for remaining amount
+                        already_staked_duration_time: 0,
+                    });
+
+                    // Adjust the original item to reflect the unstaked portion
+                    staked_item.stake_amount = remaining_unstake_amount;
+                    staked_item.stake_timestamp_start = None;
+                    remaining_unstake_amount = 0;
+                }
+            }
+
+            if let Some(new_item) = new_staked_item.clone() {
+                claimer_status.staked_list.push(new_item);
             }
+
         }
+
         Ok(())
     }
 }