pyra_margin/

limits.rs

use std::cmp;

use pyra_types::SpotMarket;

use crate::balance::calculate_value_usdc_base_units;
use crate::error::{MathError, MathResult};
use crate::math::CheckedDivCeil;
use crate::weights::{calculate_asset_weight, calculate_liability_weight, get_strict_price};

const MARGIN_PRECISION: i128 = 10_000;
const PRICE_PRECISION: i128 = 1_000_000;

/// Input data for a single position in margin calculations.
///
/// The caller is responsible for computing `token_balance` — this allows
/// different services to compose balances differently (e.g. Drift-only
/// vs Drift + deposit addresses - pending withdraws).
pub struct PositionData<'a> {
    /// Token balance in base units (positive = deposit, negative = borrow).
    /// Use `get_token_balance()` to compute from a `SpotPosition`, then add
    /// any off-chain adjustments.
    pub token_balance: i128,
    /// Current oracle price in USDC base units (1e6).
    pub price_usdc_base_units: u64,
    /// 5-minute TWAP from `SpotMarket::historical_oracle_data.last_oracle_price_twap5min`.
    pub twap5min: i64,
    /// The spot market configuration.
    pub spot_market: &'a SpotMarket,
}

/// Aggregated margin state across all positions.
///
/// Computed once from all positions, then used to derive per-market
/// withdraw/borrow limits and credit usage.
#[derive(Debug, Clone, Copy)]
pub struct MarginState {
    /// Total weighted collateral in USDC base units.
    pub total_weighted_collateral: i128,
    /// Total weighted liabilities in USDC base units (always non-negative).
    pub total_weighted_liabilities: i128,
}

impl MarginState {
    /// Calculate margin state from a set of positions.
    ///
    /// For each position: applies strict oracle pricing, computes USDC value,
    /// applies IMF-adjusted weights, and accumulates into collateral/liabilities.
    pub fn calculate(positions: &[PositionData<'_>]) -> MathResult<Self> {
        let mut total_weighted_collateral: i128 = 0;
        let mut total_weighted_liabilities: i128 = 0;

        for pos in positions {
            if pos.token_balance == 0 {
                continue;
            }

            let is_asset = pos.token_balance >= 0;
            let strict_price = get_strict_price(pos.price_usdc_base_units, pos.twap5min, is_asset);

            let value_usdc = calculate_value_usdc_base_units(
                pos.token_balance,
                strict_price,
                pos.spot_market.decimals,
            )?;

            let token_amount_unsigned = pos.token_balance.unsigned_abs();
            let weight_bps = if is_asset {
                calculate_asset_weight(
                    token_amount_unsigned,
                    pos.price_usdc_base_units,
                    pos.spot_market,
                )?
            } else {
                calculate_liability_weight(token_amount_unsigned, pos.spot_market)?
            };

            // weight_bps is u128 ≤ ~20_000, safe to cast to i128
            let weighted_value = value_usdc
                .checked_mul(weight_bps as i128)
                .ok_or(MathError::Overflow)?
                .checked_div(MARGIN_PRECISION)
                .ok_or(MathError::Overflow)?;

            if weighted_value >= 0 {
                total_weighted_collateral = total_weighted_collateral
                    .checked_add(weighted_value)
                    .ok_or(MathError::Overflow)?;
            } else {
                total_weighted_liabilities = total_weighted_liabilities
                    .checked_add(weighted_value.checked_neg().ok_or(MathError::Overflow)?)
                    .ok_or(MathError::Overflow)?;
            }
        }

        Ok(Self {
            total_weighted_collateral,
            total_weighted_liabilities,
        })
    }

    /// Free collateral = weighted collateral - weighted liabilities, clamped to 0.
    pub fn free_collateral(&self) -> u64 {
        let fc = self
            .total_weighted_collateral
            .saturating_sub(self.total_weighted_liabilities);
        clamp_to_u64(cmp::max(0, fc))
    }

    /// Credit usage in basis points (0 = no liabilities, 10_000 = 100%).
    /// Capped at 10_000 — an under-collateralized account is at 100% usage.
    /// Returns 0 if collateral is zero or negative.
    pub fn credit_usage_bps(&self) -> MathResult<u64> {
        if self.total_weighted_collateral <= 0 {
            return Ok(0);
        }
        let usage = self
            .total_weighted_liabilities
            .checked_mul(10_000)
            .ok_or(MathError::Overflow)?
            .checked_div(self.total_weighted_collateral)
            .ok_or(MathError::Overflow)?;
        // Clamp to 10_000 (100%)
        Ok(cmp::min(clamp_to_u64(cmp::max(0, usage)), 10_000))
    }
}

/// Combined withdraw and borrow limits for a single market.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PositionLimits {
    pub withdraw_limit: u64,
    pub borrow_limit: u64,
}

