program: add drift-rs feature flag to add more data to margin calc (#978

) * full margin calcs * update match * adjust directive placement * change directive name * update * cargo fmt * changelog * changelog
drift-labs · Mar 27, 2024 · fab9760 · fab9760
1 parent e93d287
commit fab9760
Show file tree

Hide file tree

Showing 5 changed files with 79 additions and 37 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,12 +10,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Features
 - program: add invariant check boolean for attempt settle revenue to insurance ([#937](https://github.com/drift-labs/protocol-v2/pull/937)) 
 - program: improve best bid/ask estimate in mark twap update  ([#975](https://github.com/drift-labs/protocol-v2/pull/975)) 
+- program: add optional margin calculations for drift-rs ([#978](https://github.com/drift-labs/protocol-v2/pull/978))
 
 ### Fixes
 
 ### Breaking
 
-## [2.74.0] - 2023-03-25
+## [2.74.0] - 2024-03-25
 
 ### Features
 
@@ -28,7 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.73.0] - 2023-03-15
+## [2.73.0] - 2024-03-15
 
 ### Features
 
@@ -38,7 +39,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.72.0] - 2023-03-14
+## [2.72.0] - 2024-03-14
 
 ### Features
 
@@ -55,7 +56,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.71.0] - 2023-03-11
+## [2.71.0] - 2024-03-11
 
 ### Features
 
@@ -71,7 +72,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - sdk: account for max confidence in isOracleValid ([#949](https://github.com/drift-labs/protocol-v2/pull/949))
 
-## [2.70.0] - 2023-03-07
+## [2.70.0] - 2024-03-07
 
 ### Features
 
@@ -81,7 +82,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.69.0] - 2023-03-06
+## [2.69.0] - 2024-03-06
 
 ### Features
 
@@ -94,7 +95,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - program: revert switchboard ([#935](https://github.com/drift-labs/protocol-v2/pull/935))
 
-## [2.68.0] - 2023-03-05
+## [2.68.0] - 2024-03-05
 
 ### Features
 
@@ -106,7 +107,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.67.0] - 2023-03-05
+## [2.67.0] - 2024-03-05
 
 ### Features
 
@@ -117,7 +118,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.66.0] - 2023-02-28
+## [2.66.0] - 2024-02-28
 
 ### Features
 
@@ -127,7 +128,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.65.0] - 2023-02-26
+## [2.65.0] - 2024-02-26
 
 ### Features
 
@@ -141,7 +142,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.64.0] - 2023-02-20
+## [2.64.0] - 2024-02-20
 
 ### Features
 
@@ -151,7 +152,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.63.0] - 2023-02-16
+## [2.63.0] - 2024-02-16
 
 ### Features
 
@@ -167,7 +168,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.62.0] - 2023-02-14
+## [2.62.0] - 2024-02-14
 
 ### Features
 
@@ -177,7 +178,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.61.0] - 2023-02-09
+## [2.61.0] - 2024-02-09
 
 ### Features
 
@@ -190,7 +191,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.60.0] - 2023-02-07
+## [2.60.0] - 2024-02-07
 
 ### Features
 
@@ -202,7 +203,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.59.0] - 2023-01-30
+## [2.59.0] - 2024-01-30
 
 ### Features
 
@@ -213,7 +214,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.58.0] - 2023-01-27
+## [2.58.0] - 2024-01-27
 
 ### Features
 
@@ -223,7 +224,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.57.0] - 2023-01-25
+## [2.57.0] - 2024-01-25
 
 ### Features
 
@@ -233,7 +234,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.56.0] - 2023-01-24
+## [2.56.0] - 2024-01-24
 
 ### Features
 
@@ -243,7 +244,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.55.0] - 2023-01-18
+## [2.55.0] - 2024-01-18
 
 ### Features
 
@@ -253,7 +254,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Breaking
 
-## [2.54.0] - 2023-01-15
+## [2.54.0] - 2024-01-15
 
 ### Features
 

diff --git a/programs/drift/Cargo.toml b/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"

diff --git a/programs/drift/src/controller/amm/tests.rs b/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;

diff --git a/programs/drift/src/math/margin.rs b/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_token_value)?;
                 }
                 Ordering::Less => {
                     validate!(
@@ -389,6 +398,9 @@ 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_token_value.unsigned_abs())?;
                 }
                 Ordering::Equal => {
                     if spot_position.has_open_order() {
@@ -403,13 +415,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 +502,11 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
 
         calculation.add_total_collateral(weighted_pnl)?;
 
+        #[cfg(drift_rs)]
+        calculation.add_perp_liability_value(worst_case_base_asset_value)?;
+        #[cfg(drift_rs)]
+        calculation.add_perp_pnl(weighted_pnl)?;
+
         let has_perp_liability = market_position.base_asset_amount != 0
             || market_position.quote_asset_amount < 0
             || market_position.has_open_order()

diff --git a/programs/drift/src/state/margin_calculation.rs b/programs/drift/src/state/margin_calculation.rs
@@ -130,6 +130,7 @@ pub struct MarginCalculation {
     pub total_spot_asset_value: i128,
     pub total_spot_liability_value: u128,
     pub total_perp_liability_value: u128,
+    pub total_perp_pnl: i128,
     pub open_orders_margin_requirement: u128,
     tracked_market_margin_requirement: u128,
 }
@@ -149,6 +150,7 @@ impl MarginCalculation {
             total_spot_asset_value: 0,
             total_spot_liability_value: 0,
             total_perp_liability_value: 0,
+            total_perp_pnl: 0,
             open_orders_margin_requirement: 0,
             tracked_market_margin_requirement: 0,
         }
@@ -204,6 +206,34 @@ 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(())
+    }
+
+    #[cfg(drift_rs)]
+    pub fn add_perp_pnl(&mut self, perp_upnl: i128) -> DriftResult {
+        self.total_perp_upnl = self.total_perp_upnl.safe_add(perp_upnl)?;
+        Ok(())
+    }
+
     pub fn update_all_oracles_valid(&mut self, valid: bool) {
         self.all_oracles_valid &= valid;
     }