pyra_margin/drift/

capacity.rs

use std::collections::HashMap;

use pyra_types::{SpotBalanceType, SpotMarket, SpotPosition};

use super::balance::{calculate_value_usdc_base_units, get_token_balance};
use super::weights::{calculate_asset_weight, calculate_liability_weight, get_strict_price};
use crate::common::usdc_base_units_to_cents;
use crate::error::{MathError, MathResult};

const MARGIN_PRECISION: i128 = 10_000;

/// Per-position info emitted during capacity calculation.
///
/// Used downstream by liquidating spend jobs to know position sizes and weights.
#[derive(Debug, Clone, PartialEq)]
pub struct PositionInfo {
    pub market_index: u16,
    /// Unsigned token balance in base units.
    pub balance: u64,
    pub position_type: SpotBalanceType,
    pub price_usdc_base_units: u64,
    /// Initial asset weight from the spot market (basis points).
    pub weight_bps: u32,
}

/// Result of spending capacity calculation.
#[derive(Debug, Clone)]
pub struct CapacityResult {
    /// Max spendable via liquidating spend (unweighted collateral - slippage - liabilities), in cents.
    /// Excludes unliquidatable assets from collateral.
    pub total_spendable_cents: u64,
    /// Available credit line (weighted collateral - weighted liabilities), in cents.
    pub available_credit_cents: u64,
    /// USDC balance (market index 0) in cents.
    pub usdc_balance_cents: u64,
    /// Total weighted collateral in USDC base units.
    pub weighted_collateral_usdc_base_units: u64,
    /// Total weighted liabilities in USDC base units.
    pub weighted_liabilities_usdc_base_units: u64,
    /// Per-position breakdown for downstream use.
    pub position_infos: Vec<PositionInfo>,
}

/// Calculate spending capacity from Drift spot positions.
///
/// This is the core calculation used for card transaction authorization.
/// It computes:
/// - **total_spendable**: max amount for liquidating spends (collateral minus slippage minus liabilities,
///   excluding unliquidatable assets)
/// - **available_credit**: credit line from weighted margin (weighted collateral minus weighted liabilities)
/// - **usdc_balance**: direct USDC holdings
///
/// Positions in `unliquidatable_market_indices` are excluded from unweighted collateral
/// (affecting `total_spendable`) but still included in weighted calculations
/// (affecting `available_credit`).
///
/// Positions whose market index is missing from `spot_market_map` or `price_map` are skipped.
pub fn calculate_capacity(
    spot_positions: &[SpotPosition],
    spot_market_map: &HashMap<u16, SpotMarket>,
    price_map: &HashMap<u16, u64>,
    unliquidatable_market_indices: &[u16],
    max_slippage_bps: u64,
) -> MathResult<CapacityResult> {
    let mut total_collateral_usdc_base_units: u64 = 0;
    let mut total_liabilities_usdc_base_units: u64 = 0;

    let mut total_weighted_collateral_usdc_base_units: u64 = 0;
    let mut total_weighted_liabilities_usdc_base_units: u64 = 0;

    let mut usdc_balance_base_units: u64 = 0;

    let mut position_infos: Vec<PositionInfo> = Vec::new();

    for position in spot_positions {
        let market_index = position.market_index;

        let Some(spot_market) = spot_market_map.get(&market_index) else {
            continue;
        };
        let Some(price_usdc_base_units) = price_map.get(&market_index).copied() else {
            continue;
        };

        // Step 1: Calculate token balance and USDC value
        let token_balance_base_units = get_token_balance(position, spot_market)?;

        let is_asset = token_balance_base_units >= 0;
        let twap5min = spot_market
            .historical_oracle_data
            .last_oracle_price_twap5min;
        let strict_price = get_strict_price(price_usdc_base_units, twap5min, is_asset);

        let value_usdc_base_units = calculate_value_usdc_base_units(
            token_balance_base_units,
            strict_price,
            spot_market.decimals,
        )?;

        // Accumulate unweighted totals (excluding unliquidatable collateral)
        let is_unliquidatable_collateral =
            unliquidatable_market_indices.contains(&market_index) && value_usdc_base_units > 0;
        if !is_unliquidatable_collateral {
            update_running_totals(
                &mut total_collateral_usdc_base_units,
                &mut total_liabilities_usdc_base_units,
                value_usdc_base_units,
            )?;
        }

        // Step 2: Apply IMF-adjusted weights
        let token_amount_unsigned = token_balance_base_units.unsigned_abs();
        let weight_bps = if is_asset {
            calculate_asset_weight(token_amount_unsigned, price_usdc_base_units, spot_market)?
                as i128
        } else {
            calculate_liability_weight(token_amount_unsigned, spot_market)? as i128
        };
        let weighted_value_usdc_base_units = value_usdc_base_units
            .checked_mul(weight_bps)
            .ok_or(MathError::Overflow)?
            .checked_div(MARGIN_PRECISION)
            .ok_or(MathError::Overflow)?;

        update_running_totals(
            &mut total_weighted_collateral_usdc_base_units,
            &mut total_weighted_liabilities_usdc_base_units,
            weighted_value_usdc_base_units,
        )?;

        // Step 3: Track USDC balance (market index 0)
        if market_index == 0 && usdc_balance_base_units == 0 && token_balance_base_units > 0 {
            usdc_balance_base_units =
                u64::try_from(token_balance_base_units).map_err(|_| MathError::Overflow)?;
        }

        // Step 4: Store position info
        let token_balance_unsigned = u64::try_from(token_balance_base_units.unsigned_abs())
            .map_err(|_| MathError::Overflow)?;
        position_infos.push(PositionInfo {
            market_index,
            balance: token_balance_unsigned,
            position_type: position.balance_type.clone(),
            price_usdc_base_units,
            weight_bps: spot_market.initial_asset_weight,
        });
    }

    // Step 5: Available credit = weighted collateral - weighted liabilities
    let available_credit_base_units = total_weighted_collateral_usdc_base_units
        .saturating_sub(total_weighted_liabilities_usdc_base_units);
    let available_credit_cents = usdc_base_units_to_cents(available_credit_base_units)?;

    // Step 6: Total spendable = collateral - slippage - liabilities (for liquidating spends)
    let max_slippage_usdc_base_units = total_collateral_usdc_base_units
        .checked_mul(max_slippage_bps)
        .ok_or(MathError::Overflow)?
        .checked_div(10_000)
        .ok_or(MathError::Overflow)?;
    let total_spendable_base_units = total_collateral_usdc_base_units
        .saturating_sub(max_slippage_usdc_base_units)
        .saturating_sub(total_liabilities_usdc_base_units);
    let total_spendable_cents = usdc_base_units_to_cents(total_spendable_base_units)?;

    let usdc_balance_cents = usdc_base_units_to_cents(usdc_balance_base_units)?;

    Ok(CapacityResult {
        total_spendable_cents,
        available_credit_cents,
        usdc_balance_cents,
        weighted_collateral_usdc_base_units: total_weighted_collateral_usdc_base_units,
        weighted_liabilities_usdc_base_units: total_weighted_liabilities_usdc_base_units,
        position_infos,
    })
}

