add if staker fucntions

programs/drift_vaults/src/cpi.rs

@@ -29,3 +29,25 @@ pub trait UpdateUserMarginTradingEnabledCPI {
 pub trait TokenTransferCPI {
     fn token_transfer(&self, amount: u64) -> Result<()>;
 }
+
+pub trait InitializeInsuranceFundStakeCPI {
+    fn drift_initialize_insurance_fund_stake(&self, market_index: u16) -> Result<()>;
+}
+
+pub trait AddInsuranceFundStakeCPI {
+    fn drift_add_insurance_fund_stake(&self, market_index: u16, amount: u64) -> Result<()>;
+}
+
+pub trait RequestRemoveInsuranceFundStakeCPI {
+    fn drift_request_remove_insurance_fund_stake(
+        &self,
+        market_index: u16,
+        amount: u64,
+    ) -> Result<()>;
+
+    fn drift_cancel_request_remove_insurance_fund_stake(&self, market_index: u16) -> Result<()>;
+}
+
+pub trait RemoveInsuranceFundStakeCPI {
+    fn drift_remove_insurance_fund_stake(&self, market_index: u16) -> Result<()>;
+}

programs/drift_vaults/src/error.rs

@@ -51,6 +51,8 @@ pub enum ErrorCode {
     InvalidVaultDeposit,
     #[msg("OngoingLiquidation")]
     OngoingLiquidation,
+    #[msg("InvalidInsuranceFundSize")]
+    InvalidAddInsuranceFundSize,
 }
 
 impl From<DriftErrorCode> for ErrorCode {

programs/drift_vaults/src/instructions/add_insurance_fund_stake.rs

@@ -0,0 +1,137 @@
+use crate::constraints::{is_manager_for_vault, is_spot_market_for_vault, is_user_stats_for_vault};
+use crate::cpi::{AddInsuranceFundStakeCPI, WithdrawCPI};
+use crate::error::ErrorCode;
+use crate::Vault;
+use crate::{declare_vault_seeds, implement_withdraw, validate, AccountMapProvider};
+use anchor_lang::prelude::*;
+use anchor_spl::token::{Token, TokenAccount};
+use drift::cpi::accounts::AddInsuranceFundStake as DriftAddInsuranceFundStake;
+use drift::cpi::accounts::Withdraw as DriftWithdraw;
+use drift::instructions::optional_accounts::AccountMaps;
+use drift::program::Drift;
+use drift::state::insurance_fund_stake::InsuranceFundStake;
+use drift::state::spot_market::SpotMarket;
+use drift::state::user::User;
+
+pub fn add_insurance_fund_stake<'info>(
+    ctx: Context<'_, '_, '_, 'info, AddInsuranceFundStake<'info>>,
+    market_index: u16,
+    amount: u64,
+) -> Result<()> {
+    let clock = &Clock::get()?;
+    let vault = ctx.accounts.vault.load()?;
+
+    let user = ctx.accounts.drift_user.load()?;
+    let spot_market_index = vault.spot_market_index;
+
+    validate!(!vault.in_liquidation(), ErrorCode::OngoingLiquidation)?;
+
+    let AccountMaps {
+        perp_market_map,
+        spot_market_map,
+        mut oracle_map,
+    } = ctx.load_maps(clock.slot, Some(spot_market_index))?;
+
+    let vault_equity =
+        vault.calculate_equity(&user, &perp_market_map, &spot_market_map, &mut oracle_map)?;
+
+    validate!(
+        amount <= vault_equity,
+        ErrorCode::InvalidAddInsuranceFundSize
+    )?;
+
+    ctx.drift_withdraw(amount)?;
+
+    ctx.drift_add_insurance_fund_stake(market_index, amount)?;
+
+    Ok(())
+}
+
+#[derive(Accounts)]
+#[instruction(market_index: u16)]
+pub struct AddInsuranceFundStake<'info> {
+    #[account(
+        constraint = is_manager_for_vault(&vault, &manager)?,
+    )]
+    pub vault: AccountLoader<'info, Vault>,
+    pub manager: Signer<'info>,
+    #[account(
+        mut,
+        token::authority = vault.key(),
+        token::mint = vault_token_account.mint
+    )]
+    pub vault_token_account: Box<Account<'info, TokenAccount>>,
+    #[account(
+        constraint = is_spot_market_for_vault(&vault, &drift_spot_market, market_index)?,
+    )]
+    pub drift_spot_market: AccountLoader<'info, SpotMarket>,
+    #[account(
+        mut,
+        seeds = [b"insurance_fund_stake", vault.key().as_ref(), market_index.to_le_bytes().as_ref()],
+        bump
+    )]
+    pub insurance_fund_stake: AccountLoader<'info, InsuranceFundStake>,
+    #[account(
+        mut,
+        seeds = [b"spot_market_vault".as_ref(), market_index.to_le_bytes().as_ref()],
+        bump,
+    )]
+    pub drift_spot_market_vault: Box<Account<'info, TokenAccount>>,
+    #[account(
+        mut,
+        seeds = [b"insurance_fund_vault".as_ref(), market_index.to_le_bytes().as_ref()],
+        bump,
+    )]
+    pub insurance_fund_vault: Box<Account<'info, TokenAccount>>,
+    #[account(
+        mut,
+        constraint = is_user_stats_for_vault(&vault, &drift_user_stats)?
+    )]
+    /// CHECK: checked in drift cpi
+    pub drift_user: AccountLoader<'info, User>,
+    /// CHECK: checked in drift cpi
+    pub drift_user_stats: AccountInfo<'info>,
+    /// CHECK: checked in drift cpi
+    pub drift_signer: AccountInfo<'info>,
+    /// CHECK: checked in drift cpi
+    pub drift_state: AccountInfo<'info>,
+    pub drift_program: Program<'info, Drift>,
+    pub token_program: Program<'info, Token>,
+}
+
+impl<'info> AddInsuranceFundStakeCPI for Context<'_, '_, '_, 'info, AddInsuranceFundStake<'info>> {
+    fn drift_add_insurance_fund_stake(&self, market_index: u16, amount: u64) -> Result<()> {
+        declare_vault_seeds!(self.accounts.vault, seeds);
+
+        let cpi_accounts = DriftAddInsuranceFundStake {
+            state: self.accounts.drift_state.clone(),
+            spot_market: self.accounts.drift_spot_market.to_account_info().clone(),
+            insurance_fund_stake: self.accounts.insurance_fund_stake.to_account_info().clone(),
+            user_stats: self.accounts.drift_user_stats.clone(),
+            authority: self.accounts.vault.to_account_info().clone(), // sign?
+            spot_market_vault: self
+                .accounts
+                .drift_spot_market_vault
+                .to_account_info()
+                .clone(),
+            insurance_fund_vault: self.accounts.insurance_fund_stake.to_account_info().clone(),
+            drift_signer: self.accounts.drift_signer.to_account_info().clone(),
+            user_token_account: self.accounts.vault_token_account.to_account_info().clone(),
+            token_program: self.accounts.token_program.to_account_info().clone(),
+        };
+
+        let drift_program = self.accounts.drift_program.to_account_info().clone();
+        let cpi_context = CpiContext::new_with_signer(drift_program, cpi_accounts, seeds)
+            .with_remaining_accounts(self.remaining_accounts.into());
+        drift::cpi::add_insurance_fund_stake(cpi_context, market_index, amount)?;
+
+        Ok(())
+    }
+}
+
+impl<'info> WithdrawCPI for Context<'_, '_, '_, 'info, AddInsuranceFundStake<'info>> {
+    fn drift_withdraw(&self, amount: u64) -> Result<()> {
+        implement_withdraw!(self, amount);
+        Ok(())
+    }
+}

