feat: shorten vesting for some accounts

staking/integration-tests/src/solana/instructions.rs

@@ -46,6 +46,34 @@ pub fn create_account(
     account.pubkey()
 }
 
+/// Creates an empty, program owned, rent exempt account at `address`.
+///
+/// Unlike `create_account` this doesn't go through the system program, so it works for addresses
+/// whose keypair we don't have. This is needed to test instructions that only accept a hardcoded
+/// set of stake accounts.
+pub fn create_account_at(
+    svm: &mut litesvm::LiteSVM,
+    address: Pubkey,
+    size: usize,
+    owner: Pubkey,
+) -> Pubkey {
+    let lamports = svm.minimum_balance_for_rent_exemption(size);
+
+    svm.set_account(
+        address,
+        solana_sdk::account::Account {
+            lamports,
+            data: vec![0; size],
+            owner,
+            executable: false,
+            rent_epoch: 0,
+        },
+    )
+    .unwrap();
+
+    address
+}
+
 pub fn create_token_account(
     svm: &mut litesvm::LiteSVM,
     payer: &Keypair,

staking/integration-tests/src/staking/helper_functions.rs

@@ -10,13 +10,15 @@ use {
         solana::instructions::{
             airdrop_spl,
             create_account,
+            create_account_at,
         },
         utils::constants::STAKED_TOKENS,
     },
     solana_sdk::{
         pubkey::Pubkey,
         signature::Keypair,
     },
+    staking::state::vesting::VestingSchedule,
 };
 
 
@@ -34,7 +36,56 @@ pub fn initialize_new_stake_account(
         staking::ID,
     );
 