/// Accumulate a signed value into positive/negative running totals.
fn update_running_totals(
    total_positive: &mut u64,
    total_negative: &mut u64,
    value: i128,
) -> MathResult<()> {
    let value_unsigned = u64::try_from(value.unsigned_abs()).map_err(|_| MathError::Overflow)?;

    if value >= 0 {
        *total_positive = total_positive
            .checked_add(value_unsigned)
            .ok_or(MathError::Overflow)?;
    } else {
        *total_negative = total_negative
            .checked_add(value_unsigned)
            .ok_or(MathError::Overflow)?;
    }

    Ok(())
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::arithmetic_side_effects
)]
mod tests {
    use super::*;
    use pyra_types::{HistoricalOracleData, InsuranceFund};

    fn make_spot_market_with_twap(
        market_index: u16,
        decimals: u32,
        initial_asset_weight: u32,
        initial_liability_weight: u32,
        twap5min: i64,
    ) -> SpotMarket {
        let precision_decrease = 10u128.pow(19u32.saturating_sub(decimals));
        SpotMarket {
            pubkey: vec![],
            market_index,
            initial_asset_weight,
            initial_liability_weight,
            imf_factor: 0,
            scale_initial_asset_weight_start: 0,
            decimals,
            cumulative_deposit_interest: precision_decrease,
            cumulative_borrow_interest: precision_decrease,
            deposit_balance: 0,
            borrow_balance: 0,
            optimal_utilization: 0,
            optimal_borrow_rate: 0,
            max_borrow_rate: 0,
            min_borrow_rate: 0,
            insurance_fund: InsuranceFund::default(),
            historical_oracle_data: HistoricalOracleData {
                last_oracle_price_twap5min: twap5min,
            },
            oracle: None,
        }
    }

    /// Convenience: creates a spot market where TWAP matches the oracle price.
    fn make_spot_market(
        market_index: u16,
        decimals: u32,
        initial_asset_weight: u32,
        initial_liability_weight: u32,
        oracle_price: u64,
    ) -> SpotMarket {
        make_spot_market_with_twap(
            market_index,
            decimals,
            initial_asset_weight,
            initial_liability_weight,
            oracle_price as i64,
        )
    }

