Withdraw: Add pre-health check

programs/mango-v4/src/instructions/token_deposit.rs

@@ -94,6 +94,9 @@ pub fn token_deposit(ctx: Context<TokenDeposit>, amount: u64) -> Result<()> {
     //
     // Health computation
     //
+    // Since depositing can only increase health, we can skip the usual pre-health computation.
+    // Also, TokenDeposit is one of the rare instructions that is allowed even during being_liquidated.
+    //
     if !account.fixed.is_in_health_region() {
         let retriever =
             new_fixed_order_account_retriever(ctx.remaining_accounts, &account.borrow())?;

programs/mango-v4/src/instructions/token_withdraw.rs

@@ -10,7 +10,6 @@ use fixed::types::I80F48;
 use crate::logs::{
     LoanOriginationFeeInstruction, TokenBalanceLog, WithdrawLoanOriginationFeeLog, WithdrawLog,
 };
-use crate::state::new_fixed_order_account_retriever;
 use crate::util::checked_math as cm;
 
 #[derive(Accounts)]
@@ -61,91 +60,107 @@ pub fn token_withdraw(ctx: Context<TokenWithdraw>, amount: u64, allow_borrow: bo
     let group = ctx.accounts.group.load()?;
     let token_index = ctx.accounts.bank.load()?.token_index;
 
-    // Get the account's position for that token index
+    // Create the account's position for that token index
     let mut account = ctx.accounts.account.load_mut()?;
+    let (_, raw_token_index, _) = account.ensure_token_position(token_index)?;
 
-    let (position, raw_token_index, _active_token_index) =
+    // Health check _after_ the token position is guaranteed to exist
-        account.ensure_token_position(token_index)?;
+    let pre_health_opt = if !account.fixed.is_in_health_region() {
-
+        let retriever =
-    // The bank will also be passed in remainingAccounts. Use an explicit scope
+            new_fixed_order_account_retriever(ctx.remaining_accounts, &account.borrow())?;
-    // to drop the &mut before we borrow it immutably again later.
+        let health_cache =
-    let (position_is_active, amount_i80f48, loan_origination_fee) = {
+            new_health_cache(&account.borrow(), &retriever).context("pre-withdraw init health")?;
-        let mut bank = ctx.accounts.bank.load_mut()?;
+        let pre_health = health_cache.health(HealthType::Init);
-        let native_position = position.native(&bank);
+        msg!("pre_health: {}", pre_health);
-
+        account
-        // Handle amount special case for withdrawing everything
+            .fixed
-        let amount = if amount == u64::MAX && !allow_borrow {
+            .maybe_recover_from_being_liquidated(pre_health);
-            if native_position.is_positive() {
-                // TODO: This rounding may mean that if we deposit and immediately withdraw
-                //       we can't withdraw the full amount!
-                native_position.floor().to_num::<u64>()
-            } else {
-                return Ok(());
-            }
-        } else {
-            amount
-        };
-
         require!(
-            allow_borrow || amount <= native_position,
+            !account.fixed.being_liquidated(),
-            MangoError::SomeError
+            MangoError::BeingLiquidated
         );
-
+        Some((health_cache, pre_health))
-        let amount_i80f48 = I80F48::from(amount);
+    } else {
-
+        None
-        // Update the bank and position
-        let (position_is_active, loan_origination_fee) =
-            bank.withdraw_with_fee(position, amount_i80f48)?;
-
-        // Provide a readable error message in case the vault doesn't have enough tokens
-        if ctx.accounts.vault.amount < amount {
-            return err!(MangoError::InsufficentBankVaultFunds).with_context(|| {
-                format!(
-                    "bank vault does not have enough tokens, need {} but have {}",
-                    amount, ctx.accounts.vault.amount
-                )
-            });
-        }
-
-        // Transfer the actual tokens
-        let group_seeds = group_seeds!(group);
-        token::transfer(
-            ctx.accounts.transfer_ctx().with_signer(&[group_seeds]),
-            amount,
-        )?;
-
-        (position_is_active, amount_i80f48, loan_origination_fee)
     };
 
-    let indexed_position = position.indexed_position;
+    let mut bank = ctx.accounts.bank.load_mut()?;
-    let bank = ctx.accounts.bank.load()?;
+    let position = account.token_position_mut_by_raw_index(raw_token_index);
-    let oracle_price = bank.oracle_price(&AccountInfoRef::borrow(ctx.accounts.oracle.as_ref())?)?;
+    let native_position = position.native(&bank);
 
-    // Update the net deposits - adjust by price so different tokens are on the same basis (in USD terms)
+    // Handle amount special case for withdrawing everything
-    let amount_usd = cm!(amount_i80f48 * oracle_price).to_num::<i64>();
+    let amount = if amount == u64::MAX && !allow_borrow {
-    account.fixed.net_deposits = cm!(account.fixed.net_deposits - amount_usd);
+        if native_position.is_positive() {
+            // TODO: This rounding may mean that if we deposit and immediately withdraw
+            //       we can't withdraw the full amount!
+            native_position.floor().to_num::<u64>()
+        } else {
+            return Ok(());
+        }
+    } else {
+        amount
+    };
+
+    require!(
+        allow_borrow || amount <= native_position,
+        MangoError::SomeError
+    );
+
+    let amount_i80f48 = I80F48::from(amount);
+
+    // Update the bank and position
+    let (position_is_active, loan_origination_fee) =
+        bank.withdraw_with_fee(position, amount_i80f48)?;
+
+    // Provide a readable error message in case the vault doesn't have enough tokens
+    if ctx.accounts.vault.amount < amount {
+        return err!(MangoError::InsufficentBankVaultFunds).with_context(|| {
+            format!(
+                "bank vault does not have enough tokens, need {} but have {}",
+                amount, ctx.accounts.vault.amount
+            )
+        });
+    }
+
+    // Transfer the actual tokens
+    let group_seeds = group_seeds!(group);
+    token::transfer(
+        ctx.accounts.transfer_ctx().with_signer(&[group_seeds]),
+        amount,
+    )?;
+
+    let native_position_after = position.native(&bank);
+    let oracle_price = bank.oracle_price(&AccountInfoRef::borrow(ctx.accounts.oracle.as_ref())?)?;
 
     emit!(TokenBalanceLog {
         mango_group: ctx.accounts.group.key(),
         mango_account: ctx.accounts.account.key(),
         token_index,
-        indexed_position: indexed_position.to_bits(),
+        indexed_position: position.indexed_position.to_bits(),
         deposit_index: bank.deposit_index.to_bits(),
         borrow_index: bank.borrow_index.to_bits(),
         price: oracle_price.to_bits(),
     });
 
+    // Update the net deposits - adjust by price so different tokens are on the same basis (in USD terms)
+    let amount_usd = cm!(amount_i80f48 * oracle_price).to_num::<i64>();
+    account.fixed.net_deposits = cm!(account.fixed.net_deposits - amount_usd);
+
     //
     // Health check
     //
-    if !account.fixed.is_in_health_region() {
+    if let Some((mut health_cache, pre_health)) = pre_health_opt {
-        let retriever =
+        health_cache
-            new_fixed_order_account_retriever(ctx.remaining_accounts, &account.borrow())?;
+            .adjust_token_balance(token_index, cm!(native_position_after - native_position))?;
-        let health = compute_health(&account.borrow(), HealthType::Init, &retriever)
+        let post_health = health_cache.health(HealthType::Init);
-            .context("post-withdraw init health")?;
+        msg!("post_health: {}", post_health);
-        msg!("health: {}", health);
+        require!(
-        require!(health >= 0, MangoError::HealthMustBePositive);
+            post_health >= 0 || post_health > pre_health,
-        account.fixed.maybe_recover_from_being_liquidated(health);
+            MangoError::HealthMustBePositive
+        );
+        account
+            .fixed
+            .maybe_recover_from_being_liquidated(post_health);
     }
 
     //