Allow manager to burns shares and distribute to depositors

programs/drift_vaults/src/error.rs

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

programs/drift_vaults/src/instructions/manager_burn_shares.rs

@@ -0,0 +1,54 @@
+use crate::constraints::{is_manager_for_vault, is_user_for_vault, is_user_stats_for_vault};
+use crate::AccountMapProvider;
+use crate::{Vault, WithdrawUnit};
+use anchor_lang::prelude::*;
+use drift::instructions::optional_accounts::AccountMaps;
+use drift::state::user::User;
+
+pub fn manager_burn_shares<'info>(
+    ctx: Context<'_, '_, '_, 'info, ManagerBurnShares<'info>>,
+    burn_amount: u64,
+    burn_unit: WithdrawUnit,
+) -> Result<()> {
+    let clock = &Clock::get()?;
+    let mut vault = ctx.accounts.vault.load_mut()?;
+    let now = clock.unix_timestamp;
+
+    let user = ctx.accounts.drift_user.load()?;
+    let spot_market_index = vault.spot_market_index;
+
+    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)?;
+
+    vault.manager_burn_shares(burn_amount, burn_unit, vault_equity, now)?;
+
+    Ok(())
+}
+
+#[derive(Accounts)]
+pub struct ManagerBurnShares<'info> {
+    #[account(
+        mut,
+        constraint = is_manager_for_vault(&vault, &manager)?
+    )]
+    pub vault: AccountLoader<'info, Vault>,
+    pub manager: Signer<'info>,
+    #[account(
+        constraint = is_user_stats_for_vault(&vault, &drift_user_stats)?
+    )]
+    /// CHECK: checked in drift cpi
+    pub drift_user_stats: AccountInfo<'info>,
+    #[account(
+        constraint = is_user_for_vault(&vault, &drift_user.key())?
+    )]
+    /// CHECK: checked in drift cpi
+    pub drift_user: AccountLoader<'info, User>,
+    /// CHECK: checked in drift cpi
+    pub drift_state: AccountInfo<'info>,
+}

programs/drift_vaults/src/instructions/mod.rs

@@ -7,6 +7,7 @@ pub use initialize_insurance_fund_stake::*;
 pub use initialize_vault::*;
 pub use initialize_vault_depositor::*;
 pub use liquidate::*;
+pub use manager_burn_shares::*;
 pub use manager_cancel_withdraw_request::*;
 pub use manager_deposit::*;
 pub use manager_request_withdraw::*;
@@ -29,6 +30,7 @@ mod initialize_insurance_fund_stake;
 mod initialize_vault;
 mod initialize_vault_depositor;
 mod liquidate;
+mod manager_burn_shares;
 mod manager_cancel_withdraw_request;
 mod manager_deposit;
 mod manager_request_withdraw;

programs/drift_vaults/src/lib.rs

@@ -110,6 +110,14 @@ pub mod drift_vaults {
         instructions::manager_withdraw(ctx)
     }
 
+    pub fn manager_burn_shares<'info>(
+        ctx: Context<'_, '_, '_, 'info, ManagerBurnShares<'info>>,
+        burn_amount: u64,
+        burn_unit: WithdrawUnit,
+    ) -> Result<()> {
+        instructions::manager_burn_shares(ctx, burn_amount, burn_unit)
+    }
+
     pub fn apply_profit_share<'info>(
         ctx: Context<'_, '_, '_, 'info, ApplyProfitShare<'info>>,
     ) -> Result<()> {

programs/drift_vaults/src/state/events.rs

@@ -49,3 +49,16 @@ impl Default for VaultDepositorAction {
         VaultDepositorAction::Deposit
     }
 }
+
+#[event]
+#[derive(Default)]
+pub struct BurnVaultSharesRecord {
+    pub ts: i64,
+    pub vault: Pubkey,
+    pub depositor_authority: Pubkey,
+    pub amount: u64,
+    pub spot_market_index: u16,
+    pub vault_equity: u64,
+    pub total_vault_shares_before: u128,
+    pub total_vault_shares_after: u128,
+}

programs/drift_vaults/src/state/vault.rs

@@ -1,4 +1,5 @@
 use crate::events::{VaultDepositorAction, VaultDepositorRecord};
+use crate::state::BurnVaultSharesRecord;
 use crate::{validate, Size, VaultDepositor, WithdrawUnit};
 
 use crate::constants::TIME_FOR_LIQUIDATION;
@@ -500,4 +501,53 @@ impl Vault {
         self.liquidation_delegate = Pubkey::default();
         self.liquidation_start_ts = 0;
     }