    fn make_position(market_index: u16, scaled_balance: u64, is_deposit: bool) -> SpotPosition {
        SpotPosition {
            market_index,
            scaled_balance,
            balance_type: if is_deposit {
                SpotBalanceType::Deposit
            } else {
                SpotBalanceType::Borrow
            },
            ..Default::default()
        }
    }

    // --- calculate_capacity ---

    #[test]
    fn empty_positions() {
        let result = calculate_capacity(&[], &HashMap::new(), &HashMap::new(), &[], 0).unwrap();
        assert_eq!(result.total_spendable_cents, 0);
        assert_eq!(result.available_credit_cents, 0);
        assert_eq!(result.usdc_balance_cents, 0);
        assert_eq!(result.weighted_collateral_usdc_base_units, 0);
        assert_eq!(result.weighted_liabilities_usdc_base_units, 0);
        assert!(result.position_infos.is_empty());
    }

    #[test]
    fn single_usdc_deposit() {
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let positions = vec![make_position(0, 100_000_000, true)]; // 100 USDC

        let mut markets = HashMap::new();
        markets.insert(0, usdc);
        let mut prices = HashMap::new();
        prices.insert(0, 1_000_000u64);

        let result = calculate_capacity(&positions, &markets, &prices, &[], 0).unwrap();

        assert_eq!(result.usdc_balance_cents, 10_000); // $100
        assert_eq!(result.total_spendable_cents, 10_000);
        assert_eq!(result.available_credit_cents, 10_000);
        assert_eq!(result.weighted_collateral_usdc_base_units, 100_000_000);
        assert_eq!(result.weighted_liabilities_usdc_base_units, 0);
        assert_eq!(result.position_infos.len(), 1);
        assert_eq!(result.position_infos[0].market_index, 0);
    }

    #[test]
    fn deposit_and_borrow() {
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let positions = vec![
            make_position(0, 100_000_000, true), // 100 USDC deposit
            make_position(0, 50_000_000, false), // 50 USDC borrow
        ];

        let mut markets = HashMap::new();
        markets.insert(0, usdc);
        let mut prices = HashMap::new();
        prices.insert(0, 1_000_000u64);

        let result = calculate_capacity(&positions, &markets, &prices, &[], 0).unwrap();

        assert_eq!(result.usdc_balance_cents, 10_000); // 100 USDC deposit
        assert_eq!(result.total_spendable_cents, 5_000); // 100 - 50 = 50 USDC
        assert_eq!(result.available_credit_cents, 5_000);
    }

    #[test]
    fn unliquidatable_excluded_from_spendable() {
        // Market 4 is unliquidatable
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let weth = make_spot_market(4, 9, 8_000, 12_000, 100_000_000);

        let positions = vec![
            make_position(0, 10_000_000, true),    // 10 USDC
            make_position(4, 1_000_000_000, true), // 1 wETH at $100
        ];

        let mut markets = HashMap::new();
        markets.insert(0, usdc);
        markets.insert(4, weth);
        let mut prices = HashMap::new();
        prices.insert(0, 1_000_000u64);
        prices.insert(4, 100_000_000u64); // $100

        let unliquidatable = vec![4u16];
        let result = calculate_capacity(&positions, &markets, &prices, &unliquidatable, 0).unwrap();

        // total_spendable: only USDC (10M base = 1000 cents), wETH excluded
        assert_eq!(result.total_spendable_cents, 1_000);
        // available_credit: includes wETH weighted (100M * 80% = 80M) + USDC (10M * 100% = 10M) = 90M = 9000 cents
        assert_eq!(result.available_credit_cents, 9_000);
        assert_eq!(result.position_infos.len(), 2);
    }

    #[test]
    fn slippage_reduces_spendable() {
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let positions = vec![make_position(0, 100_000_000, true)]; // 100 USDC

        let mut markets = HashMap::new();
        markets.insert(0, usdc);
        let mut prices = HashMap::new();
        prices.insert(0, 1_000_000u64);

        // 10% slippage = 1000 bps
        let result = calculate_capacity(&positions, &markets, &prices, &[], 1_000).unwrap();

        // 100 USDC collateral, 10% slippage = 10 USDC, spendable = 90 USDC = 9000 cents
        assert_eq!(result.total_spendable_cents, 9_000);
        // available_credit not affected by slippage
        assert_eq!(result.available_credit_cents, 10_000);
    }

