bigz/improve-isolated-contract-tier-checks

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Features
 
+- program: make isolated perp contract tier more ergonomic ([#913](https://github.com/drift-labs/protocol-v2/pull/913))
+
 ### Fixes
 
 ### Breaking

programs/drift/src/math/margin.rs

@@ -385,14 +385,14 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
                     )?;
 
                     calculation.add_spot_liability()?;
-                    calculation.update_with_isolated_liability(
+                    calculation.update_with_spot_isolated_liability(
                         spot_market.asset_tier == AssetTier::Isolated,
                     );
                 }
                 Ordering::Equal => {
                     if spot_position.has_open_order() {
                         calculation.add_spot_liability()?;
-                        calculation.update_with_isolated_liability(
+                        calculation.update_with_spot_isolated_liability(
                             spot_market.asset_tier == AssetTier::Isolated,
                         );
                     }
@@ -488,8 +488,9 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
 
         if has_perp_liability {
             calculation.add_perp_liability()?;
-            calculation
-                .update_with_isolated_liability(market.contract_tier == ContractTier::Isolated);
+            calculation.update_with_perp_isolated_liability(
+                market.contract_tier == ContractTier::Isolated,
+            );
         }
 
         if has_perp_liability || calculation.context.margin_type != MarginRequirementType::Initial {
@@ -505,6 +506,45 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
     Ok(calculation)
 }
 
+pub fn validate_any_isolated_tier_requirements(
+    user: &User,
+    calculation: MarginCalculation,
+) -> DriftResult {
+    if calculation.with_perp_isolated_liability && !user.is_reduce_only() {
+        validate!(
+            calculation.num_perp_liabilities <= 1,
+            ErrorCode::IsolatedAssetTierViolation,
+            "User attempting to increase perp liabilities above 1 with a isolated tier liability"
+        )?;
+
+        validate!(
+            !user.is_margin_trading_enabled,
+            ErrorCode::IsolatedAssetTierViolation,
+            "User attempting isolated tier liability with margin trading enabled"
+        )?;
+
+        if calculation.num_spot_liabilities > 0 {
+            let quote_spot_position = user.get_quote_spot_position();
+            validate!(
+                    (calculation.num_spot_liabilities == 1 && quote_spot_position.is_borrow()
+                    ),
+                    ErrorCode::IsolatedAssetTierViolation,
+                    "User attempting to increase spot liabilities beyond usdc with a isolated tier liability"
+                )?;
+        }
+    }
+
+    if calculation.with_spot_isolated_liability && !user.is_reduce_only() {
+        validate!(
+            calculation.num_perp_liabilities == 0 && calculation.num_spot_liabilities == 1,
+            ErrorCode::IsolatedAssetTierViolation,
+            "User attempting to increase perp liabilities above 0 with a isolated tier liability"
+        )?;
+    }
+
+    Ok(())
+}
+
 pub fn meets_withdraw_margin_requirement(
     user: &User,
     perp_market_map: &PerpMarketMap,
@@ -531,13 +571,7 @@ pub fn meets_withdraw_margin_requirement(
         )?;
     }
 
-    if calculation.get_num_of_liabilities()? > 1 {
-        validate!(
-            !calculation.with_isolated_liability,
-            ErrorCode::IsolatedAssetTierViolation,
-            "User attempting to increase number of liabilities above 1 with a isolated tier liability"
-        )?;
-    }
+    validate_any_isolated_tier_requirements(user, calculation)?;
 
     validate!(
         calculation.meets_margin_requirement(),
@@ -582,13 +616,7 @@ pub fn meets_place_order_margin_requirement(
         return Err(ErrorCode::InsufficientCollateral);
     }
 
-    if calculation.get_num_of_liabilities()? > 1 {
-        validate!(
-            !calculation.with_isolated_liability,
-            ErrorCode::IsolatedAssetTierViolation,
-            "User attempting to increase number of liabilities above 1 with a isolated tier liability"
-        )?;
-    }
+    validate_any_isolated_tier_requirements(user, calculation)?;
 
     Ok(())
 }

programs/drift/src/math/margin/tests.rs

@@ -1322,6 +1322,7 @@ mod calculate_margin_requirement_and_total_collateral_and_liability_info {
     use anchor_lang::Owner;
     use solana_program::pubkey::Pubkey;
 
+    use crate::controller::position::PositionDirection;
     use crate::math::constants::{
         AMM_RESERVE_PRECISION, LIQUIDATION_FEE_PRECISION, PEG_PRECISION, QUOTE_PRECISION,
         SPOT_BALANCE_PRECISION, SPOT_BALANCE_PRECISION_U64, SPOT_CUMULATIVE_INTEREST_PRECISION,
@@ -1333,7 +1334,7 @@ mod calculate_margin_requirement_and_total_collateral_and_liability_info {
     use crate::state::margin_calculation::{MarginCalculation, MarginContext};
     use crate::state::oracle::{HistoricalOracleData, OracleSource};
     use crate::state::oracle_map::OracleMap;
-    use crate::state::perp_market::{MarketStatus, PerpMarket, AMM};
+    use crate::state::perp_market::{ContractTier, MarketStatus, PerpMarket, AMM};
     use crate::state::perp_market_map::PerpMarketMap;
     use crate::state::spot_market::{SpotBalanceType, SpotMarket};
     use crate::state::spot_market_map::SpotMarketMap;
@@ -1450,6 +1451,165 @@ mod calculate_margin_requirement_and_total_collateral_and_liability_info {
 
         assert_eq!(calculation.margin_requirement, QUOTE_PRECISION / 100);
         assert_eq!(calculation.get_num_of_liabilities().unwrap(), 1);
+        assert_eq!(calculation.with_perp_isolated_liability, false);
+        assert_eq!(calculation.with_spot_isolated_liability, false);
+    }
+
+    #[test]
+    fn isolated_contract_tier_count_and_check() {
+        let slot = 0_u64;
+
+        let mut sol_oracle_price = get_pyth_price(100, 6);
+        let sol_oracle_price_key =
+            Pubkey::from_str("J83w4HKfqxwcq3BEMMkPFSppX3gqekLyLJBexebFVkix").unwrap();
+        let pyth_program = crate::ids::pyth_program::id();
+        create_account_info!(
+            sol_oracle_price,
+            &sol_oracle_price_key,
+            &pyth_program,
+            oracle_account_info
+        );
+        let mut oracle_map = OracleMap::load_one(&oracle_account_info, slot, None).unwrap();
+
+        let mut market = PerpMarket {
+            amm: AMM {
+                base_asset_reserve: 100 * AMM_RESERVE_PRECISION,
+                quote_asset_reserve: 100 * AMM_RESERVE_PRECISION,
+                bid_base_asset_reserve: 101 * AMM_RESERVE_PRECISION,
+                bid_quote_asset_reserve: 99 * AMM_RESERVE_PRECISION,
+                ask_base_asset_reserve: 99 * AMM_RESERVE_PRECISION,
+                ask_quote_asset_reserve: 101 * AMM_RESERVE_PRECISION,
+                sqrt_k: 100 * AMM_RESERVE_PRECISION,
+                peg_multiplier: 100 * PEG_PRECISION,
+                order_step_size: 10000000,
+                oracle: sol_oracle_price_key,
+                ..AMM::default()
+            },
+            margin_ratio_initial: 1000,
+            margin_ratio_maintenance: 500,
+            status: MarketStatus::Initialized,
+            contract_tier: ContractTier::Isolated,
+            ..PerpMarket::default()
+        };
+        create_anchor_account_info!(market, PerpMarket, market_account_info);
+        let perp_market_map = PerpMarketMap::load_one(&market_account_info, true).unwrap();
+
+        let mut usdc_spot_market = SpotMarket {
+            market_index: 0,
+            oracle_source: OracleSource::QuoteAsset,
+            cumulative_deposit_interest: SPOT_CUMULATIVE_INTEREST_PRECISION,
+            decimals: 6,
+            initial_asset_weight: SPOT_WEIGHT_PRECISION,
+            maintenance_asset_weight: SPOT_WEIGHT_PRECISION,
+            deposit_balance: 10000 * SPOT_BALANCE_PRECISION,
+            liquidator_fee: 0,
+            historical_oracle_data: HistoricalOracleData::default_quote_oracle(),
+            ..SpotMarket::default()
+        };
+        create_anchor_account_info!(usdc_spot_market, SpotMarket, usdc_spot_market_account_info);
+        let mut sol_spot_market = SpotMarket {
+            market_index: 1,
+            oracle_source: OracleSource::Pyth,
+            oracle: sol_oracle_price_key,
+            cumulative_deposit_interest: SPOT_CUMULATIVE_INTEREST_PRECISION,
+            cumulative_borrow_interest: SPOT_CUMULATIVE_INTEREST_PRECISION,
+            decimals: 9,
+            initial_asset_weight: 8 * SPOT_WEIGHT_PRECISION / 10,
+            maintenance_asset_weight: 9 * SPOT_WEIGHT_PRECISION / 10,
+            initial_liability_weight: 12 * SPOT_WEIGHT_PRECISION / 10,
+            maintenance_liability_weight: 11 * SPOT_WEIGHT_PRECISION / 10,
+            liquidator_fee: LIQUIDATION_FEE_PRECISION / 1000,
+            ..SpotMarket::default()
+        };
+        create_anchor_account_info!(sol_spot_market, SpotMarket, sol_spot_market_account_info);
+        let spot_market_account_infos = Vec::from([
+            &usdc_spot_market_account_info,
+            &sol_spot_market_account_info,
+        ]);
+        let spot_market_map =
+            SpotMarketMap::load_multiple(spot_market_account_infos, true).unwrap();
+
+        // just usdc and resting limit order
+        let mut spot_positions = [SpotPosition::default(); 8];
+        spot_positions[0] = SpotPosition {
+            market_index: 0,
+            balance_type: SpotBalanceType::Deposit,
+            scaled_balance: 100 * SPOT_BALANCE_PRECISION_U64,
+            ..SpotPosition::default()
+        };
+
+        let user = User {
+            orders: get_orders(Order {
+                order_type: OrderType::Limit,
+                base_asset_amount: 1,
+                base_asset_amount_filled: 0,
+                ..Order::default()
+            }),
+            perp_positions: get_positions(PerpPosition {
+                market_index: 0,
+                base_asset_amount: 0,
+                open_orders: 1,
+                ..PerpPosition::default()
+            }),
+            spot_positions,
+            ..User::default()
+        };
+
+        let calculation: MarginCalculation =
+            calculate_margin_requirement_and_total_collateral_and_liability_info(
+                &user,
+                &perp_market_map,
+                &spot_market_map,
+                &mut oracle_map,
+                MarginContext::standard(MarginRequirementType::Initial),
+            )
+            .unwrap();
+
+        assert_eq!(calculation.margin_requirement, QUOTE_PRECISION / 100);
+        assert_eq!(calculation.get_num_of_liabilities().unwrap(), 1);
+        assert_eq!(calculation.with_spot_isolated_liability, false);
+        assert_eq!(calculation.with_perp_isolated_liability, true);
+
+        // just usdc, long iso perp, resting limit order to close
+        let mut spot_positions = [SpotPosition::default(); 8];
+        spot_positions[0] = SpotPosition {
+            market_index: 0,
+            balance_type: SpotBalanceType::Deposit,
+            scaled_balance: 100 * SPOT_BALANCE_PRECISION_U64,
+            ..SpotPosition::default()
+        };
+
+        let user = User {
+            orders: get_orders(Order {
+                order_type: OrderType::Limit,
+                base_asset_amount: market.amm.order_step_size,
+                base_asset_amount_filled: 0,
+                direction: PositionDirection::Short,
+                ..Order::default()
+            }),
+            perp_positions: get_positions(PerpPosition {
+                market_index: 0,
+                base_asset_amount: market.amm.order_step_size as i64,
+                open_orders: 1,
+                ..PerpPosition::default()
+            }),
+            spot_positions,
+            ..User::default()
+        };
+
+        let calculation: MarginCalculation =
+            calculate_margin_requirement_and_total_collateral_and_liability_info(
+                &user,
+                &perp_market_map,
+                &spot_market_map,
+                &mut oracle_map,
+                MarginContext::standard(MarginRequirementType::Initial),
+            )
+            .unwrap();
+
+        assert_eq!(calculation.margin_requirement, 110000);
+        assert_eq!(calculation.get_num_of_liabilities().unwrap(), 1);
+        assert_eq!(calculation.with_perp_isolated_liability, true);
     }
 
     #[test]

programs/drift/src/state/margin_calculation.rs

@@ -125,7 +125,8 @@ pub struct MarginCalculation {
     pub num_spot_liabilities: u8,
     pub num_perp_liabilities: u8,
     pub all_oracles_valid: bool,
-    pub with_isolated_liability: bool,
+    pub with_perp_isolated_liability: bool,
+    pub with_spot_isolated_liability: bool,
     pub total_spot_asset_value: i128,
     pub total_spot_liability_value: u128,
     pub total_perp_liability_value: u128,
@@ -143,7 +144,8 @@ impl MarginCalculation {
             num_spot_liabilities: 0,
             num_perp_liabilities: 0,
             all_oracles_valid: true,
-            with_isolated_liability: false,
+            with_perp_isolated_liability: false,
+            with_spot_isolated_liability: false,
             total_spot_asset_value: 0,
             total_spot_liability_value: 0,
             total_perp_liability_value: 0,
@@ -206,8 +208,12 @@ impl MarginCalculation {
         self.all_oracles_valid &= valid;
     }
 
-    pub fn update_with_isolated_liability(&mut self, isolated: bool) {
-        self.with_isolated_liability &= isolated;
+    pub fn update_with_spot_isolated_liability(&mut self, isolated: bool) {
+        self.with_spot_isolated_liability |= isolated;
+    }
+
+    pub fn update_with_perp_isolated_liability(&mut self, isolated: bool) {
+        self.with_perp_isolated_liability |= isolated;
     }
 
     pub fn validate_num_spot_liabilities(&self) -> DriftResult {