From da6a129a69bb6475c2f6191abd8045d127f22114 Mon Sep 17 00:00:00 2001
From: Chris Heaney <chrisheaney30@gmail.com>
Date: Tue, 7 Jan 2025 12:47:10 -0500
Subject: [PATCH] program: account for fuel when there is full withdraw

---
 programs/drift/src/math/margin.rs             | 23 +++++++++++++++++++
 .../drift/src/state/margin_calculation.rs     |  4 ++++
 2 files changed, 27 insertions(+)

diff --git a/programs/drift/src/math/margin.rs b/programs/drift/src/math/margin.rs
index be883b6ed..e25565fbb 100644
--- a/programs/drift/src/math/margin.rs
+++ b/programs/drift/src/math/margin.rs
@@ -254,6 +254,29 @@ pub fn calculate_margin_requirement_and_total_collateral_and_liability_info(
         validation::position::validate_spot_position(spot_position)?;
 
         if spot_position.is_available() {
+            // might need to update fuel to account for spot market deltas
+            if calculation.context.has_fuel_spot_delta(spot_position.market_index) {
+                let spot_market = spot_market_map.get_ref(&spot_position.market_index)?;
+                let (oracle_price_data, oracle_validity) = oracle_map.get_price_data_and_validity(
+                    MarketType::Spot,
+                    spot_market.market_index,
+                    &spot_market.oracle_id(),
+                    spot_market.historical_oracle_data.last_oracle_price_twap,
+                    spot_market.get_max_confidence_interval_multiplier()?,
+                )?;
+
+                let strict_oracle_price = StrictOraclePrice::new(
+                    oracle_price_data.price,
+                    spot_market
+                        .historical_oracle_data
+                        .last_oracle_price_twap_5min,
+                    calculation.context.strict,
+                );
+                strict_oracle_price.validate()?;
+
+                calculation.update_fuel_spot_bonus(&spot_market, 0, &strict_oracle_price)?;
+            }
+
             continue;
         }
 
diff --git a/programs/drift/src/state/margin_calculation.rs b/programs/drift/src/state/margin_calculation.rs
index 103d87808..dd95ef08e 100644
--- a/programs/drift/src/state/margin_calculation.rs
+++ b/programs/drift/src/state/margin_calculation.rs
@@ -105,6 +105,10 @@ impl MarginContext {
         self
     }
 
+    pub fn has_fuel_spot_delta(&self, market_index: u16) -> bool {
+        self.fuel_spot_deltas.iter().any(|(index, amount)| *index == market_index && *amount != 0)
+    }
+
     pub fn fuel_numerator(mut self, user: &User, now: i64) -> Self {
         self.fuel_bonus_numerator = user.get_fuel_bonus_numerator(now).unwrap();
         self