programs/drift_vaults/src/instructions/constraints.rs

@@ -1,6 +1,7 @@
 use crate::{Vault, VaultDepositor};
 
 use anchor_lang::prelude::*;
+use drift::state::spot_market::SpotMarket;
 
 pub fn is_vault_for_vault_depositor(
     vault_depositor: &AccountLoader<VaultDepositor>,
@@ -43,3 +44,14 @@ pub fn is_user_stats_for_vault(
 ) -> anchor_lang::Result<bool> {
     Ok(vault_depositor.load()?.user_stats.eq(user_stats.key))
 }
+
+pub fn is_spot_market_for_vault(
+    vault_depositor: &AccountLoader<Vault>,
+    drift_spot_market: &AccountLoader<SpotMarket>,
+    market_index: u16,
+) -> anchor_lang::Result<bool> {
+    Ok(
+        (&vault_depositor.load()?.spot_market_index).eq(&drift_spot_market.load()?.market_index)
+            && (&vault_depositor.load()?.spot_market_index).eq(&market_index),
+    )
+}

programs/drift_vaults/src/instructions/deposit.rs

@@ -3,7 +3,7 @@ use crate::constraints::{
 };
 use crate::cpi::{DepositCPI, TokenTransferCPI};
 use crate::error::ErrorCode;
-use crate::{declare_vault_seeds, validate, AccountMapProvider};
+use crate::{declare_vault_seeds, implement_deposit, validate, AccountMapProvider};
 use crate::{Vault, VaultDepositor};
 use anchor_lang::prelude::*;
 use anchor_spl::token::{self, Token, TokenAccount, Transfer};
@@ -110,27 +110,7 @@ impl<'info> TokenTransferCPI for Context<'_, '_, '_, 'info, Deposit<'info>> {
 
 impl<'info> DepositCPI for Context<'_, '_, '_, 'info, Deposit<'info>> {
     fn drift_deposit(&self, amount: u64) -> Result<()> {
-        declare_vault_seeds!(self.accounts.vault, seeds);
-        let spot_market_index = self.accounts.vault.load()?.spot_market_index;
-
-        let cpi_program = self.accounts.drift_program.to_account_info().clone();
-        let cpi_accounts = DriftDeposit {
-            state: self.accounts.drift_state.clone(),
-            user: self.accounts.drift_user.to_account_info().clone(),
-            user_stats: self.accounts.drift_user_stats.clone(),
-            authority: self.accounts.vault.to_account_info().clone(),
-            spot_market_vault: self
-                .accounts
-                .drift_spot_market_vault
-                .to_account_info()
-                .clone(),
-            user_token_account: self.accounts.vault_token_account.to_account_info().clone(),
-            token_program: self.accounts.token_program.to_account_info().clone(),
-        };
-        let cpi_context = CpiContext::new_with_signer(cpi_program, cpi_accounts, seeds)
-            .with_remaining_accounts(self.remaining_accounts.into());
-        drift::cpi::deposit(cpi_context, spot_market_index, amount, false)?;
-
+        implement_deposit!(self, amount);
         Ok(())
     }
 }

programs/drift_vaults/src/instructions/initialize_insurance_fund_stake.rs

@@ -0,0 +1,78 @@
+use crate::constraints::{is_manager_for_vault, is_spot_market_for_vault, is_user_stats_for_vault};
+use crate::cpi::InitializeInsuranceFundStakeCPI;
+use crate::declare_vault_seeds;
+use crate::Vault;
+use anchor_lang::prelude::*;
+use drift::cpi::accounts::InitializeInsuranceFundStake as DriftInitializeInsuranceFundStake;
+use drift::program::Drift;
+use drift::state::insurance_fund_stake::InsuranceFundStake;
+use drift::state::spot_market::SpotMarket;
+
+pub fn initialize_insurance_fund_stake<'info>(
+    ctx: Context<'_, '_, '_, 'info, InitializeInsuranceFundStake<'info>>,
+    market_index: u16,
+) -> Result<()> {
+    ctx.drift_initialize_insurance_fund_stake(market_index)?;
+    Ok(())
+}
+
+#[derive(Accounts)]
+#[instruction(market_index: u16)]
+pub struct InitializeInsuranceFundStake<'info> {
+    #[account(
+        mut,
+        constraint = is_manager_for_vault(&vault, &manager)?,
+    )]
+    pub vault: AccountLoader<'info, Vault>,
+    pub manager: Signer<'info>,
+    #[account(mut)]
+    pub payer: Signer<'info>,
+    pub rent: Sysvar<'info, Rent>,
+    pub system_program: Program<'info, System>,
+
+    #[account(
+        constraint = is_spot_market_for_vault(&vault, &drift_spot_market, market_index)?,
+    )]
+    pub drift_spot_market: AccountLoader<'info, SpotMarket>,
+    #[account(
+        mut,
+        seeds = [b"insurance_fund_stake", vault.key().as_ref(), market_index.to_le_bytes().as_ref()],
+        bump
+    )]
+    pub insurance_fund_stake: AccountLoader<'info, InsuranceFundStake>,
+    #[account(
+        mut,
+        constraint = is_user_stats_for_vault(&vault, &drift_user_stats)?
+    )]
+    /// CHECK: checked in drift cpi
+    pub drift_user_stats: AccountInfo<'info>,
+    /// CHECK: checked in drift cpi
+    pub drift_state: AccountInfo<'info>,
+    pub drift_program: Program<'info, Drift>,
+}
+
+impl<'info> InitializeInsuranceFundStakeCPI
+    for Context<'_, '_, '_, 'info, InitializeInsuranceFundStake<'info>>
+{
+    fn drift_initialize_insurance_fund_stake(&self, market_index: u16) -> Result<()> {
+        declare_vault_seeds!(self.accounts.vault, seeds);
+
+        let cpi_accounts = DriftInitializeInsuranceFundStake {
+            spot_market: self.accounts.drift_spot_market.to_account_info().clone(),
+            insurance_fund_stake: self.accounts.insurance_fund_stake.to_account_info().clone(),
+            user_stats: self.accounts.drift_user_stats.clone(),
+            state: self.accounts.drift_state.clone(),
+            authority: self.accounts.vault.to_account_info().clone(), // sign?
+            payer: self.accounts.payer.to_account_info().clone(),
+            rent: self.accounts.rent.to_account_info().clone(),
+            system_program: self.accounts.system_program.to_account_info().clone(),
+        };
+
+        let drift_program = self.accounts.drift_program.to_account_info().clone();
+        let cpi_context = CpiContext::new_with_signer(drift_program, cpi_accounts, seeds)
+            .with_remaining_accounts(self.remaining_accounts.into());
+        drift::cpi::initialize_insurance_fund_stake(cpi_context, market_index)?;
+
+        Ok(())
+    }
+}

