From a75183dda7db444ba4ec493f49ee7d4d78248483 Mon Sep 17 00:00:00 2001 From: Jayant Krishnamurthy Date: Wed, 8 Nov 2023 09:03:52 -0800 Subject: [PATCH] Split accounts test (#258) * vesting tests * more tests * better invariant tests --- staking/programs/staking/src/error.rs | 7 +- staking/programs/staking/src/lib.rs | 20 +- .../programs/staking/src/state/positions.rs | 89 --------- staking/programs/staking/src/state/vesting.rs | 184 +++++++++++++++++- staking/target/idl/staking.json | 5 + staking/target/types/staking.ts | 10 + 6 files changed, 210 insertions(+), 105 deletions(-) diff --git a/staking/programs/staking/src/error.rs b/staking/programs/staking/src/error.rs index e4b59393..95e21b44 100644 --- a/staking/programs/staking/src/error.rs +++ b/staking/programs/staking/src/error.rs @@ -69,8 +69,11 @@ pub enum ErrorCode { SplitZeroTokens, #[msg("Can't split more tokens than are in the account")] // 6032 SplitTooManyTokens, - #[msg("Sanity check failed")] // 6033 + #[msg("Can't split a token account with staking positions. Unstake your tokens first.")] + // 6033 + SplitWithStake, + #[msg("Sanity check failed")] // 6034 SanityCheckFailed, - #[msg("Other")] //6031 + #[msg("Other")] //6035 Other, } diff --git a/staking/programs/staking/src/lib.rs b/staking/programs/staking/src/lib.rs index b3b34e97..2d384858 100644 --- a/staking/programs/staking/src/lib.rs +++ b/staking/programs/staking/src/lib.rs @@ -607,6 +607,14 @@ pub mod staking { config.unlocking_duration, )?; + // Check that there aren't any positions (i.e., staked tokens) in the source account. + // This check allows us to create an empty positions account on behalf of the recipient and + // not worry about moving positions from the source account to the new account. + require!( + source_stake_account_metadata.next_index == 0, + ErrorCode::SplitWithStake + ); + require!(split_request.amount > 0, ErrorCode::SplitZeroTokens); require!( split_request.amount < source_stake_account_custody.amount, @@ -626,18 +634,6 @@ pub mod staking { source_stake_account_metadata.set_lock(source_vesting_account); new_stake_account_metadata.set_lock(new_vesting_account); - // Split positions - source_stake_account_positions.split( - new_stake_account_positions, - &mut source_stake_account_metadata.next_index, - &mut new_stake_account_metadata.next_index, - remaining_amount, - split_request.amount, - source_stake_account_custody.amount, - current_epoch, - config.unlocking_duration, - )?; - { transfer( diff --git a/staking/programs/staking/src/state/positions.rs b/staking/programs/staking/src/state/positions.rs index 466994fa..0c9ea672 100644 --- a/staking/programs/staking/src/state/positions.rs +++ b/staking/programs/staking/src/state/positions.rs @@ -114,94 +114,6 @@ impl PositionData { } Ok(exposure) } - - pub fn split( - &mut self, - dest_position_data: &mut PositionData, - src_next_index: &mut u8, - dest_next_index: &mut u8, - remaining_amount: u64, - transferred_amount: u64, - total_amount: u64, - current_epoch: u64, - unlocking_duration: u8, - ) -> Result<()> { - require!( - transferred_amount - .checked_add(remaining_amount) - .ok_or(ErrorCode::Other)? - == total_amount, - ErrorCode::SanityCheckFailed - ); - let governance_exposure = - self.get_target_exposure(&Target::VOTING, current_epoch, unlocking_duration)?; - require!( - governance_exposure <= total_amount, - ErrorCode::SanityCheckFailed - ); - - if remaining_amount < governance_exposure { - // We need to transfer some positions over to the new account - let mut excess_governance_exposure = - governance_exposure.saturating_sub(remaining_amount); - - while excess_governance_exposure > 0 && *src_next_index > 0 { - let index = TryInto::::try_into(*src_next_index - 1) - .map_err(|_| ErrorCode::GenericOverflow)?; - match self.read_position(index)? { - Some(position) => { - match position.get_current_position(current_epoch, unlocking_duration)? { - PositionState::UNLOCKED => self.make_none(index, src_next_index)?, - PositionState::LOCKING - | PositionState::LOCKED - | PositionState::PREUNLOCKING - | PositionState::UNLOCKING => { - if excess_governance_exposure < position.amount { - // We need to split the position - self.write_position( - index, - &Position { - amount: position - .amount - .saturating_sub(excess_governance_exposure), - ..position - }, - )?; - - let new_position = Position { - amount: excess_governance_exposure, - ..position - }; - - let new_index = - dest_position_data.reserve_new_index(dest_next_index)?; - dest_position_data.write_position(new_index, &new_position)?; - - excess_governance_exposure = 0; - } else { - // We need to transfer the whole position - let new_index = - dest_position_data.reserve_new_index(dest_next_index)?; - dest_position_data.write_position(new_index, &position)?; - - self.make_none(index, src_next_index)?; - excess_governance_exposure = - excess_governance_exposure.saturating_sub(position.amount); - } - } - } - } - None => { - // This should never happen - return Err(error!(ErrorCode::SanityCheckFailed)); - } - } - } - } - - - Ok(()) - } } pub trait TryBorsh { @@ -526,7 +438,6 @@ pub mod tests { } } - #[quickcheck] fn prop(input: Vec) -> bool { let mut position_data = PositionData::default(); diff --git a/staking/programs/staking/src/state/vesting.rs b/staking/programs/staking/src/state/vesting.rs index a9fb9898..29e2b7f0 100644 --- a/staking/programs/staking/src/state/vesting.rs +++ b/staking/programs/staking/src/state/vesting.rs @@ -20,7 +20,7 @@ use { /// Represents how a given initial balance vests over time /// It is unit-less, but units must be consistent #[repr(u8)] -#[derive(AnchorSerialize, AnchorDeserialize, Debug, Clone, Copy, BorshSchema)] +#[derive(AnchorSerialize, AnchorDeserialize, Debug, Clone, Copy, BorshSchema, PartialEq)] pub enum VestingSchedule { /// No vesting, i.e. balance is fully vested at all time FullyVested, @@ -247,6 +247,10 @@ impl VestingSchedule { == total_amount, ErrorCode::SanityCheckFailed ); + // Note that the arithmetic below may lose precision. The calculations round down + // the number of vesting tokens for both of the new accounts, which means that splitting + // may vest some dust (1 of the smallest decimal point) of PYTH for both the source and + // destination accounts. match self { VestingSchedule::FullyVested => { Ok((VestingSchedule::FullyVested, VestingSchedule::FullyVested)) @@ -299,8 +303,14 @@ pub mod tests { use { crate::state::vesting::{ VestingEvent, - VestingSchedule, + VestingSchedule::{ + self, + PeriodicVesting, + PeriodicVestingAfterListing, + }, }, + quickcheck::TestResult, + quickcheck_macros::quickcheck, std::convert::TryInto, }; @@ -629,4 +639,174 @@ pub mod tests { None ); } + + const START_TIMESTAMP: i64 = 10; + const PERIOD_DURATION: u64 = 5; + const NUM_PERIODS: u64 = 4; + + #[quickcheck] + fn test_split_props(transferred: u64, total: u64, initial_balance: u64) -> TestResult { + if transferred > total || total == 0 { + return TestResult::discard(); + } + let received = total - transferred; + + let schedule = VestingSchedule::FullyVested; + let (remaining_schedule, transferred_schedule) = schedule + .split_vesting_schedule(received, transferred, total) + .unwrap(); + + assert_eq!(remaining_schedule, VestingSchedule::FullyVested); + assert_eq!(transferred_schedule, VestingSchedule::FullyVested); + + let schedule = PeriodicVesting { + initial_balance, + // all of these fields should be preserved in the result + start_date: START_TIMESTAMP, + period_duration: PERIOD_DURATION, + num_periods: NUM_PERIODS, + }; + let (remaining_schedule, transferred_schedule) = schedule + .split_vesting_schedule(received, transferred, total) + .unwrap(); + + match (remaining_schedule, transferred_schedule) { + ( + PeriodicVesting { + initial_balance: r, .. + }, + PeriodicVesting { + initial_balance: t, .. + }, + ) => { + let sum = r + t; + assert!(initial_balance.saturating_sub(2) <= sum && sum <= initial_balance); + } + _ => { + panic!("Test failed"); + } + } + + let schedule = PeriodicVestingAfterListing { + initial_balance, + // all of these fields should be preserved in the result + period_duration: PERIOD_DURATION, + num_periods: NUM_PERIODS, + }; + let (remaining_schedule, transferred_schedule) = schedule + .split_vesting_schedule(received, transferred, total) + .unwrap(); + + match (remaining_schedule, transferred_schedule) { + ( + PeriodicVestingAfterListing { + initial_balance: r, .. + }, + PeriodicVestingAfterListing { + initial_balance: t, .. + }, + ) => { + let sum = r + t; + assert!(initial_balance.saturating_sub(2) <= sum && sum <= initial_balance); + } + _ => { + panic!("Test failed"); + } + } + + for timestamp in 0..(START_TIMESTAMP + (PERIOD_DURATION * NUM_PERIODS + 1) as i64) { + let initial_unvested = schedule + .get_unvested_balance(timestamp, Some(START_TIMESTAMP)) + .unwrap(); + let remaining_unvested = remaining_schedule + .get_unvested_balance(timestamp, Some(START_TIMESTAMP)) + .unwrap(); + let transferred_unvested = transferred_schedule + .get_unvested_balance(timestamp, Some(START_TIMESTAMP)) + .unwrap(); + + assert!( + initial_unvested.saturating_sub(2) <= (remaining_unvested + transferred_unvested) + && (remaining_unvested + transferred_unvested) <= initial_unvested + ); + + if initial_unvested <= total { + assert!(transferred_unvested <= transferred); + assert!(remaining_unvested <= received); + } + } + + TestResult::passed() + } + + fn test_split_helper( + transferred: u64, + total: u64, + initial_balance: u64, + expected_remaining: u64, + expected_transferred: u64, + ) { + let schedule = PeriodicVesting { + initial_balance, + // all of these fields should be preserved in the result + start_date: 203, + period_duration: 100, + num_periods: 12, + }; + let (remaining_schedule, transferred_schedule) = schedule + .split_vesting_schedule(total - transferred, transferred, total) + .unwrap(); + + let t = PeriodicVesting { + initial_balance: expected_transferred, + start_date: 203, + period_duration: 100, + num_periods: 12, + }; + let r = PeriodicVesting { + initial_balance: expected_remaining, + start_date: 203, + period_duration: 100, + num_periods: 12, + }; + + assert_eq!(remaining_schedule, r); + assert_eq!(transferred_schedule, t); + + let schedule = PeriodicVestingAfterListing { + initial_balance, + period_duration: 100, + num_periods: 12, + }; + let (remaining_schedule, transferred_schedule) = schedule + .split_vesting_schedule(total - transferred, transferred, total) + .unwrap(); + + let t = PeriodicVestingAfterListing { + initial_balance: expected_transferred, + period_duration: 100, + num_periods: 12, + }; + let r = PeriodicVestingAfterListing { + initial_balance: expected_remaining, + period_duration: 100, + num_periods: 12, + }; + + assert_eq!(remaining_schedule, r); + assert_eq!(transferred_schedule, t); + } + + #[test] + fn test_split() { + test_split_helper(10, 100, 100, 90, 10); + test_split_helper(10, 1000, 100, 99, 1); + test_split_helper(1, 1000, 100, 99, 0); + + test_split_helper(10, 10, 1000, 0, 1000); + test_split_helper(9, 10, 1000, 100, 900); + test_split_helper(10, 100, 1000, 900, 100); + + test_split_helper(1, 3, 1000, 666, 333); + } } diff --git a/staking/target/idl/staking.json b/staking/target/idl/staking.json index d086ac1a..e1201267 100644 --- a/staking/target/idl/staking.json +++ b/staking/target/idl/staking.json @@ -1916,6 +1916,11 @@ "code": 6034, "name": "Other", "msg": "Other" + }, + { + "code": 6035, + "name": "SplitWithStake", + "msg": "Can't split a token account with staking positions. Unstake your tokens first." } ] } \ No newline at end of file diff --git a/staking/target/types/staking.ts b/staking/target/types/staking.ts index 8fcb4954..b84b6354 100644 --- a/staking/target/types/staking.ts +++ b/staking/target/types/staking.ts @@ -1916,6 +1916,11 @@ export type Staking = { "code": 6034, "name": "Other", "msg": "Other" + }, + { + "code": 6035, + "name": "SplitWithStake", + "msg": "Can't split a token account with staking positions. Unstake your tokens first." } ] }; @@ -3838,6 +3843,11 @@ export const IDL: Staking = { "code": 6034, "name": "Other", "msg": "Other" + }, + { + "code": 6035, + "name": "SplitWithStake", + "msg": "Can't split a token account with staking positions. Unstake your tokens first." } ] };