frank/full margin calcs

programs/drift/Cargo.toml

@@ -15,6 +15,7 @@ cpi = ["no-entrypoint"]
 mainnet-beta=[]
 anchor-test= []
 default=["mainnet-beta"]
+drift-rs=[]
 
 [dependencies]
 anchor-lang = "0.27.0"

programs/drift/src/controller/amm/tests.rs

@@ -825,26 +825,13 @@ fn update_pool_balances_revenue_to_fee_test() {
     assert_eq!(spot_market.insurance_fund.revenue_settle_period, 0);
 
     spot_market.insurance_fund.revenue_settle_period = 0;
-    let res = settle_revenue_to_insurance_fund(
-        0,
-        0,
-        &mut spot_market,
-        now + 3600,
-        true,
-    )
-    .unwrap();
+    let res = settle_revenue_to_insurance_fund(0, 0, &mut spot_market, now + 3600, true).unwrap();
     assert_eq!(res, 0);
     spot_market.insurance_fund.revenue_settle_period = 1;
 
     spot_market.revenue_pool.scaled_balance = 0;
-    let res = settle_revenue_to_insurance_fund(
-        200000000,
-        0,
-        &mut spot_market,
-        now + 1,
-        false,
-    )
-    .unwrap();
+    let res =
+        settle_revenue_to_insurance_fund(200000000, 0, &mut spot_market, now + 1, false).unwrap();
     assert_eq!(res, 0);
     spot_market.revenue_pool.scaled_balance = 100 * SPOT_BALANCE_PRECISION;
     now += 2;

programs/drift/src/math/margin.rs

@@ -299,6 +299,9 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
             match spot_position.balance_type {
                 SpotBalanceType::Deposit => {
                     calculation.add_total_collateral(token_value)?;
+
+                    #[cfg(drift_rs)]
+                    calculation.add_spot_asset_value(token_value)?;
                 }
                 SpotBalanceType::Borrow => {
                     let token_value = token_value.unsigned_abs();
@@ -318,6 +321,9 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
                     )?;
 
                     calculation.add_spot_liability()?;
+
+                    #[cfg(drift_rs)]
+                    calculation.add_spot_liability_value(token_value)?;
                 }
             }
         } else {
@@ -362,6 +368,9 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
                 Ordering::Greater => {
                     calculation
                         .add_total_collateral(worst_case_weighted_token_value.cast::<i128>()?)?;
+
+                    #[cfg(drift_rs)]
+                    calculation.add_spot_asset_value(worst_case_weighted_token_value)?;
                 }
                 Ordering::Less => {
                     validate!(
@@ -389,6 +398,10 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
                     calculation.update_with_spot_isolated_liability(
                         spot_market.asset_tier == AssetTier::Isolated,
                     );
+
+                    #[cfg(drift_rs)]
+                    calculation
+                        .add_spot_liability_value(worst_case_weighted_token_value.unsigned_abs())?;
                 }
                 Ordering::Equal => {
                     if spot_position.has_open_order() {
@@ -403,13 +416,19 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
             match worst_case_orders_value.cmp(&0) {
                 Ordering::Greater => {
                     calculation.add_total_collateral(worst_case_orders_value.cast::<i128>()?)?;
+
+                    #[cfg(drift_rs)]
+                    calculation.add_spot_asset_value(worst_case_orders_value)?;
                 }
                 Ordering::Less => {
                     calculation.add_margin_requirement(
                         worst_case_orders_value.unsigned_abs(),
                         worst_case_orders_value.unsigned_abs(),
                         MarketIdentifier::spot(0),
                     )?;
+
+                    #[cfg(drift_rs)]
+                    calculation.add_spot_liability_value(worst_case_orders_value.unsigned_abs())?;
                 }
                 Ordering::Equal => {}
             }
@@ -484,6 +503,14 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
 
         calculation.add_total_collateral(weighted_pnl)?;
 
+        #[cfg(drift_rs)]
+        match weighted_pnl.cmp(&0) {
+            Ordering::Less => {
+                calculation.add_perp_liability_value(weighted_pnl.unsigned_abs())?;
+            }
+            _ => {}
+        }
+
         let has_perp_liability = market_position.base_asset_amount != 0
             || market_position.quote_asset_amount < 0
             || market_position.has_open_order()

programs/drift/src/state/margin_calculation.rs

@@ -204,6 +204,28 @@ impl MarginCalculation {
         Ok(())
     }
 
+    #[cfg(drift_rs)]
+    pub fn add_spot_asset_value(&mut self, spot_asset_value: i128) -> DriftResult {
+        self.total_spot_asset_value = self.total_spot_asset_value.safe_add(spot_asset_value)?;
+        Ok(())
+    }
+
+    #[cfg(drift_rs)]
+    pub fn add_spot_liability_value(&mut self, spot_liability_value: u128) -> DriftResult {
+        self.total_spot_liability_value = self
+            .total_spot_liability_value
+            .safe_add(spot_liability_value)?;
+        Ok(())
+    }
+
+    #[cfg(drift_rs)]
+    pub fn add_perp_liability_value(&mut self, perp_liability_value: u128) -> DriftResult {
+        self.total_perp_liability_value = self
+            .total_perp_liability_value
+            .safe_add(perp_liability_value)?;
+        Ok(())
+    }
+
     pub fn update_all_oracles_valid(&mut self, valid: bool) {
         self.all_oracles_valid &= valid;
     }