Update fuel.py

src/driftpy/math/fuel.py

@@ -1,56 +1,100 @@
 from driftpy.types import SpotMarketAccount, PerpMarketAccount
 from driftpy.constants.numeric_constants import QUOTE_PRECISION, FUEL_WINDOW
 
+def _calculate_fuel_bonus(
+    value: int, 
+    boost_factor: int, 
+    bonus_numerator: int
+) -> int:
+    """Calculate fuel bonus using unified formula.
+
+    Args:
+        value: Absolute value of the position/stake
+        boost_factor: Market-specific boost multiplier
+        bonus_numerator: Protocol-level bonus parameter
+
+    Returns:
+        Daily fuel bonus scaled by QUOTE_PRECISION
+    """
+    if abs(value) <= QUOTE_PRECISION:  # Dust threshold
+        return 0
+
+    return (
+        abs(value) 
+        * bonus_numerator 
+        * boost_factor 
+    ) // (FUEL_WINDOW * (QUOTE_PRECISION // 10))  # Combined divisions
+
 
 def calculate_insurance_fuel_bonus(
-    spot_market: SpotMarketAccount, token_stake_amount: int, fuel_bonus_numerator: int
+    spot_market: SpotMarketAccount, 
+    token_stake_amount: int, 
+    fuel_bonus_numerator: int
 ) -> int:
-    insurance_fund_fuel = (
-        abs(token_stake_amount) * fuel_bonus_numerator
-    ) * spot_market.fuel_boost_insurance
-    insurace_fund_fuel_per_day = insurance_fund_fuel // FUEL_WINDOW
-    insurance_fund_fuel_scaled = insurace_fund_fuel_per_day // (QUOTE_PRECISION // 10)
-
-    return insurance_fund_fuel_scaled
+    """Calculate insurance fund fuel bonus.
+
+    Args:
+        spot_market: Spot market account data
+        token_stake_amount: Staked token amount (absolute value)
+        fuel_bonus_numerator: Protocol-level bonus parameter
+
+    Returns:
+        Scaled daily insurance fuel bonus
+    """
+    return _calculate_fuel_bonus(
+        token_stake_amount,
+        spot_market.fuel_boost_insurance,
+        fuel_bonus_numerator
+    )
 
 
 def calculate_spot_fuel_bonus(
-    spot_market: SpotMarketAccount, signed_token_value: int, fuel_bonus_numerator: int
+    spot_market: SpotMarketAccount, 
+    signed_token_value: int, 
+    fuel_bonus_numerator: int
 ) -> int:
-    spot_fuel_scaled: int
-
-    # dust
-    if abs(signed_token_value) <= QUOTE_PRECISION:
-        spot_fuel_scaled = 0
-    elif signed_token_value > 0:
-        deposit_fuel = (
-            abs(signed_token_value) * fuel_bonus_numerator
-        ) * spot_market.fuel_boost_deposits
-        deposit_fuel_per_day = deposit_fuel // FUEL_WINDOW
-        spot_fuel_scaled = deposit_fuel_per_day // (QUOTE_PRECISION // 10)
-    else:
-        borrow_fuel = (
-            abs(signed_token_value) * fuel_bonus_numerator
-        ) * spot_market.fuel_boost_borrows
-        borrow_fuel_per_day = borrow_fuel // FUEL_WINDOW
-        spot_fuel_scaled = borrow_fuel_per_day // (QUOTE_PRECISION // 10)
-
-    return spot_fuel_scaled
+    """Calculate spot market fuel bonus (deposits/borrows).
+
+    Args:
+        spot_market: Spot market account data
+        signed_token_value: Signed position value 
+            (positive for deposits, negative for borrows)
+        fuel_bonus_numerator: Protocol-level bonus parameter
+
+    Returns:
+        Scaled daily fuel bonus based on position type
+    """
+    if signed_token_value == 0:
+        return 0
+
+    boost_factor = (
+        spot_market.fuel_boost_deposits if signed_token_value > 0 
+        else spot_market.fuel_boost_borrows
+    )
+    return _calculate_fuel_bonus(
+        signed_token_value,
+        boost_factor,
+        fuel_bonus_numerator
+    )
 
 
 def calculate_perp_fuel_bonus(
-    perp_market: PerpMarketAccount, base_asset_value: int, fuel_bonus_numerator: int
+    perp_market: PerpMarketAccount, 
+    base_asset_value: int, 
+    fuel_bonus_numerator: int
 ) -> int:
-    perp_fuel_scaled: int
-
-    # dust
-    if abs(base_asset_value) <= QUOTE_PRECISION:
-        perp_fuel_scaled = 0
-    else:
-        perp_fuel = (
-            abs(base_asset_value) * fuel_bonus_numerator
-        ) * perp_market.fuel_boost_position
-        perp_fuel_per_day = perp_fuel // FUEL_WINDOW
-        perp_fuel_scaled = perp_fuel_per_day // (QUOTE_PRECISION // 10)
-
-    return perp_fuel_scaled
+    """Calculate perp market position fuel bonus.
+
+    Args:
+        perp_market: Perp market account data
+        base_asset_value: Absolute position size
+        fuel_bonus_numerator: Protocol-level bonus parameter
+
+    Returns:
+        Scaled daily perp position fuel bonus
+    """
+    return _calculate_fuel_bonus(
+        base_asset_value,
+        perp_market.fuel_boost_position,
+        fuel_bonus_numerator
+    )