program: add percent shares withdraw unit (#21)

* program: add percent shares withdraw unit * add ts enum
drift-labs · Aug 7, 2023 · 2c7f37b · 2c7f37b
1 parent de3cfe5
commit 2c7f37b
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 1 deletion.
diff --git a/programs/drift_vaults/src/error.rs b/programs/drift_vaults/src/error.rs
@@ -45,6 +45,8 @@ pub enum ErrorCode {
     PermissionedVault,
     #[msg("WithdrawInProgress")]
     WithdrawInProgress,
+    #[msg("SharesPercentTooLarge")]
+    SharesPercentTooLarge,
 }
 
 impl From<DriftErrorCode> for ErrorCode {

diff --git a/programs/drift_vaults/src/state/vault.rs b/programs/drift_vaults/src/state/vault.rs
@@ -161,6 +161,11 @@ impl Vault {
         ))
     }
 
+    pub fn get_manager_shares(&self) -> VaultResult<u128> {
+        let manager_shares = self.total_shares.safe_sub(self.user_shares)?;
+        Ok(manager_shares)
+    }
+
     pub fn apply_rebase(&mut self, vault_equity: u64) -> Result<()> {
         if vault_equity != 0 && vault_equity.cast::<u128>()? < self.total_shares {
             let (expo_diff, rebase_divisor) =
@@ -305,6 +310,16 @@ impl Vault {
                         .min(vault_equity);
                 (n_tokens, n_shares)
             }
+            WithdrawUnit::SharesPercent => {
+                let n_shares: u128 = WithdrawUnit::get_shares_from_percent(
+                    withdraw_amount.cast()?,
+                    self.get_manager_shares()?,
+                )?;
+                let n_tokens: u64 =
+                    depositor_shares_to_vault_amount(n_shares, self.total_shares, vault_equity)?
+                        .min(vault_equity);
+                (n_tokens, n_shares)
+            }
         };
 
         self.total_withdraws = self.total_withdraws.saturating_add(n_tokens);

diff --git a/programs/drift_vaults/src/state/vault_depositor.rs b/programs/drift_vaults/src/state/vault_depositor.rs
@@ -1,4 +1,4 @@
-use crate::error::ErrorCode;
+use crate::error::{ErrorCode, VaultResult};
 use crate::Size;
 use anchor_lang::prelude::*;
 use borsh::{BorshDeserialize, BorshSerialize};
@@ -399,6 +399,14 @@ impl VaultDepositor {
                         .min(vault_equity);
                 (withdraw_value, n_shares)
             }
+            WithdrawUnit::SharesPercent => {
+                let n_shares =
+                    WithdrawUnit::get_shares_from_percent(withdraw_amount, self.vault_shares)?;
+                let withdraw_value: u64 =
+                    depositor_shares_to_vault_amount(n_shares, vault.total_shares, vault_equity)?
+                        .min(vault_equity);
+                (withdraw_value, n_shares)
+            }
         };
 
         validate!(
@@ -687,6 +695,19 @@ impl VaultDepositor {
 pub enum WithdrawUnit {
     Shares,
     Token,
+    SharesPercent,
+}
+
+const MAX_WITHDRAW_PERCENT: u128 = 100_000;
+impl WithdrawUnit {
+    pub fn get_shares_from_percent(percent: u128, shares: u128) -> VaultResult<u128> {
+        validate!(
+            percent <= MAX_WITHDRAW_PERCENT,
+            ErrorCode::SharesPercentTooLarge
+        )?;
+        let shares = shares.safe_mul(percent)?.safe_div(MAX_WITHDRAW_PERCENT)?;
+        Ok(shares)
+    }
 }
 
 #[cfg(test)]

diff --git a/ts/sdk/src/types/drift_vaults.ts b/ts/sdk/src/types/drift_vaults.ts
@@ -1004,6 +1004,9 @@ export type DriftVaults = {
 					},
 					{
 						name: 'Token';
+					},
+					{
+						name: 'SharesPercent';
 					}
 				];
 			};
@@ -1208,6 +1211,16 @@ export type DriftVaults = {
 			code: 6017;
 			name: 'PermissionedVault';
 			msg: 'PermissionedVault';
+		},
+		{
+			code: 6018;
+			name: 'WithdrawInProgress';
+			msg: 'WithdrawInProgress';
+		},
+		{
+			code: 6019;
+			name: 'SharesPercentTooLarge';
+			msg: 'SharesPercentTooLarge';
 		}
 	];
 };
@@ -2219,6 +2232,9 @@ export const IDL: DriftVaults = {
 					{
 						name: 'Token',
 					},
+					{
+						name: 'SharesPercent',
+					},
 				],
 			},
 		},
@@ -2423,5 +2439,15 @@ export const IDL: DriftVaults = {
 			name: 'PermissionedVault',
 			msg: 'PermissionedVault',
 		},
+		{
+			code: 6018,
+			name: 'WithdrawInProgress',
+			msg: 'WithdrawInProgress',
+		},
+		{
+			code: 6019,
+			name: 'SharesPercentTooLarge',
+			msg: 'SharesPercentTooLarge',
+		},
 	],
 };
diff --git a/ts/sdk/src/types/types.ts b/ts/sdk/src/types/types.ts
@@ -10,6 +10,7 @@ export const VAULT_PROGRAM_ID = new PublicKey(
 export class WithdrawUnit {
 	static readonly SHARES = { shares: {} };
 	static readonly TOKEN = { token: {} };
+	static readonly SHARES_PERCENT = { sharesPercent: {} };
 }
 
 export type Vault = {