program: add max margin ratio per position

programs/drift/src/controller/position.rs

@@ -49,8 +49,19 @@ pub fn add_new_position(
         .position(|market_position| market_position.is_available())
         .ok_or(ErrorCode::MaxNumberOfPositions)?;
 
+    let max_margin_ratio = {
+        let old_position = &user_positions[new_position_index];
+
+        if old_position.market_index == market_index {
+            old_position.max_margin_ratio
+        } else {
+            0_u16
+        }
+    };
+
     let new_market_position = PerpPosition {
         market_index,
+        max_margin_ratio,
         ..PerpPosition::default()
     };
 

programs/drift/src/instructions/keeper.rs

@@ -1,4 +1,5 @@
 use std::cell::RefMut;
+use std::collections::BTreeMap;
 use std::convert::TryFrom;
 
 use anchor_lang::prelude::*;
@@ -26,6 +27,7 @@ use crate::instructions::constraints::*;
 use crate::instructions::optional_accounts::{load_maps, AccountMaps};
 use crate::math::casting::Cast;
 use crate::math::constants::QUOTE_SPOT_MARKET_INDEX;
+use crate::math::margin::get_margin_calculation_for_disable_high_leverage_mode;
 use crate::math::margin::{calculate_user_equity, meets_settle_pnl_maintenance_margin_requirement};
 use crate::math::orders::{estimate_price_from_side, find_bids_and_asks_from_users};
 use crate::math::position::calculate_base_asset_value_and_pnl_with_oracle_price;
@@ -2803,20 +2805,13 @@ pub fn handle_disable_user_high_leverage_mode<'c: 'info, 'info>(
         }
     }
 
-    let custom_margin_ratio_before = user.max_margin_ratio;
-    user.max_margin_ratio = 0;
-
-    let margin_calc = calculate_margin_requirement_and_total_collateral_and_liability_info(
-        &user,
+    let margin_calc = get_margin_calculation_for_disable_high_leverage_mode(
+        &mut user,
         &perp_market_map,
         &spot_market_map,
         &mut oracle_map,
-        MarginContext::standard(MarginRequirementType::Initial)
-            .margin_buffer(MARGIN_PRECISION / 100), // 1% buffer
     )?;
 
-    user.max_margin_ratio = custom_margin_ratio_before;
-
     if margin_calc.num_perp_liabilities > 0 {
         let mut requires_invariant_check = false;
 

programs/drift/src/instructions/user.rs

@@ -2945,6 +2945,20 @@ pub fn handle_update_user_custom_margin_ratio(
     Ok(())
 }
 
+pub fn handle_update_user_perp_position_custom_margin_ratio(
+    ctx: Context<UpdateUserPerpPositionCustomMarginRatio>,
+    _sub_account_id: u16,
+    perp_market_index: u16,
+    margin_ratio: u16,
+) -> Result<()> {
+    let mut user = load_mut!(ctx.accounts.user)?;
+
+    let perp_position = user.force_get_perp_position_mut(perp_market_index)?;
+    perp_position.max_margin_ratio = margin_ratio;
+    Ok(())
+}
+
+
 pub fn handle_update_user_margin_trading_enabled<'c: 'info, 'info>(
     ctx: Context<'_, '_, 'c, 'info, UpdateUser<'info>>,
     _sub_account_id: u16,
@@ -4432,6 +4446,21 @@ pub struct UpdateUser<'info> {
     pub authority: Signer<'info>,
 }
 
+#[derive(Accounts)]
+#[instruction(
+    sub_account_id: u16,
+)]
+pub struct UpdateUserPerpPositionCustomMarginRatio<'info> {
+    #[account(
+        mut,
+        seeds = [b"user", authority.key.as_ref(), sub_account_id.to_le_bytes().as_ref()],
+        bump,
+        constraint = can_sign_for_user(&user, &authority)?
+    )]
+    pub user: AccountLoader<'info, User>,
+    pub authority: Signer<'info>,
+}
+
 #[derive(Accounts)]
 pub struct DeleteUser<'info> {
     #[account(

