Feat: quote and ixs for prefund_withdraw_stake

Cargo.toml

@@ -13,7 +13,7 @@ members = [
 [workspace.dependencies]
 anyhow = "^1.0"
 bincode = "^1.0"
-borsh = "^1"
+borsh = ">=0.9"
 clap = "^4"
 itertools = ">=0.1"
 jupiter-amm-interface = "~0.4.0"
@@ -33,7 +33,7 @@ spl-associated-token-account = { version = ">=1", features = ["no-entrypoint"] }
 # but 2.0.1 patches those breaking changes.
 # Ideally, this vers requirement should be "anything >=1.0 except 2.0.0"
 # but theres no way to express that in cargo dep vers syntax
-spl-stake-pool = { version = ">=1.0", features = ["no-entrypoint"] } 
+spl-stake-pool = { git = "https://github.com/igneous-labs/stake-pool.git", branch = "loosen-deps", features = ["no-entrypoint"] } 
 spl-token = ">=3.0"
 thiserror = "^1.0"
 tokio = "^1.0"

common/src/withdraw_stake.rs

@@ -37,6 +37,10 @@ impl WithdrawStakeQuote {
         self.lamports_out == 0
     }
 
+    pub fn is_rent_exempt(&self) -> bool {
+        self.lamports_out.saturating_sub(self.lamports_staked) >= STAKE_ACCOUNT_RENT_EXEMPT_LAMPORTS
+    }
+
     pub fn from_lamports_and_voter(stake_acc_lamports: u64, voter: Pubkey) -> Self {
         let (lamports_out, lamports_staked) =
             if stake_acc_lamports > STAKE_ACCOUNT_RENT_EXEMPT_LAMPORTS {

jup_interface/src/pool_pair/common.rs

@@ -245,28 +245,40 @@ pub fn first_avail_prefund_quote<W: WithdrawStake + ?Sized, D: DepositStake + ?S
     let withdraw_quote_iter = withdraw_from.withdraw_stake_quote_iter_dyn(withdraw_amount);
     for wsq in withdraw_quote_iter {
         let wsq = prefund_transform_wsq(wsq);
-        let mut wsq_after_prefund = wsq;
-        wsq_after_prefund.lamports_out = wsq.lamports_out.saturating_sub(prefund_split_lamports);
-        wsq_after_prefund.lamports_staked =
-            wsq.lamports_staked.saturating_sub(prefund_split_lamports);
-        if wsq_after_prefund.is_zero_out() {
+        let wsq_after_repaying_prefund = wsq_post_prefund_repay(wsq, prefund_split_lamports);
+        if !wsq_after_repaying_prefund.is_rent_exempt() {
             continue;
         }
-        let dsq = deposit_to.get_deposit_stake_quote(wsq_after_prefund)?;
+        let dsq = deposit_to.get_deposit_stake_quote(wsq_after_repaying_prefund)?;
+
         if !dsq.is_zero_out() {
             return Ok((wsq, dsq));
         }
     }
     Err(SwapViaStakeQuoteErr::NoRouteFound)
 }
 
-/// Since we're prefunding bridge stake with the rent, we need to add it to the output stake account
-fn prefund_transform_wsq(mut wsq: WithdrawStakeQuote) -> WithdrawStakeQuote {
+/// Returns the state of the stake account with rent-exempt lamports prefunded.
+///
+/// This is basically adding [`STAKE_ACCOUNT_RENT_EXEMPT_LAMPORTS`] non-staked lamports to the stake account
+/// while treating its original [`WithdrawStakeQuote::lamports_out`] as the [`WithdrawStakeQuote::lamports_staked`]
+pub fn prefund_transform_wsq(mut wsq: WithdrawStakeQuote) -> WithdrawStakeQuote {
     wsq.lamports_staked = wsq.lamports_out;
     wsq.lamports_out += STAKE_ACCOUNT_RENT_EXEMPT_LAMPORTS;
     wsq
 }
 
+/// Returns the state of the stake account after `prefund_split_lamports` is split
+/// from it to repay the prefund flash loan
+pub fn wsq_post_prefund_repay(
+    mut wsq: WithdrawStakeQuote,
+    prefund_split_lamports: u64,
+) -> WithdrawStakeQuote {
+    wsq.lamports_out = wsq.lamports_out.saturating_sub(prefund_split_lamports);
+    wsq.lamports_staked = wsq.lamports_staked.saturating_sub(prefund_split_lamports);
+    wsq
+}
+
 /// Calculates approximate fees charged in terms of out token given known
 /// amt_after_fee in terms of out token and fee ratio
 fn approx_fees_charged_out_token(amt_after_fee: u64, fee_num: u64, fee_denom: u64) -> Result<u64> {

stakedex_sdk/src/lib.rs

@@ -3,32 +3,35 @@ use std::collections::HashSet;
 use anyhow::{anyhow, Result};
 use itertools::Itertools;
 use jupiter_amm_interface::{
-    AccountMap, Amm, AmmContext, KeyedAccount, Quote, QuoteParams, SwapParams,
+    AccountMap, Amm, AmmContext, KeyedAccount, Quote, QuoteParams, SwapMode, SwapParams,
 };
 use lazy_static::lazy_static;
 use sanctum_lst_list::{PoolInfo, SanctumLst};
-use solana_sdk::{instruction::Instruction, pubkey::Pubkey, system_program};
+use solana_sdk::{instruction::Instruction, pubkey::Pubkey, stake, system_program, sysvar};
 use spl_token::native_mint;
 use stakedex_interface::{
     DepositStakeKeys, PrefundSwapViaStakeIxArgs, PrefundSwapViaStakeKeys,
     PrefundWithdrawStakeIxArgs, PrefundWithdrawStakeKeys, StakeWrappedSolIxArgs,
     StakeWrappedSolKeys, SwapViaStakeArgs, WithdrawWrappedSolIxArgs, WithdrawWrappedSolKeys,
 };
 use stakedex_jup_interface::{
-    manual_concat_get_account_metas, prefund_get_account_metas, quote_pool_pair, DepositSolWrapper,
-    DepositWithdrawSolWrapper, OneWayPoolPair, PrefundRepayParams, TwoWayPoolPair,
+    manual_concat_get_account_metas, prefund_get_account_metas, prefund_transform_wsq,
+    quote_pool_pair, wsq_post_prefund_repay, DepositSolWrapper, DepositWithdrawSolWrapper,
+    OneWayPoolPair, PrefundRepayParams, TwoWayPoolPair,
 };
 use stakedex_lido::LidoStakedex;
 use stakedex_marinade::MarinadeStakedex;
 use stakedex_sdk_common::{
-    find_fee_token_acc, lido_state, marinade_state, msol,
+    find_bridge_stake, find_fee_token_acc, lido_state, marinade_state, msol,
+    slumdog_stake_create_with_seed,
     stakedex_program::{self, WSOL_FEE_TOKEN_ACCOUNT_ID},
-    stsol, unstake_it_program, wsol, wsol_bridge_in, BaseStakePoolAmm, DepositSol, DepositStake,
-    DepositStakeInfo, DepositStakeQuote, InitFromKeyedAccount, WithdrawSol, WithdrawStake,
-    WithdrawStakeQuote, DEPOSIT_STAKE_DST_TOKEN_ACCOUNT_INDEX,
+    stsol, unstake_it_pool, unstake_it_program, wsol, wsol_bridge_in, BaseStakePoolAmm, DepositSol,
+    DepositStake, DepositStakeInfo, DepositStakeQuote, InitFromKeyedAccount, WithdrawSol,
+    WithdrawStake, WithdrawStakeQuote, WithdrawStakeQuoteErr,
+    DEPOSIT_STAKE_DST_TOKEN_ACCOUNT_INDEX,
 };
 use stakedex_spl_stake_pool::{SplStakePoolStakedex, SplStakePoolStakedexWithWithdrawSol};
-use stakedex_unstake_it::{UnstakeItStakedex, UnstakeItStakedexPrefund};
+use stakedex_unstake_it::{find_stake_account_record, UnstakeItStakedex, UnstakeItStakedexPrefund};
 
 pub use sanctum_lst_list::SanctumLstList;
 pub use stakedex_interface::ID as stakedex_program_id;
@@ -300,6 +303,8 @@ impl Stakedex {
         )
     }
 
+    /// Returns [PrefundWithdrawStake, DepositStake] instructions
+    /// to create a SwapViaStake action without running into CPI account write-lock limits
     pub fn manual_concat_prefund_swap_via_stake_ixs(
         &self,
         swap_params: &SwapParams,
@@ -564,6 +569,94 @@ impl Stakedex {
         Ok(ix)
     }
 
+    /// `output_mint` can either be
+    /// - the desired vote account of the stake account to withdraw, or
+    /// - the stake program to indicate no preference of voter
+    pub fn quote_prefund_withdraw_stake(
+        &self,
+        quote_params: &QuoteParams,
+    ) -> Result<WithdrawStakeQuote> {
+        let prefund_split_lamports = self.prefund_repay_params().prefund_split_lamports()?;
+        let withdraw_from = self
+            .get_withdraw_stake_pool(&quote_params.input_mint)
+            .ok_or_else(|| anyhow!("pool not found for input mint {}", quote_params.input_mint))?;
+        if !withdraw_from.can_accept_stake_withdrawals() {
+            return Err(WithdrawStakeQuoteErr::CannotAcceptStakeWithdrawals.into());
+        }
+        let itr = withdraw_from.withdraw_stake_quote_iter_dyn(quote_params.amount);
+
+        for wsq in itr {
+            let wsq = prefund_transform_wsq(wsq);
+            let wsq_after_repaying_prefund = wsq_post_prefund_repay(wsq, prefund_split_lamports);
+            if wsq_after_repaying_prefund.is_rent_exempt()
+                && (quote_params.output_mint == wsq.voter
+                    || quote_params.output_mint == stake::program::ID)
+            {
+                return Ok(wsq_after_repaying_prefund);
+            }
+        }
+
+        Err(anyhow!("No route found"))
+    }
+
+    /// `destination_mint` can either be
+    /// - the desired vote account of the stake account to withdraw, or
+    /// - the stake program to indicate no preference of voter
+    pub fn prefund_withdraw_stake_ix(
+        &self,
+        SwapParams {
+            in_amount,
+            source_mint,
+            destination_mint,
+            source_token_account,
+            token_transfer_authority,
+            ..
+        }: &SwapParams,
+        bridge_stake_seed: u32,
+    ) -> Result<Instruction> {
+        let withdraw_from = self
+            .get_withdraw_stake_pool(source_mint)
+            .ok_or_else(|| anyhow!("pool not found for input mint {}", source_mint))?;
+        let wsq = self.quote_prefund_withdraw_stake(&QuoteParams {
+            amount: *in_amount,
+            input_mint: *source_mint,
+            output_mint: *destination_mint,
+            swap_mode: SwapMode::ExactIn,
+        })?;
+        let bridge_stake_seed_le_bytes = bridge_stake_seed.to_le_bytes();
+        let bridge_stake =
+            find_bridge_stake(token_transfer_authority, &bridge_stake_seed_le_bytes).0;
+        let slumdog_stake = slumdog_stake_create_with_seed(&bridge_stake)?;
+        let mut ix = stakedex_interface::prefund_withdraw_stake_ix(
+            PrefundWithdrawStakeKeys {
+                user: *token_transfer_authority,
+                src_token_from: *source_token_account,
+                bridge_stake,
+                src_token_mint: *source_mint,
+                prefunder: stakedex_program::PREFUNDER_ID,
+                slumdog_stake,
+                unstakeit_program: unstake_it_program::ID,
+                unstake_pool: unstake_it_pool::ID,
+                pool_sol_reserves: unstake_it_program::SOL_RESERVES_ID,
+                unstake_fee: unstake_it_program::FEE_ID,
+                slumdog_stake_acc_record: find_stake_account_record(&slumdog_stake).0,
+                unstake_protocol_fee: unstake_it_program::PROTOCOL_FEE_ID,
+                unstake_protocol_fee_dest: self.prefund_repay_params().protocol_fee_dest,
+                clock: sysvar::clock::ID,
+                stake_program: stake::program::ID,
+                system_program: system_program::ID,
+            },
+            PrefundWithdrawStakeIxArgs {
+                args: SwapViaStakeArgs {
+                    amount: *in_amount,
+                    bridge_stake_seed,
+                },
+            },
+        )?;
+        ix.accounts.extend(withdraw_from.virtual_ix(&wsq)?.accounts);
+        Ok(ix)
+    }
+
     /// Creates all possible Amms from the underlying available Stakedexes
     pub fn get_amms(self) -> Vec<Box<dyn Amm + Send + Sync>> {
         #[derive(Clone)]