-    create_stake_account(svm, payer, pyth_token_mint, stake_account_positions).unwrap();
+    initialize_stake_account(
+        svm,
+        payer,
+        pyth_token_mint,
+        join_dao,
+        airdrop,
+        stake_account_positions,
+        VestingSchedule::FullyVested,
+    )
+}
+
+/// Same as `initialize_new_stake_account`, but the positions account is placed at
+/// `stake_account_positions` instead of at a fresh keypair, and the lock is configurable.
+pub fn initialize_new_stake_account_at(
+    svm: &mut litesvm::LiteSVM,
+    payer: &Keypair,
+    pyth_token_mint: &Keypair,
+    join_dao: bool,
+    airdrop: bool,
+    stake_account_positions: Pubkey,
+    lock: VestingSchedule,
+) -> Pubkey {
+    create_account_at(
+        svm,
+        stake_account_positions,
+        staking::state::positions::PositionData::LEN,
+        staking::ID,
+    );
+
+    initialize_stake_account(
+        svm,
+        payer,
+        pyth_token_mint,
+        join_dao,
+        airdrop,
+        stake_account_positions,
+        lock,
+    )
+}
+
+fn initialize_stake_account(
+    svm: &mut litesvm::LiteSVM,
+    payer: &Keypair,
+    pyth_token_mint: &Keypair,
+    join_dao: bool,
+    airdrop: bool,
+    stake_account_positions: Pubkey,
+    lock: VestingSchedule,
+) -> Pubkey {
+    create_stake_account(svm, payer, pyth_token_mint, stake_account_positions, lock).unwrap();
 
     if join_dao {
         join_dao_llc(svm, payer, stake_account_positions).unwrap();

staking/integration-tests/src/staking/instructions.rs

@@ -41,6 +41,7 @@ use {
     staking::state::{
         global_config::GlobalConfig,
         positions::TargetWithParameters,
+        vesting::VestingSchedule,
         voter_weight_record::VoterWeightAction,
     },
 };
@@ -257,6 +258,7 @@ pub fn create_stake_account(
     payer: &Keypair,
     pyth_token_mint: &Keypair,
     stake_account_positions: Pubkey,
+    lock: VestingSchedule,
 ) -> TransactionResult {
     let stake_account_metadata = get_stake_account_metadata_address(stake_account_positions);
     let stake_account_custody = get_stake_account_custody_address(stake_account_positions);
@@ -265,7 +267,7 @@ pub fn create_stake_account(
 
     let create_stake_account_data = staking::instruction::CreateStakeAccount {
         owner: payer.pubkey(),
-        lock:  staking::state::vesting::VestingSchedule::FullyVested,
+        lock,
     };
     let create_stake_account_accs = staking::accounts::CreateStakeAccount {
         payer: payer.pubkey(),
@@ -545,6 +547,38 @@ pub fn transfer_account(
     svm.send_transaction(tx)
 }
 
+/// The instruction is permissionless, `payer` only pays the fee.
+/// `stake_account_metadata_override` lets a test pass a metadata account that doesn't belong to
+/// `stake_account_positions`; when `None` the correct PDA is derived.
+pub fn shorten_vesting_schedule(
+    svm: &mut litesvm::LiteSVM,
+    payer: &Keypair,
+    stake_account_positions: Pubkey,
+    stake_account_metadata_override: Option<Pubkey>,
+) -> TransactionResult {
+    let stake_account_metadata = stake_account_metadata_override
+        .unwrap_or_else(|| get_stake_account_metadata_address(stake_account_positions));
+
+    let accs = staking::accounts::ShortenVestingSchedule {
+        stake_account_positions,
+        stake_account_metadata,
+    };
+
+    let ix = Instruction::new_with_bytes(
+        staking::ID,
+        &staking::instruction::ShortenVestingSchedule {}.data(),
+        accs.to_account_metas(None),
+    );
+    let tx = Transaction::new_signed_with_payer(
+        &[ix],
+        Some(&payer.pubkey()),
+        &[&payer],
+        svm.latest_blockhash(),
+    );
+
+    svm.send_transaction(tx)
+}
+
 pub fn create_voter_record(
     svm: &mut litesvm::LiteSVM,
     payer: &Keypair,

staking/integration-tests/tests/shorten_vesting_schedule.rs

@@ -0,0 +1,184 @@
+use {
+    anchor_lang::error::ErrorCode,
+    integration_tests::{
+        assert_anchor_program_error,
+        setup::{
+            setup,
+            SetupProps,
+            SetupResult,
+        },
+        solana::utils::fetch_account_data,
+        staking::{
+            helper_functions::{
+                initialize_new_stake_account,
+                initialize_new_stake_account_at,
+            },
+            instructions::shorten_vesting_schedule,
+            pda::get_stake_account_metadata_address,
+        },
+    },
+    solana_sdk::{
+        native_token::LAMPORTS_PER_SOL,
+        pubkey::Pubkey,
+        signature::Keypair,
+        signer::Signer,
+    },
+    staking::{
+        error::ErrorCode as StakingError,
+        state::{
+            stake_account::StakeAccountMetadataV2,
+            vesting::VestingSchedule,
+        },
+        ADDRESSES_TO_SHORTEN_VESTING_SCHEDULE,
+    },
+};
+
+const INITIAL_BALANCE: u64 = 1_000_000;
+const START_DATE: i64 = 100;
+const PERIOD_DURATION: u64 = 3600;
+const NUM_PERIODS: u64 = 4;
+
+fn periodic_vesting(num_periods: u64) -> VestingSchedule {
+    VestingSchedule::PeriodicVesting {
+        initial_balance: INITIAL_BALANCE,
+        start_date: START_DATE,
+        period_duration: PERIOD_DURATION,
+        num_periods,
+    }
+}
+
+fn fetch_lock(svm: &mut litesvm::LiteSVM, stake_account_positions: Pubkey) -> VestingSchedule {
+    let metadata: StakeAccountMetadataV2 = fetch_account_data(
+        svm,
+        &get_stake_account_metadata_address(stake_account_positions),
+    );
+    metadata.lock
+}
+
+#[test]
+fn test_shorten_vesting_schedule() {
+    let SetupResult {
+        mut svm,
+        payer,
+        pyth_token_mint,
+        publisher_keypair: _,
+        pool_data_pubkey: _,
+        reward_program_authority: _,
+        maybe_publisher_index: _,
+    } = setup(SetupProps {
+        init_config:            true,
+        init_target:            true,
+        init_mint:              true,
+        init_pool_data:         true,
+        init_publishers:        true,
+        reward_amount_override: None,
+    });
+
+    let owner = Keypair::new();
+    svm.airdrop(&owner.pubkey(), LAMPORTS_PER_SOL).unwrap();
+
+    // A stake account at one of the hardcoded addresses, i.e. eligible for shortening.
+    let whitelisted_stake_account_positions = initialize_new_stake_account_at(
+        &mut svm,
+        &owner,
+        &pyth_token_mint,
+        true,
+        true,
+        ADDRESSES_TO_SHORTEN_VESTING_SCHEDULE[0],
+        periodic_vesting(NUM_PERIODS),
+    );
+
+    // A stake account at an arbitrary address, i.e. not eligible for shortening.
+    let other_stake_account_positions =
+        initialize_new_stake_account(&mut svm, &owner, &pyth_token_mint, true, true);
+
+    // Wrong stake account address: the account isn't in the hardcoded list.
+    assert_anchor_program_error!(
+        shorten_vesting_schedule(&mut svm, &payer, other_stake_account_positions, None),
+        StakingError::UnauthorizedVestingScheduleShortening,
+        0
+    );
+
+    // Wrong metadata account: the metadata of another stake account doesn't match the seeds.
+    assert_anchor_program_error!(
+        shorten_vesting_schedule(
+            &mut svm,
+            &payer,
+            whitelisted_stake_account_positions,
+            Some(get_stake_account_metadata_address(
+                other_stake_account_positions
+            )),
+        ),
+        ErrorCode::ConstraintSeeds,
+        0
+    );
+
+    // Neither failed instruction touched the vesting schedules.
+    assert_eq!(
+        fetch_lock(&mut svm, whitelisted_stake_account_positions),
+        periodic_vesting(NUM_PERIODS)
+    );
+    assert_eq!(
+        fetch_lock(&mut svm, other_stake_account_positions),
+        VestingSchedule::FullyVested
+    );
+
+    // Happy path: only `num_periods` changes, and it becomes 1.
+    shorten_vesting_schedule(&mut svm, &payer, whitelisted_stake_account_positions, None).unwrap();
+    assert_eq!(
+        fetch_lock(&mut svm, whitelisted_stake_account_positions),
+        periodic_vesting(1)
+    );
+
+    // Idempotence: shortening an already shortened schedule is a no-op.
+    svm.expire_blockhash();
+    shorten_vesting_schedule(&mut svm, &payer, whitelisted_stake_account_positions, None).unwrap();
+    assert_eq!(
+        fetch_lock(&mut svm, whitelisted_stake_account_positions),
+        periodic_vesting(1)
+    );
+}
+
+#[test]
+fn test_shorten_vesting_schedule_is_a_noop_for_other_schedules() {
+    let SetupResult {
+        mut svm,
+        payer,
+        pyth_token_mint,
+        publisher_keypair: _,
+        pool_data_pubkey: _,
+        reward_program_authority: _,
+        maybe_publisher_index: _,
+    } = setup(SetupProps {
+        init_config:            true,
+        init_target:            true,
+        init_mint:              true,
+        init_pool_data:         true,
+        init_publishers:        true,
+        reward_amount_override: None,
+    });
+
+    let owner = Keypair::new();
+    svm.airdrop(&owner.pubkey(), LAMPORTS_PER_SOL).unwrap();
+
+    let lock = VestingSchedule::PeriodicVestingAfterListing {
+        initial_balance: INITIAL_BALANCE,
+        period_duration: PERIOD_DURATION,
+        num_periods:     NUM_PERIODS,
+    };
+
+    let stake_account_positions = initialize_new_stake_account_at(
+        &mut svm,
+        &owner,
+        &pyth_token_mint,
+        true,
+        true,
+        ADDRESSES_TO_SHORTEN_VESTING_SCHEDULE[1],
+        lock,
+    );
+
+    shorten_vesting_schedule(&mut svm, &payer, stake_account_positions, None).unwrap();
+
+    // Only `PeriodicVesting` is shortened, everything else is left alone.
+    assert_eq!(fetch_lock(&mut svm, stake_account_positions), lock);
+}

staking/programs/staking/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pyth-staking-program"
-version = "2.1.0"
+version = "2.2.0"
 description = "Created with Anchor"
 edition = "2018"