/// Calculate withdraw and borrow limits for a single market.
///
/// When `reduce_only` is true, borrowing is not allowed — `borrow_limit == withdraw_limit`.
/// `token_balance` should be the same balance used in `MarginState::calculate`.
pub fn calculate_position_limits(
    margin_state: &MarginState,
    spot_market: &SpotMarket,
    price_usdc_base_units: u64,
    token_balance: i128,
    reduce_only: bool,
) -> MathResult<PositionLimits> {
    if price_usdc_base_units == 0 {
        return Ok(PositionLimits {
            withdraw_limit: 0,
            borrow_limit: 0,
        });
    }

    let free_collateral = cmp::max(
        0,
        margin_state
            .total_weighted_collateral
            .saturating_sub(margin_state.total_weighted_liabilities),
    );
    let token_deposit_balance = clamp_to_u64(cmp::max(0, token_balance));
    let asset_weight = spot_market.initial_asset_weight;
    let (numerator_scale, denominator_scale) = decimal_scale(spot_market.decimals)?;

    // Withdraw limit: how much can be withdrawn without breaching margin.
    // No liabilities or zero asset weight -> can withdraw entire deposit.
    let withdraw_limit = if asset_weight == 0 || margin_state.total_weighted_liabilities == 0 {
        token_deposit_balance
    } else {
        let withdraw_limit_i128 = free_collateral
            .checked_mul(MARGIN_PRECISION)
            .and_then(|v| v.checked_div_ceil(asset_weight as i128))
            .and_then(|v| v.checked_mul(PRICE_PRECISION))
            .and_then(|v| v.checked_div_ceil(price_usdc_base_units as i128))
            .and_then(|v| v.checked_mul(numerator_scale as i128))
            .and_then(|v| v.checked_div(denominator_scale as i128))
            .ok_or(MathError::Overflow)?;

        cmp::min(
            token_deposit_balance,
            clamp_to_u64(cmp::max(0, withdraw_limit_i128)),
        )
    };

    if reduce_only {
        return Ok(PositionLimits {
            withdraw_limit,
            borrow_limit: withdraw_limit,
        });
    }

    // Borrow limit: withdraw_limit + max additional liability.
    // Match ts-sdk: subtract the full weighted value of the entire deposit position
    // from free collateral before computing max liability.
    let free_collateral_after = if token_balance > 0 {
        let position_value_usdc = calculate_value_usdc_base_units(
            token_balance,
            price_usdc_base_units,
            spot_market.decimals,
        )?;
        let weighted_position_value = position_value_usdc
            .checked_mul(asset_weight as i128)
            .and_then(|v| v.checked_div(MARGIN_PRECISION))
            .ok_or(MathError::Overflow)?;
        cmp::max(0, free_collateral.saturating_sub(weighted_position_value))
    } else {
        free_collateral
    };

    // Max additional liability the remaining collateral can support
    let liability_weight = spot_market.initial_liability_weight as i128;
    let max_liability = free_collateral_after
        .checked_mul(MARGIN_PRECISION)
        .and_then(|v| v.checked_div(liability_weight))
        .and_then(|v| v.checked_mul(PRICE_PRECISION))
        .and_then(|v| v.checked_div(price_usdc_base_units as i128))
        .and_then(|v| v.checked_mul(numerator_scale as i128))
        .and_then(|v| v.checked_div(denominator_scale as i128))
        .ok_or(MathError::Overflow)?;

    let borrow_limit_unclamped = (withdraw_limit as i128)
        .checked_add(max_liability)
        .ok_or(MathError::Overflow)?;

    Ok(PositionLimits {
        withdraw_limit,
        borrow_limit: clamp_to_u64(cmp::max(0, borrow_limit_unclamped)),
    })
}