+
+    pub fn manager_burn_shares(
+        &mut self,
+        burn_amount: u64,
+        burn_unit: WithdrawUnit,
+        vault_equity: u64,
+        now: i64,
+    ) -> Result<()> {
+        self.apply_rebase(vault_equity)?;
+        self.apply_management_fee(vault_equity, now)?;
+
+        let (burn_value, new_total_shares) = burn_unit.get_transfer_value_with_new_total(
+            burn_amount,
+            vault_equity,
+            self.user_shares,
+            self.total_shares,
+        )?;
+
+        validate!(
+            new_total_shares < self.total_shares,
+            ErrorCode::InvalidBurnSharesAmount,
+            "new_total_shares > total_shares"
+        )?;
+
+        let total_vault_shares_before = self.total_shares;
+
+        msg!(
+            "burn_value: {}, total_shares: {} -> {}",
+            burn_value,
+            total_vault_shares_before,
+            new_total_shares
+        );
+        self.total_shares = new_total_shares;
+
+        emit!(BurnVaultSharesRecord {
+            ts: now,
+            vault: self.pubkey,
+            depositor_authority: self.manager,
+            amount: burn_value,
+            spot_market_index: self.spot_market_index,
+            vault_equity,
+            total_vault_shares_before,
+            total_vault_shares_after: self.total_shares,
+        });
+
+        self.apply_rebase(vault_equity)?;
+
+        Ok(())
+    }
 }

programs/drift_vaults/src/state/withdraw_unit.rs

@@ -52,6 +52,53 @@ impl WithdrawUnit {
         }
     }
 
+    /// returns the amount and shares transfered as a result of transfering `amount` and changing
+    /// `total_shares` to `new_total_shares`
+    pub fn get_transfer_value_with_new_total(
+        &self,
+        amount: u64,
+        vault_equity: u64,
+        user_shares: u128,
+        total_shares: u128,
+    ) -> VaultResult<(u64, u128)> {
+        match self {
+            WithdrawUnit::Token => {
+                let transfer_value = amount;
+                let new_total_shares = user_shares
+                    .safe_mul(vault_equity.cast::<u128>()?)?
+                    .safe_div(vault_equity.safe_sub(amount)?.cast::<u128>()?)?;
+
+                Ok((transfer_value, new_total_shares))
+            }
+            WithdrawUnit::Shares => {
+                let transfer_value = depositor_shares_to_vault_amount(
+                    amount.cast::<u128>()?,
+                    total_shares,
+                    vault_equity,
+                )?;
+
+                let new_total_shares = user_shares
+                    .safe_mul(vault_equity.cast::<u128>()?)?
+                    .safe_div(vault_equity.safe_sub(transfer_value)?.cast::<u128>()?)?;
+                Ok((transfer_value, new_total_shares))
+            }
+            WithdrawUnit::SharesPercent => {
+                let manager_shares = total_shares.safe_sub(user_shares)?;
+                let transfer_shares =
+                    WithdrawUnit::get_shares_from_percent(amount.cast()?, manager_shares)?;
+                let transfer_value = depositor_shares_to_vault_amount(
+                    transfer_shares.cast::<u128>()?,
+                    total_shares,
+                    vault_equity,
+                )?;
+                let new_total_shares = user_shares
+                    .safe_mul(vault_equity.cast::<u128>()?)?
+                    .safe_div(vault_equity.safe_sub(transfer_value)?.cast::<u128>()?)?;
+                Ok((transfer_value, new_total_shares))
+            }
+        }
+    }
+
     fn get_shares_from_percent(percent: u128, shares: u128) -> VaultResult<u128> {
         validate!(
             percent <= MAX_WITHDRAW_PERCENT,
@@ -61,3 +108,56 @@ impl WithdrawUnit {
         Ok(shares)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_get_value_and_shares_with_new_total() {
+        // want shares to transfer 100 tokens
+        let withdraw_unit = WithdrawUnit::Token;
+        let result = withdraw_unit.get_transfer_value_with_new_total(
+            50_000_000, // 50 tokens
+            1000_000_000,
+            90_000_000, // manager shares = 100 - 90 = 10
+            100_000_000,
+        );
+        assert!(result.is_ok());
+        let (value, new_total_shares) = result.unwrap();
+        assert_eq!(value, 50_000_000);
+        // before: manager owns (100-90)/100 * 1000 = 100
+        // after: manager owns (94.736842-90)/94.736842 * 1000 = 49.999998
+        assert_eq!(new_total_shares, 94736842);
+
+        // want shares to transfer 5 shares
+        let withdraw_unit = WithdrawUnit::Shares;
+        let result = withdraw_unit.get_transfer_value_with_new_total(
+            5_000_000, // 5 tokens
+            1000_000_000,
+            90_000_000, // manager shares = 100 - 90 = 10
+            100_000_000,
+        );
+        assert!(result.is_ok());
+        let (value, new_total_shares) = result.unwrap();
+        assert_eq!(value, 50_000_000);
+        // before: manager owns (100-90)/100 * 1000 = 100
+        // after: manager owns (94.736842-90)/94.736842 * 1000 = 49.999998
+        assert_eq!(new_total_shares, 94736842);
+
+        // want shares to transfer 50% of shares
+        let withdraw_unit = WithdrawUnit::SharesPercent;
+        let result = withdraw_unit.get_transfer_value_with_new_total(
+            500_000, // 0.50
+            1000_000_000,
+            90_000_000, // manager shares = 100 - 90 = 10
+            100_000_000,
+        );
+        assert!(result.is_ok());
+        let (value, new_total_shares) = result.unwrap();
+        assert_eq!(value, 50_000_000);
+        // before: manager owns (100-90)/100 * 1000 = 100
+        // after: manager owns (94.736842-90)/94.736842 * 1000 = 49.999998
+        assert_eq!(new_total_shares, 94736842);
+    }
+}