    #[test]
    fn missing_market_skipped() {
        let positions = vec![make_position(5, 1_000_000, true)];

        let result =
            calculate_capacity(&positions, &HashMap::new(), &HashMap::new(), &[], 0).unwrap();

        assert_eq!(result.total_spendable_cents, 0);
        assert!(result.position_infos.is_empty());
    }

    #[test]
    fn missing_price_skipped() {
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let positions = vec![make_position(0, 1_000_000, true)];

        let mut markets = HashMap::new();
        markets.insert(0, usdc);

        let result = calculate_capacity(&positions, &markets, &HashMap::new(), &[], 0).unwrap();

        assert_eq!(result.total_spendable_cents, 0);
        assert!(result.position_infos.is_empty());
    }

    #[test]
    fn multi_position_with_unliquidatable_and_slippage() {
        let usdc = make_spot_market(0, 6, 10_000, 10_000, 1_000_000);
        let m4 = make_spot_market(4, 9, 8_000, 12_000, 200_000_000); // unliquidatable
        let m5 = make_spot_market(5, 9, 8_000, 12_000, 100_000_000);
        let m6 = make_spot_market(6, 6, 10_000, 10_000, 1_000_000);

        let positions = vec![
            make_position(0, 50_000_000, true),    // 50 USDC
            make_position(4, 1_000_000_000, true), // 1 token @ $200 (unliquidatable)
            make_position(5, 500_000_000, true),   // 0.5 token @ $100
            make_position(6, 20_000_000, false),   // 20 USDC-like borrow
        ];

        let mut markets = HashMap::new();
        markets.insert(0, usdc);
        markets.insert(4, m4);
        markets.insert(5, m5);
        markets.insert(6, m6);
        let mut prices = HashMap::new();
        prices.insert(0, 1_000_000u64);
        prices.insert(4, 200_000_000u64);
        prices.insert(5, 100_000_000u64);
        prices.insert(6, 1_000_000u64);

        let unliquidatable = vec![4u16, 32u16];
        let result =
            calculate_capacity(&positions, &markets, &prices, &unliquidatable, 500).unwrap();

        // Unweighted collateral (excluding market 4): 50M (USDC) + 50M (m5: 0.5 * $100) = 100M
        // Unweighted liabilities: 20M (m6 borrow)
        // Slippage: 100M * 500/10000 = 5M
        // total_spendable = 100M - 5M - 20M = 75M base = 7500 cents
        assert_eq!(result.total_spendable_cents, 7_500);

        // Weighted collateral: 50M*100% + 200M*80% + 50M*80% = 50M + 160M + 40M = 250M
        // Weighted liabilities: 20M*100% = 20M
        // available_credit = 250M - 20M = 230M base = 23000 cents
        assert_eq!(result.available_credit_cents, 23_000);
        assert_eq!(result.usdc_balance_cents, 5_000);
        assert_eq!(result.position_infos.len(), 4);
    }

    // --- update_running_totals ---

    #[test]
    fn running_totals_positive() {
        let mut pos = 0u64;
        let mut neg = 0u64;
        update_running_totals(&mut pos, &mut neg, 100).unwrap();
        assert_eq!(pos, 100);
        assert_eq!(neg, 0);
    }

    #[test]
    fn running_totals_negative() {
        let mut pos = 0u64;
        let mut neg = 0u64;
        update_running_totals(&mut pos, &mut neg, -50).unwrap();
        assert_eq!(pos, 0);
        assert_eq!(neg, 50);
    }

    #[test]
    fn running_totals_accumulate() {
        let mut pos = 10u64;
        let mut neg = 5u64;
        update_running_totals(&mut pos, &mut neg, 20).unwrap();
        update_running_totals(&mut pos, &mut neg, -15).unwrap();
        assert_eq!(pos, 30);
        assert_eq!(neg, 20);
    }
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::arithmetic_side_effects
)]
mod proptests {
    use super::*;
    use proptest::prelude::*;

    proptest! {
        #[test]
        fn spendable_le_collateral_minus_liabilities(
            collateral_base in 0u64..=1_000_000_000_000u64,
            liabilities_base in 0u64..=500_000_000_000u64,
        ) {
            // Spendable should never exceed collateral - liabilities (without slippage)
            let collateral_cents = usdc_base_units_to_cents(collateral_base).unwrap();
            let liabilities_cents = usdc_base_units_to_cents(liabilities_base).unwrap();
            let max_possible = collateral_cents.saturating_sub(liabilities_cents);

            // Since we're testing the formula directly, we can verify the invariant
            let spendable_base = collateral_base.saturating_sub(liabilities_base);
            let spendable_cents = usdc_base_units_to_cents(spendable_base).unwrap();
            prop_assert!(spendable_cents <= max_possible + 1, "rounding violation");
        }
    }
}