feat: allow withdrawing without cleaning (#502)

* feat: allow withdraw without cleaning up unlocked positions * clean * clean * remove parantheses
pyth-network · Aug 28, 2024 · 4d400e4 · 4d400e4
1 parent b08f91f
commit 4d400e4
Show file tree

Hide file tree

Showing 5 changed files with 154 additions and 109 deletions.
diff --git a/staking/programs/staking/src/lib.rs b/staking/programs/staking/src/lib.rs
@@ -226,6 +226,8 @@ pub mod staking {
             stake_account_positions,
             stake_account_custody.amount,
             unvested_balance,
+            current_epoch,
+            config.unlocking_duration,
         )?;
 
         if let Some(target_account) = maybe_target_account {
@@ -405,6 +407,7 @@ pub mod staking {
         let destination_account = &ctx.accounts.destination;
         let signer = &ctx.accounts.owner;
         let config = &ctx.accounts.config;
+        let current_epoch = get_current_epoch(config)?;
 
 
         let unvested_balance = ctx
@@ -426,8 +429,14 @@ pub mod staking {
             .amount
             .checked_sub(amount)
             .ok_or_else(|| error!(ErrorCode::InsufficientWithdrawableBalance))?;
-        if utils::risk::validate(stake_account_positions, remaining_balance, unvested_balance)
-            .is_err()
+        if utils::risk::validate(
+            stake_account_positions,
+            remaining_balance,
+            unvested_balance,
+            current_epoch,
+            config.unlocking_duration,
+        )
+        .is_err()
         {
             return Err(error!(ErrorCode::InsufficientWithdrawableBalance));
         }
@@ -447,6 +456,8 @@ pub mod staking {
             stake_account_positions,
             ctx.accounts.stake_account_custody.amount,
             unvested_balance,
+            current_epoch,
+            config.unlocking_duration,
         )
         .is_err()
         {
@@ -488,6 +499,8 @@ pub mod staking {
             stake_account_positions,
             stake_account_custody.amount,
             unvested_balance,
+            current_epoch,
+            config.unlocking_duration,
         )?;
 
         let epoch_of_snapshot: u64;
@@ -637,6 +650,7 @@ pub mod staking {
      */
     pub fn accept_split(ctx: Context<AcceptSplit>, amount: u64, recipient: Pubkey) -> Result<()> {
         let config = &ctx.accounts.config;
+        let current_epoch = get_current_epoch(config)?;
 
         let split_request = &ctx.accounts.source_stake_account_split_request;
         require!(
@@ -673,6 +687,8 @@ pub mod staking {
                     utils::clock::get_current_time(config),
                     config.pyth_token_list_time,
                 )?,
+            current_epoch,
+            config.unlocking_duration,
         )?;
 
         // Check that there aren't any positions (i.e., staked tokens) in the source account.
@@ -726,6 +742,8 @@ pub mod staking {
                     utils::clock::get_current_time(config),
                     config.pyth_token_list_time,
                 )?,
+            current_epoch,
+            config.unlocking_duration,
         )?;
 
         utils::risk::validate(
@@ -738,6 +756,8 @@ pub mod staking {
                     utils::clock::get_current_time(config),
                     config.pyth_token_list_time,
                 )?,
+            current_epoch,
+            config.unlocking_duration,
         )?;
 
         // Delete current request

diff --git a/staking/programs/staking/src/state/positions.rs b/staking/programs/staking/src/state/positions.rs
@@ -1,9 +1,6 @@
 use {
     super::target::TargetMetadata,
-    crate::{
-        error::ErrorCode,
-        utils::risk::calculate_governance_exposure,
-    },
+    crate::error::ErrorCode,
     anchor_lang::{
         prelude::{
             borsh::BorshSchema,
@@ -212,6 +209,28 @@ impl<'a> DynamicPositionArray<'a> {
         Ok(false)
     }
 
+    pub fn get_target_exposure(
+        &self,
+        target: &Target,
+        current_epoch: u64,
+        unlocking_duration: u8,
+    ) -> Result<u64> {
+        let mut exposure: u64 = 0;
+        for i in 0..self.get_position_capacity() {
+            if let Some(position) = self.read_position(i)? {
+                if position.target_with_parameters.get_target() == *target
+                    && position.get_current_position(current_epoch, unlocking_duration)?
+                        != PositionState::UNLOCKED
+                {
+                    exposure = exposure
+                        .checked_add(position.amount)
+                        .ok_or_else(|| error!(ErrorCode::GenericOverflow))?;
+                }
+            }
+        }
+        Ok(exposure)
+    }
+
     /// This function is used to reduce the number of positions in the array by merging equivalent
     /// positions. Sometimes some positions have the same `target_with_parameters`,
     /// `activation_epoch` and `unlocking_start`. These can obviously be merged, but this is not
@@ -337,7 +356,8 @@ impl<'a> DynamicPositionArray<'a> {
             i += 1;
         }
 
-        let governance_exposure = calculate_governance_exposure(self)?;
+        let governance_exposure =
+            self.get_target_exposure(&Target::Voting, current_epoch, unlocking_duration)?;
 
         let total_slashed = locked_slashed + unlocking_slashed + preunlocking_slashed;
         if let Some(mut remaining) =
@@ -351,7 +371,9 @@ impl<'a> DynamicPositionArray<'a> {
                     let current_state =
                         position.get_current_position(current_epoch, unlocking_duration)?;
 
-                    if position.target_with_parameters == TargetWithParameters::Voting {
+                    if position.target_with_parameters == TargetWithParameters::Voting
+                        && current_state != PositionState::UNLOCKED
+                    {
                         let to_slash = remaining.min(position.amount);
                         remaining -= to_slash;
 
@@ -610,21 +632,19 @@ impl std::fmt::Display for PositionState {
 pub mod tests {
     use {
         super::DynamicPositionArray,
-        crate::{
-            state::{
-                positions::{
-                    DynamicPositionArrayAccount,
-                    Position,
-                    PositionData,
-                    PositionState,
-                    SlashedAmounts,
-                    TargetWithParameters,
-                    TryBorsh,
-                    POSITION_BUFFER_SIZE,
-                },
-                target::TargetMetadata,
+        crate::state::{
+            positions::{
+                DynamicPositionArrayAccount,
+                Position,
+                PositionData,
+                PositionState,
+                SlashedAmounts,
+                Target,
+                TargetWithParameters,
+                TryBorsh,
+                POSITION_BUFFER_SIZE,
             },
-            utils::risk::calculate_governance_exposure,
+            target::TargetMetadata,
         },
         anchor_lang::prelude::*,
         quickcheck::{
@@ -1151,7 +1171,9 @@ pub mod tests {
             }
         };
 
-        let governance_exposure = calculate_governance_exposure(&dynamic_position_array).unwrap();
+        let governance_exposure = dynamic_position_array
+            .get_target_exposure(&Target::Voting, epoch, 1)
+            .unwrap();
 
         let publisher_1_exposure = pre_position_buckets
             .iter()
@@ -1250,8 +1272,10 @@ pub mod tests {
         }
 
         // check governance exposure has been reduced by the correct amount
-        let post_governance_exposure =
-            calculate_governance_exposure(&dynamic_position_array).unwrap();
+        let post_governance_exposure = dynamic_position_array
+            .get_target_exposure(&Target::Voting, epoch, 1)
+            .unwrap();
+
         if post_governance_exposure
             != std::cmp::min(governance_exposure, custody_account_amount - total_slashed)
         {
@@ -1385,6 +1409,14 @@ pub mod tests {
             {
                 return false;
             }
+
+            if target == &TargetWithParameters::Voting
+                && curr_state == &PositionState::UNLOCKED
+                && pre_position_buckets.get(&(*target, *prev_state, *curr_state))
+                    != post_position_buckets.get(&(*target, *prev_state, *curr_state))
+            {
+                return false;
+            }
         }
 
         // the returned values match the position updates