refactor voter_log_info logic && refactor DepositEntry:: methods

program-states/src/state/deposit_entry.rs

@@ -39,28 +39,23 @@ impl DepositEntry {
 
     /// Returns native tokens still locked.
     #[inline(always)]
-    pub fn amount_locked(&self, curr_ts: u64) -> u64 {
-        let unlocked_tokens = if self.lockup.expired(curr_ts) {
+    pub fn amount_locked(&self) -> u64 {
+        if self.is_staked() {
             self.amount_deposited_native
         } else {
             0
-        };
-        self.amount_deposited_native
-            .checked_sub(unlocked_tokens)
-            .unwrap()
+        }
     }
 
     /// Returns native tokens that are unlocked given current vesting
     /// and previous withdraws.
     #[inline(always)]
-    pub fn amount_unlocked(&self, curr_ts: u64) -> u64 {
-        if let LockupKind::None = self.lockup.kind {
-            return self.amount_deposited_native;
+    pub fn amount_unlocked(&self) -> u64 {
+        if self.is_staked() {
+            0
+        } else {
+            self.amount_deposited_native
         }
-
-        self.amount_deposited_native
-            .checked_sub(self.amount_locked(curr_ts))
-            .unwrap()
     }
 
     /// Returns the weighted stake for the given deposit at the specified timestamp.
@@ -80,6 +75,7 @@ impl DepositEntry {
         self.is_used
             && self.lockup.kind.ne(&LockupKind::None)
             && self.lockup.period.ne(&LockupPeriod::None)
+            && !self.lockup.cooldown_requested
     }
 }
 
@@ -116,7 +112,7 @@ mod tests {
         assert_eq!(baseline_vote_weight, 20_000);
 
         // The timestamp 100_000 is very far before the lockup_start timestamp
-        let withdrawable = deposit.amount_unlocked(100_000);
+        let withdrawable = deposit.amount_unlocked();
         assert_eq!(withdrawable, 0);
 
         let voting_power = deposit.voting_power().unwrap();

programs/voter-stake-registry/src/instructions/close_voter.rs

@@ -1,4 +1,4 @@
-use crate::{clock_unix_timestamp, cpi_instructions};
+use crate::cpi_instructions;
 use anchor_lang::prelude::*;
 use anchor_spl::token::{self, CloseAccount, Token, TokenAccount, Transfer};
 use bytemuck::bytes_of_mut;
@@ -58,11 +58,10 @@ pub fn close_voter<'key, 'accounts, 'remaining, 'info>(
     ctx: Context<'key, 'accounts, 'remaining, 'info, CloseVoter<'info>>,
 ) -> Result<()> {
     let registrar = ctx.accounts.registrar.load()?;
-    let curr_ts = clock_unix_timestamp();
     {
         let voter = ctx.accounts.voter.load()?;
 
-        let any_locked = voter.deposits.iter().any(|d| d.amount_locked(curr_ts) > 0);
+        let any_locked = voter.deposits.iter().any(|d| d.amount_locked() > 0);
         require!(!any_locked, VsrError::DepositStillLocked);
 
         let active_deposit_entries = voter.deposits.iter().filter(|d| d.is_used).count();

programs/voter-stake-registry/src/instructions/extend_stake.rs

@@ -26,7 +26,7 @@ pub fn extend_stake(
 
     let source = voter.active_deposit_mut(source_deposit_entry_index)?;
     let source_mint_idx = source.voting_mint_config_idx;
-    let source_available_tokens = source.amount_unlocked(curr_ts);
+    let source_available_tokens = source.amount_unlocked();
     require!(
         source.lockup.kind == LockupKind::None,
         VsrError::LockingIsForbidded

programs/voter-stake-registry/src/instructions/log_voter_info.rs

@@ -1,6 +1,6 @@
-use crate::{clock_unix_timestamp, events::*};
+use crate::events::*;
 use anchor_lang::prelude::*;
-use mplx_staking_states::state::{LockupKind, Registrar, Voter};
+use mplx_staking_states::state::{Registrar, Voter};
 
 #[derive(Accounts)]
 pub struct LogVoterInfo<'info> {
@@ -23,7 +23,6 @@ pub fn log_voter_info(
     deposit_entry_count: u8,
 ) -> Result<()> {
     let voter = ctx.accounts.voter.load()?;
-    let curr_ts = clock_unix_timestamp();
     let deposit_entry_begin = deposit_entry_begin as usize;
     let deposit_entry_count = deposit_entry_count as usize;
 
@@ -41,19 +40,20 @@ pub fn log_voter_info(
         {
             continue;
         }
-        let lockup = &deposit.lockup;
-        let seconds_left = lockup.seconds_left(curr_ts);
-        let end_ts = curr_ts + seconds_left;
-        let locking_info = (seconds_left > 0).then(|| LockingInfo {
-            amount: deposit.amount_locked(curr_ts),
-            end_timestamp: (lockup.kind != LockupKind::Constant).then_some(end_ts),
-            vesting: None,
-        });
+        let locking_info = if !deposit.lockup.cooldown_requested {
+            Some(LockingInfo {
+                amount: deposit.amount_locked(),
+                end_timestamp: Some(deposit.lockup.end_ts),
+                vesting: None,
+            })
+        } else {
+            None
+        };
 
         emit!(DepositEntryInfo {
             deposit_entry_index: deposit_index as u8,
             voting_mint_config_index: deposit.voting_mint_config_idx,
-            unlocked: deposit.amount_unlocked(curr_ts),
+            unlocked: deposit.amount_unlocked(),
             voting_power: deposit.voting_power()?,
             voting_power_baseline: deposit.amount_deposited_native,
             locking: locking_info,

programs/voter-stake-registry/src/instructions/stake.rs

@@ -1,4 +1,4 @@
-use crate::{clock_unix_timestamp, cpi_instructions, Stake};
+use crate::{cpi_instructions, Stake};
 use anchor_lang::prelude::*;
 use mplx_staking_states::{error::VsrError, state::LockupKind};
 
@@ -15,7 +15,6 @@ pub fn stake(
 ) -> Result<()> {
     let registrar = &ctx.accounts.registrar.load()?;
     let voter = &mut ctx.accounts.voter.load_mut()?;
-    let curr_ts = clock_unix_timestamp();
 
     let source = voter.active_deposit_mut(source_deposit_entry_index)?;
     let source_mint_idx = source.voting_mint_config_idx;
@@ -26,7 +25,7 @@ pub fn stake(
 
     // Reduce source amounts
     require_gte!(
-        source.amount_unlocked(curr_ts),
+        source.amount_unlocked(),
         amount,
         VsrError::InsufficientUnlockedTokens
     );

programs/voter-stake-registry/src/instructions/withdraw.rs

@@ -123,7 +123,7 @@ pub fn withdraw(ctx: Context<Withdraw>, deposit_entry_index: u8, amount: u64) ->
     }
 
     require_gte!(
-        deposit_entry.amount_unlocked(curr_ts),
+        deposit_entry.amount_unlocked(),
         amount,
         VsrError::InsufficientUnlockedTokens
     );

programs/voter-stake-registry/tests/program_test/mod.rs

@@ -19,7 +19,7 @@ use std::{
     str::FromStr,
     sync::{Arc, RwLock},
 };
-pub use utils::{assert_custom_on_chain_error::AssertCustomOnChainErr, *};
+pub use utils::*;
 
 pub mod addin;
 pub mod cookies;

programs/voter-stake-registry/tests/test_basic.rs

@@ -1,4 +1,5 @@
 use anchor_spl::token::TokenAccount;
+use assert_custom_on_chain_error::AssertCustomOnChainErr;
 use mplx_staking_states::{
     error::VsrError,
     state::{LockupKind, LockupPeriod},

programs/voter-stake-registry/tests/test_change_delegate.rs

@@ -1,4 +1,5 @@
 use anchor_spl::token::TokenAccount;
+use assert_custom_on_chain_error::AssertCustomOnChainErr;
 use mplx_staking_states::{
     error::VsrError,
     state::{LockupKind, LockupPeriod},

programs/voter-stake-registry/tests/test_deposit_constant.rs

@@ -1,4 +1,5 @@
 use anchor_spl::token::TokenAccount;
+use assert_custom_on_chain_error::AssertCustomOnChainErr;
 use mplx_staking_states::{
     error::VsrError,
     state::{LockupKind, LockupPeriod},

programs/voter-stake-registry/tests/test_deposit_no_locking.rs

@@ -1,4 +1,5 @@
 use anchor_spl::token::TokenAccount;
+use assert_custom_on_chain_error::AssertCustomOnChainErr;
 use mplx_staking_states::{
     error::VsrError,
     state::{LockupKind, LockupPeriod},

programs/voter-stake-registry/tests/test_extend_deposit.rs

@@ -1,4 +1,5 @@
 use anchor_spl::token::TokenAccount;
+use assert_custom_on_chain_error::AssertCustomOnChainErr;
 use mplx_staking_states::{
     error::VsrError,
     state::{LockupKind, LockupPeriod},

programs/voter-stake-registry/tests/test_lockup.rs

@@ -1,4 +1,5 @@
 use anchor_spl::token::TokenAccount;
+use assert_custom_on_chain_error::AssertCustomOnChainErr;
 use mplx_staking_states::{
     error::VsrError,
     state::{LockupKind, LockupPeriod},