programs/drift_vaults/src/instructions/mod.rs

@@ -1,14 +1,18 @@
+pub use add_insurance_fund_stake::*;
 pub use apply_profit_share::*;
 pub use cancel_withdraw_request::*;
 pub use deposit::*;
 pub use force_withdraw::*;
+pub use initialize_insurance_fund_stake::*;
 pub use initialize_vault::*;
 pub use initialize_vault_depositor::*;
 pub use liquidate::*;
 pub use manager_cancel_withdraw_request::*;
 pub use manager_deposit::*;
 pub use manager_request_withdraw::*;
 pub use manager_withdraw::*;
+pub use remove_insurance_fund_stake::*;
+pub use request_remove_insurance_fund_stake::*;
 
 pub use request_withdraw::*;
 pub use reset_delegate::*;
@@ -17,18 +21,22 @@ pub use update_margin_trading_enabled::*;
 pub use update_vault::*;
 pub use withdraw::*;
 
+mod add_insurance_fund_stake;
 mod apply_profit_share;
 mod cancel_withdraw_request;
 pub mod constraints;
 mod deposit;
 mod force_withdraw;
+mod initialize_insurance_fund_stake;
 mod initialize_vault;
 mod initialize_vault_depositor;
 mod liquidate;
 mod manager_cancel_withdraw_request;
 mod manager_deposit;
 mod manager_request_withdraw;
 mod manager_withdraw;
+mod remove_insurance_fund_stake;
+mod request_remove_insurance_fund_stake;
 mod request_withdraw;
 mod reset_delegate;
 mod update_delegate;