Split accounts test (#258)

* vesting tests

* more tests

* better invariant tests

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,
 }

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(

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::<usize>::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<DataOperation>) -> bool {
         let mut position_data = PositionData::default();

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);
+    }
 }

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."
     }
   ]
 }

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."
     }
   ]
 };