programs/drift/src/lib.rs

@@ -367,6 +367,15 @@ pub mod drift {
         handle_update_user_custom_margin_ratio(ctx, _sub_account_id, margin_ratio)
     }
 
+    pub fn update_user_perp_position_custom_margin_ratio(
+        ctx: Context<UpdateUserPerpPositionCustomMarginRatio>,
+        _sub_account_id: u16,
+        perp_market_index: u16,
+        margin_ratio: u16,
+    ) -> Result<()> {
+        handle_update_user_perp_position_custom_margin_ratio(ctx, _sub_account_id, perp_market_index, margin_ratio)
+    }
+
     pub fn update_user_margin_trading_enabled<'c: 'info, 'info>(
         ctx: Context<'_, '_, 'c, 'info, UpdateUser<'info>>,
         _sub_account_id: u16,

programs/drift/src/math/margin.rs

@@ -6,6 +6,7 @@ use crate::math::constants::{
 };
 use crate::math::position::calculate_base_asset_value_and_pnl_with_oracle_price;
 
+use crate::MARGIN_PRECISION;
 use crate::{validate, PRICE_PRECISION_I128};
 use crate::{validation, PRICE_PRECISION_I64};
 
@@ -27,6 +28,7 @@ use crate::state::spot_market_map::SpotMarketMap;
 use crate::state::user::{MarketType, OrderFillSimulation, PerpPosition, User};
 use num_integer::Roots;
 use std::cmp::{max, min, Ordering};
+use std::collections::BTreeMap;
 
 #[cfg(test)]
 mod tests;
@@ -535,6 +537,12 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
             0,
         )?;
 
+        let perp_position_custom_margin_ratio = if context.margin_type == MarginRequirementType::Initial {
+            market_position.max_margin_ratio as u32
+        } else {
+            0_u32
+        };
+
         let (
             perp_margin_requirement,
             weighted_pnl,
@@ -547,7 +555,7 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
             oracle_price_data,
             &strict_quote_price,
             context.margin_type,
-            user_custom_margin_ratio,
+            user_custom_margin_ratio.max(perp_position_custom_margin_ratio),
             user_high_leverage_mode,
             calculation.track_open_orders_fraction(),
         )?;
@@ -884,6 +892,42 @@ pub fn validate_spot_margin_trading(
     Ok(())
 }
 
+pub fn get_margin_calculation_for_disable_high_leverage_mode(
+    user: &mut User,
+    perp_market_map: &PerpMarketMap,
+    spot_market_map: &SpotMarketMap,
+    oracle_map: &mut OracleMap,
+) -> DriftResult<MarginCalculation> {
+    let custom_margin_ratio_before = user.max_margin_ratio;
+
+
+    let mut perp_position_max_margin_ratio_map = BTreeMap::new();
+    for (index, position) in user.perp_positions.iter_mut().enumerate() {
+        if position.max_margin_ratio == 0 {
+            continue;
+        }
+
+        perp_position_max_margin_ratio_map.insert(index, position.max_margin_ratio);
+        position.max_margin_ratio = 0;
+    }
+
+    let margin_buffer = MARGIN_PRECISION / 100; // 1% buffer
+    let margin_calc = calculate_margin_requirement_and_total_collateral_and_liability_info(
+        user,
+        perp_market_map,
+        spot_market_map,
+        oracle_map,
+        MarginContext::standard(MarginRequirementType::Initial).margin_buffer(margin_buffer),
+    )?;
+
+    user.max_margin_ratio = custom_margin_ratio_before;
+    for (index, position) in perp_position_max_margin_ratio_map.iter() {
+        user.perp_positions[*index].max_margin_ratio = *position;
+    }
+
+    Ok(margin_calc)
+}
+
 pub fn calculate_user_equity(
     user: &User,
     perp_market_map: &PerpMarketMap,