/// Compute decimal scaling factors for converting between token precision and USDC (6 decimals).
fn decimal_scale(token_decimals: u32) -> MathResult<(u32, u32)> {
    if token_decimals > 6 {
        let numerator = 10u32
            .checked_pow(token_decimals.checked_sub(6).ok_or(MathError::Overflow)?)
            .ok_or(MathError::Overflow)?;
        Ok((numerator, 1))
    } else {
        let denominator = 10u32
            .checked_pow(
                6u32.checked_sub(token_decimals)
                    .ok_or(MathError::Overflow)?,
            )
            .ok_or(MathError::Overflow)?;
        Ok((1, denominator))
    }
}

/// Clamp an i128 value to u64 range.
fn clamp_to_u64(value: i128) -> u64 {
    if value < 0 {
        0
    } else if value > u64::MAX as i128 {
        u64::MAX
    } else {
        value as u64
    }
}

#[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(
        market_index: u16,
        decimals: u32,
        initial_asset_weight: u32,
        initial_liability_weight: u32,
    ) -> 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: 1_000_000,
            },
            oracle: None,
        }
    }

    fn usdc_market() -> SpotMarket {
        make_spot_market(0, 6, 10_000, 10_000)
    }

    fn sol_market() -> SpotMarket {
        make_spot_market(1, 9, 8_000, 12_000)
    }

    // --- MarginState tests ---

    #[test]
    fn empty_positions() {
        let state = MarginState::calculate(&[]).unwrap();
        assert_eq!(state.total_weighted_collateral, 0);
        assert_eq!(state.total_weighted_liabilities, 0);
        assert_eq!(state.free_collateral(), 0);
        assert_eq!(state.credit_usage_bps().unwrap(), 0);
    }

    #[test]
    fn single_deposit() {
        let market = usdc_market();
        let positions = [PositionData {
            token_balance: 1_000_000, // 1 USDC
            price_usdc_base_units: 1_000_000,
            twap5min: 1_000_000,
            spot_market: &market,
        }];
        let state = MarginState::calculate(&positions).unwrap();
        assert_eq!(state.total_weighted_collateral, 1_000_000);
        assert_eq!(state.total_weighted_liabilities, 0);
        assert_eq!(state.free_collateral(), 1_000_000);
        assert_eq!(state.credit_usage_bps().unwrap(), 0);
    }

    #[test]
    fn deposit_and_borrow() {
        let market = usdc_market();
        let positions = [
            PositionData {
                token_balance: 1_000_000, // 1 USDC deposit
                price_usdc_base_units: 1_000_000,
                twap5min: 1_000_000,
                spot_market: &market,
            },
            PositionData {
                token_balance: -500_000, // 0.5 USDC borrow
                price_usdc_base_units: 1_000_000,
                twap5min: 1_000_000,
                spot_market: &market,
            },
        ];
        let state = MarginState::calculate(&positions).unwrap();
        assert_eq!(state.total_weighted_collateral, 1_000_000);
        assert_eq!(state.total_weighted_liabilities, 500_000);
        assert_eq!(state.free_collateral(), 500_000);
        assert_eq!(state.credit_usage_bps().unwrap(), 5_000); // 50%
    }

    #[test]
    fn multi_market_positions() {
        let usdc = usdc_market();
        let sol = sol_market(); // 80% asset weight
        let positions = [
            PositionData {
                token_balance: 10_000_000, // 10 USDC
                price_usdc_base_units: 1_000_000,
                twap5min: 1_000_000,
                spot_market: &usdc,
            },
            PositionData {
                token_balance: 1_000_000_000,       // 1 SOL
                price_usdc_base_units: 100_000_000, // $100
                twap5min: 100_000_000,
                spot_market: &sol,
            },
        ];
        let state = MarginState::calculate(&positions).unwrap();
        // 10 USDC * 100% = 10 USDC weighted
        // 1 SOL * $100 * 80% = $80 weighted
        assert_eq!(state.total_weighted_collateral, 10_000_000 + 80_000_000);
        assert_eq!(state.total_weighted_liabilities, 0);
    }

    #[test]
    fn strict_pricing_for_assets() {
        let market = usdc_market();
        // TWAP is lower than oracle -> use TWAP for asset
        let positions = [PositionData {
            token_balance: 1_000_000,
            price_usdc_base_units: 1_100_000,
            twap5min: 1_000_000,
            spot_market: &market,
        }];
        let state = MarginState::calculate(&positions).unwrap();
        // Should use min(1_100_000, 1_000_000) = 1_000_000
        assert_eq!(state.total_weighted_collateral, 1_000_000);
    }

    #[test]
    fn zero_balance_skipped() {
        let market = usdc_market();
        let positions = [PositionData {
            token_balance: 0,
            price_usdc_base_units: 1_000_000,
            twap5min: 1_000_000,
            spot_market: &market,
        }];
        let state = MarginState::calculate(&positions).unwrap();
        assert_eq!(state.total_weighted_collateral, 0);
    }

    // --- free_collateral clamped to zero ---

    #[test]
    fn free_collateral_clamped_to_zero() {
        let state = MarginState {
            total_weighted_collateral: 5_000_000,
            total_weighted_liabilities: 10_000_000,
        };
        assert_eq!(state.free_collateral(), 0);
    }

    // --- credit_usage capped at 10_000 ---

    #[test]
    fn credit_usage_capped_at_10000() {
        let state = MarginState {
            total_weighted_collateral: 5_000_000,
            total_weighted_liabilities: 10_000_000,
        };
        assert_eq!(state.credit_usage_bps().unwrap(), 10_000);
    }

    // --- Position limits tests ---

    #[test]
    fn withdraw_limit_no_liabilities() {
        let market = usdc_market();
        let state = MarginState {
            total_weighted_collateral: 10_000_000,
            total_weighted_liabilities: 0,
        };
        let limits =
            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
        // No liabilities -> full deposit
        assert_eq!(limits.withdraw_limit, 10_000_000);
    }

    #[test]
    fn withdraw_limit_zero_asset_weight() {
        let market = make_spot_market(0, 6, 0, 10_000);
        let state = MarginState {
            total_weighted_collateral: 10_000_000,
            total_weighted_liabilities: 5_000_000,
        };
        let limits =
            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
        assert_eq!(limits.withdraw_limit, 10_000_000);
    }

    #[test]
    fn withdraw_limit_with_liabilities() {
        let market = usdc_market(); // 100% asset weight
        let state = MarginState {
            total_weighted_collateral: 10_000_000,
            total_weighted_liabilities: 5_000_000,
        };
        // free_collateral = 5_000_000
        // withdraw_limit = 5M * 10k / 10k * 1M / 1M * 1 / 1 = 5_000_000
        let limits =
            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
        assert_eq!(limits.withdraw_limit, 5_000_000);
    }

    #[test]
    fn withdraw_limit_capped_at_deposit() {
        let market = usdc_market();
        let state = MarginState {
            total_weighted_collateral: 100_000_000,
            total_weighted_liabilities: 1_000_000,
        };
        let deposit = 2_000_000i128;
        let limits = calculate_position_limits(&state, &market, 1_000_000, deposit, false).unwrap();
        assert_eq!(limits.withdraw_limit, deposit as u64);
    }

    #[test]
    fn withdraw_limit_zero_price() {
        let market = usdc_market();
        let state = MarginState {
            total_weighted_collateral: 10_000_000,
            total_weighted_liabilities: 5_000_000,
        };
        let limits = calculate_position_limits(&state, &market, 0, 10_000_000, false).unwrap();
        assert_eq!(limits.withdraw_limit, 0);
        assert_eq!(limits.borrow_limit, 0);
    }

    #[test]
    fn withdraw_limit_sol_decimals() {
        let market = sol_market(); // 9 decimals, 80% weight
        let state = MarginState {
            total_weighted_collateral: 100_000_000, // $100
            total_weighted_liabilities: 20_000_000, // $20
        };
        // free_collateral = $80
        // withdraw_limit = 80M * 10k / 8k * 1M / 100M * 1000 / 1
        //                = 100M * 0.01 * 1000 = 1_000_000_000 (1 SOL)
        let limits = calculate_position_limits(
            &state,
            &market,
            100_000_000,   // $100
            2_000_000_000, // 2 SOL deposit
            false,
        )
        .unwrap();
        assert_eq!(limits.withdraw_limit, 1_000_000_000); // 1 SOL
    }

    // --- Borrow limit tests ---

    #[test]
    fn borrow_limit_basic() {
        let market = usdc_market(); // 100% asset, 100% liability weight
        let state = MarginState {
            total_weighted_collateral: 10_000_000,
            total_weighted_liabilities: 0,
        };
        let limits =
            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
        // After subtracting full weighted position value (10M * 100% = 10M):
        // free_collateral_after = 10M - 10M = 0
        // max_liability = 0 -> borrow = withdraw_limit + 0 = 10M
        assert_eq!(limits.withdraw_limit, 10_000_000);
        assert_eq!(limits.borrow_limit, 10_000_000);
    }

    #[test]
    fn borrow_limit_with_collateral_headroom() {
        // SOL collateral backing USDC borrows
        let usdc = usdc_market();
        let state = MarginState {
            total_weighted_collateral: 80_000_000, // $80 weighted (1 SOL at $100, 80% weight)
            total_weighted_liabilities: 0,
        };
        // No USDC deposits -> token_balance = 0, so free_collateral_after = free_collateral
        let limits = calculate_position_limits(&state, &usdc, 1_000_000, 0, false).unwrap();
        // free_collateral_after = $80 (no position to subtract)
        // max_liability = 80M * 10k / 10k * 1M / 1M = 80M
        // borrow = 0 + 80M = $80
        assert_eq!(limits.withdraw_limit, 0);
        assert_eq!(limits.borrow_limit, 80_000_000);
    }

    #[test]
    fn borrow_limit_zero_asset_weight() {
        let market = make_spot_market(0, 6, 0, 10_000);
        let state = MarginState {
            total_weighted_collateral: 10_000_000,
            total_weighted_liabilities: 0,
        };
        let limits =
            calculate_position_limits(&state, &market, 1_000_000, 5_000_000, false).unwrap();
        // asset_weight=0 -> can withdraw full deposit
        assert_eq!(limits.withdraw_limit, 5_000_000);
        // Weighted position value is 0 (asset_weight=0), so full $10M collateral
        // (from other positions) remains to back borrowing:
        // borrow = 5M withdraw + 10M max_liability = 15M
        assert_eq!(limits.borrow_limit, 15_000_000);
    }

    // --- reduce_only tests ---

    #[test]
    fn usdc_reduce_only() {
        let market = usdc_market();
        let state = MarginState {
            total_weighted_collateral: 100_000_000,
            total_weighted_liabilities: 0,
        };
        let limits =
            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, true).unwrap();
        assert_eq!(limits.borrow_limit, limits.withdraw_limit);
    }

    // --- decimal_scale tests ---

    #[test]
    fn decimal_scale_usdc() {
        let (n, d) = decimal_scale(6).unwrap();
        assert_eq!((n, d), (1, 1));
    }

    #[test]
    fn decimal_scale_sol() {
        let (n, d) = decimal_scale(9).unwrap();
        assert_eq!((n, d), (1_000, 1));
    }

    #[test]
    fn decimal_scale_small() {
        let (n, d) = decimal_scale(4).unwrap();
        assert_eq!((n, d), (1, 100));
    }

    // --- clamp_to_u64 tests ---

    #[test]
    fn clamp_negative() {
        assert_eq!(clamp_to_u64(-100), 0);
    }

    #[test]
    fn clamp_overflow() {
        assert_eq!(clamp_to_u64(i128::from(u64::MAX) + 1), u64::MAX);
    }

    #[test]
    fn clamp_normal() {
        assert_eq!(clamp_to_u64(42), 42);
    }
}

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

    fn arb_usdc_market() -> SpotMarket {
        SpotMarket {
            pubkey: vec![],
            market_index: 0,
            initial_asset_weight: 10_000,
            initial_liability_weight: 10_000,
            imf_factor: 0,
            scale_initial_asset_weight_start: 0,
            decimals: 6,
            cumulative_deposit_interest: 10_000_000_000_000,
            cumulative_borrow_interest: 10_000_000_000_000,
            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: 1_000_000,
            },
            oracle: None,
        }
    }

    proptest! {
        #[test]
        fn withdraw_limit_le_deposit(
            collateral in 0i128..=1_000_000_000_000i128,
            liabilities in 0i128..=1_000_000_000_000i128,
            deposit in 0i128..=1_000_000_000_000i128,
        ) {
            let market = arb_usdc_market();
            let state = MarginState {
                total_weighted_collateral: collateral,
                total_weighted_liabilities: liabilities,
            };
            let limits = calculate_position_limits(&state, &market, 1_000_000, deposit, false).unwrap();
            let deposit_u64 = clamp_to_u64(std::cmp::max(0, deposit));
            prop_assert!(limits.withdraw_limit <= deposit_u64, "withdraw {} > deposit {}", limits.withdraw_limit, deposit_u64);
        }

        #[test]
        fn borrow_limit_ge_withdraw_limit(
            collateral in 1i128..=1_000_000_000_000i128,
            liabilities in 0i128..=500_000_000_000i128,
            deposit in 0i128..=1_000_000_000_000i128,
        ) {
            let market = arb_usdc_market();
            let state = MarginState {
                total_weighted_collateral: collateral,
                total_weighted_liabilities: liabilities,
            };
            let limits = calculate_position_limits(&state, &market, 1_000_000, deposit, false).unwrap();
            prop_assert!(limits.borrow_limit >= limits.withdraw_limit, "borrow {} < withdraw {}", limits.borrow_limit, limits.withdraw_limit);
        }

        #[test]
        fn credit_usage_bounded(
            collateral in 1i128..=1_000_000_000_000i128,
            liabilities in 0i128..=1_000_000_000_000i128,
        ) {
            let state = MarginState {
                total_weighted_collateral: collateral,
                total_weighted_liabilities: liabilities,
            };
            let usage = state.credit_usage_bps().unwrap();
            prop_assert!(usage <= 10_000, "usage {} > 10_000", usage);
        }

        #[test]
        fn free_collateral_matches_components(
            collateral in 0i128..=i128::MAX / 2,
            liabilities in 0i128..=i128::MAX / 2,
        ) {
            let state = MarginState {
                total_weighted_collateral: collateral,
                total_weighted_liabilities: liabilities,
            };
            let fc = state.free_collateral();
            let expected = collateral.saturating_sub(liabilities);
            let expected_u64 = if expected < 0 { 0u64 } else { clamp_to_u64(expected) };
            prop_assert_eq!(fc, expected_u64);
        }
    }
}