Skip to main content

pyra_margin/
limits.rs

1use std::cmp;
2
3use pyra_types::SpotMarket;
4
5use crate::balance::calculate_value_usdc_base_units;
6use crate::error::{MathError, MathResult};
7use crate::math::CheckedDivCeil;
8use crate::weights::{calculate_asset_weight, calculate_liability_weight, get_strict_price};
9
10const MARGIN_PRECISION: i128 = 10_000;
11const PRICE_PRECISION: i128 = 1_000_000;
12
13/// Input data for a single position in margin calculations.
14///
15/// The caller is responsible for computing `token_balance` — this allows
16/// different services to compose balances differently (e.g. Drift-only
17/// vs Drift + deposit addresses - pending withdraws).
18pub struct PositionData<'a> {
19    /// Token balance in base units (positive = deposit, negative = borrow).
20    /// Use `get_token_balance()` to compute from a `SpotPosition`, then add
21    /// any off-chain adjustments.
22    pub token_balance: i128,
23    /// Current oracle price in USDC base units (1e6).
24    pub price_usdc_base_units: u64,
25    /// 5-minute TWAP from `SpotMarket::historical_oracle_data.last_oracle_price_twap5min`.
26    pub twap5min: i64,
27    /// The spot market configuration.
28    pub spot_market: &'a SpotMarket,
29}
30
31/// Aggregated margin state across all positions.
32///
33/// Computed once from all positions, then used to derive per-market
34/// withdraw/borrow limits and credit usage.
35#[derive(Debug, Clone, Copy)]
36pub struct MarginState {
37    /// Total weighted collateral in USDC base units (after IMF weight application).
38    pub total_weighted_collateral: i128,
39    /// Total weighted liabilities in USDC base units (always non-negative, after IMF weight).
40    pub total_weighted_liabilities: i128,
41    /// Total unweighted collateral in USDC base units (raw position value before weights).
42    pub total_collateral: i128,
43    /// Total unweighted liabilities in USDC base units (always non-negative, before weights).
44    pub total_liabilities: i128,
45}
46
47impl MarginState {
48    /// Calculate margin state from a set of positions.
49    ///
50    /// For each position: applies strict oracle pricing, computes USDC value,
51    /// applies IMF-adjusted weights, and accumulates into collateral/liabilities.
52    pub fn calculate(positions: &[PositionData<'_>]) -> MathResult<Self> {
53        let mut total_weighted_collateral: i128 = 0;
54        let mut total_weighted_liabilities: i128 = 0;
55        let mut total_collateral: i128 = 0;
56        let mut total_liabilities: i128 = 0;
57
58        for pos in positions {
59            if pos.token_balance == 0 {
60                continue;
61            }
62
63            let is_asset = pos.token_balance >= 0;
64            let strict_price = get_strict_price(pos.price_usdc_base_units, pos.twap5min, is_asset);
65
66            let value_usdc = calculate_value_usdc_base_units(
67                pos.token_balance,
68                strict_price,
69                pos.spot_market.decimals,
70            )?;
71
72            // Accumulate unweighted totals
73            if value_usdc >= 0 {
74                total_collateral = total_collateral
75                    .checked_add(value_usdc)
76                    .ok_or(MathError::Overflow)?;
77            } else {
78                total_liabilities = total_liabilities
79                    .checked_add(value_usdc.checked_neg().ok_or(MathError::Overflow)?)
80                    .ok_or(MathError::Overflow)?;
81            }
82
83            let token_amount_unsigned = pos.token_balance.unsigned_abs();
84            let weight_bps = if is_asset {
85                calculate_asset_weight(
86                    token_amount_unsigned,
87                    pos.price_usdc_base_units,
88                    pos.spot_market,
89                )?
90            } else {
91                calculate_liability_weight(token_amount_unsigned, pos.spot_market)?
92            };
93
94            // weight_bps is u128 ≤ ~20_000, safe to cast to i128
95            let weighted_value = value_usdc
96                .checked_mul(weight_bps as i128)
97                .ok_or(MathError::Overflow)?
98                .checked_div(MARGIN_PRECISION)
99                .ok_or(MathError::Overflow)?;
100
101            if weighted_value >= 0 {
102                total_weighted_collateral = total_weighted_collateral
103                    .checked_add(weighted_value)
104                    .ok_or(MathError::Overflow)?;
105            } else {
106                total_weighted_liabilities = total_weighted_liabilities
107                    .checked_add(weighted_value.checked_neg().ok_or(MathError::Overflow)?)
108                    .ok_or(MathError::Overflow)?;
109            }
110        }
111
112        Ok(Self {
113            total_weighted_collateral,
114            total_weighted_liabilities,
115            total_collateral,
116            total_liabilities,
117        })
118    }
119
120    /// Free collateral = weighted collateral - weighted liabilities, clamped to 0.
121    pub fn free_collateral(&self) -> u64 {
122        let fc = self
123            .total_weighted_collateral
124            .saturating_sub(self.total_weighted_liabilities);
125        clamp_to_u64(cmp::max(0, fc))
126    }
127
128    /// Credit usage in basis points (0 = no liabilities, 10_000 = 100%).
129    /// Capped at 10_000 — an under-collateralized account is at 100% usage.
130    /// Returns 0 if collateral is zero or negative.
131    pub fn credit_usage_bps(&self) -> MathResult<u64> {
132        if self.total_weighted_collateral <= 0 {
133            return Ok(0);
134        }
135        let usage = self
136            .total_weighted_liabilities
137            .checked_mul(10_000)
138            .ok_or(MathError::Overflow)?
139            .checked_div(self.total_weighted_collateral)
140            .ok_or(MathError::Overflow)?;
141        // Clamp to 10_000 (100%)
142        Ok(cmp::min(clamp_to_u64(cmp::max(0, usage)), 10_000))
143    }
144}
145
146/// Combined withdraw and borrow limits for a single market.
147#[derive(Debug, Clone, Copy, PartialEq, Eq)]
148pub struct PositionLimits {
149    pub withdraw_limit: u64,
150    pub borrow_limit: u64,
151}
152
153/// Calculate withdraw and borrow limits for a single market.
154///
155/// When `reduce_only` is true, borrowing is not allowed — `borrow_limit == withdraw_limit`.
156/// `token_balance` should be the same balance used in `MarginState::calculate`.
157pub fn calculate_position_limits(
158    margin_state: &MarginState,
159    spot_market: &SpotMarket,
160    price_usdc_base_units: u64,
161    token_balance: i128,
162    reduce_only: bool,
163) -> MathResult<PositionLimits> {
164    if price_usdc_base_units == 0 {
165        return Ok(PositionLimits {
166            withdraw_limit: 0,
167            borrow_limit: 0,
168        });
169    }
170
171    let free_collateral = cmp::max(
172        0,
173        margin_state
174            .total_weighted_collateral
175            .saturating_sub(margin_state.total_weighted_liabilities),
176    );
177    let token_deposit_balance = clamp_to_u64(cmp::max(0, token_balance));
178    let asset_weight = spot_market.initial_asset_weight;
179    let (numerator_scale, denominator_scale) = decimal_scale(spot_market.decimals)?;
180
181    // Withdraw limit: how much can be withdrawn without breaching margin.
182    // No liabilities or zero asset weight -> can withdraw entire deposit.
183    let withdraw_limit = if asset_weight == 0 || margin_state.total_weighted_liabilities == 0 {
184        token_deposit_balance
185    } else {
186        let withdraw_limit_i128 = free_collateral
187            .checked_mul(MARGIN_PRECISION)
188            .and_then(|v| v.checked_div_ceil(asset_weight as i128))
189            .and_then(|v| v.checked_mul(PRICE_PRECISION))
190            .and_then(|v| v.checked_div_ceil(price_usdc_base_units as i128))
191            .and_then(|v| v.checked_mul(numerator_scale as i128))
192            .and_then(|v| v.checked_div(denominator_scale as i128))
193            .ok_or(MathError::Overflow)?;
194
195        cmp::min(
196            token_deposit_balance,
197            clamp_to_u64(cmp::max(0, withdraw_limit_i128)),
198        )
199    };
200
201    if reduce_only {
202        return Ok(PositionLimits {
203            withdraw_limit,
204            borrow_limit: withdraw_limit,
205        });
206    }
207
208    // Borrow limit: withdraw_limit + max additional liability.
209    // Match ts-sdk: subtract the full weighted value of the entire deposit position
210    // from free collateral before computing max liability.
211    let free_collateral_after = if token_balance > 0 {
212        let position_value_usdc = calculate_value_usdc_base_units(
213            token_balance,
214            price_usdc_base_units,
215            spot_market.decimals,
216        )?;
217        let weighted_position_value = position_value_usdc
218            .checked_mul(asset_weight as i128)
219            .and_then(|v| v.checked_div(MARGIN_PRECISION))
220            .ok_or(MathError::Overflow)?;
221        cmp::max(0, free_collateral.saturating_sub(weighted_position_value))
222    } else {
223        free_collateral
224    };
225
226    // Max additional liability the remaining collateral can support
227    let liability_weight = spot_market.initial_liability_weight as i128;
228    let max_liability = free_collateral_after
229        .checked_mul(MARGIN_PRECISION)
230        .and_then(|v| v.checked_div(liability_weight))
231        .and_then(|v| v.checked_mul(PRICE_PRECISION))
232        .and_then(|v| v.checked_div(price_usdc_base_units as i128))
233        .and_then(|v| v.checked_mul(numerator_scale as i128))
234        .and_then(|v| v.checked_div(denominator_scale as i128))
235        .ok_or(MathError::Overflow)?;
236
237    let borrow_limit_unclamped = (withdraw_limit as i128)
238        .checked_add(max_liability)
239        .ok_or(MathError::Overflow)?;
240
241    Ok(PositionLimits {
242        withdraw_limit,
243        borrow_limit: clamp_to_u64(cmp::max(0, borrow_limit_unclamped)),
244    })
245}
246
247/// Compute decimal scaling factors for converting between token precision and USDC (6 decimals).
248fn decimal_scale(token_decimals: u32) -> MathResult<(u32, u32)> {
249    if token_decimals > 6 {
250        let numerator = 10u32
251            .checked_pow(token_decimals.checked_sub(6).ok_or(MathError::Overflow)?)
252            .ok_or(MathError::Overflow)?;
253        Ok((numerator, 1))
254    } else {
255        let denominator = 10u32
256            .checked_pow(
257                6u32.checked_sub(token_decimals)
258                    .ok_or(MathError::Overflow)?,
259            )
260            .ok_or(MathError::Overflow)?;
261        Ok((1, denominator))
262    }
263}
264
265/// Clamp an i128 value to u64 range.
266fn clamp_to_u64(value: i128) -> u64 {
267    if value < 0 {
268        0
269    } else if value > u64::MAX as i128 {
270        u64::MAX
271    } else {
272        value as u64
273    }
274}
275
276#[cfg(test)]
277#[allow(
278    clippy::unwrap_used,
279    clippy::expect_used,
280    clippy::panic,
281    clippy::arithmetic_side_effects
282)]
283mod tests {
284    use super::*;
285    use pyra_types::{HistoricalOracleData, InsuranceFund};
286
287    fn make_spot_market(
288        market_index: u16,
289        decimals: u32,
290        initial_asset_weight: u32,
291        initial_liability_weight: u32,
292    ) -> SpotMarket {
293        let precision_decrease = 10u128.pow(19u32.saturating_sub(decimals));
294        SpotMarket {
295            pubkey: vec![],
296            market_index,
297            initial_asset_weight,
298            initial_liability_weight,
299            imf_factor: 0,
300            scale_initial_asset_weight_start: 0,
301            decimals,
302            cumulative_deposit_interest: precision_decrease,
303            cumulative_borrow_interest: precision_decrease,
304            deposit_balance: 0,
305            borrow_balance: 0,
306            optimal_utilization: 0,
307            optimal_borrow_rate: 0,
308            max_borrow_rate: 0,
309            min_borrow_rate: 0,
310            insurance_fund: InsuranceFund::default(),
311            historical_oracle_data: HistoricalOracleData {
312                last_oracle_price_twap5min: 1_000_000,
313            },
314            oracle: None,
315        }
316    }
317
318    fn usdc_market() -> SpotMarket {
319        make_spot_market(0, 6, 10_000, 10_000)
320    }
321
322    fn sol_market() -> SpotMarket {
323        make_spot_market(1, 9, 8_000, 12_000)
324    }
325
326    // --- MarginState tests ---
327
328    #[test]
329    fn empty_positions() {
330        let state = MarginState::calculate(&[]).unwrap();
331        assert_eq!(state.total_weighted_collateral, 0);
332        assert_eq!(state.total_weighted_liabilities, 0);
333        assert_eq!(state.free_collateral(), 0);
334        assert_eq!(state.credit_usage_bps().unwrap(), 0);
335    }
336
337    #[test]
338    fn single_deposit() {
339        let market = usdc_market();
340        let positions = [PositionData {
341            token_balance: 1_000_000, // 1 USDC
342            price_usdc_base_units: 1_000_000,
343            twap5min: 1_000_000,
344            spot_market: &market,
345        }];
346        let state = MarginState::calculate(&positions).unwrap();
347        assert_eq!(state.total_weighted_collateral, 1_000_000);
348        assert_eq!(state.total_weighted_liabilities, 0);
349        assert_eq!(state.total_collateral, 1_000_000);
350        assert_eq!(state.total_liabilities, 0);
351        assert_eq!(state.free_collateral(), 1_000_000);
352        assert_eq!(state.credit_usage_bps().unwrap(), 0);
353    }
354
355    #[test]
356    fn deposit_and_borrow() {
357        let market = usdc_market();
358        let positions = [
359            PositionData {
360                token_balance: 1_000_000, // 1 USDC deposit
361                price_usdc_base_units: 1_000_000,
362                twap5min: 1_000_000,
363                spot_market: &market,
364            },
365            PositionData {
366                token_balance: -500_000, // 0.5 USDC borrow
367                price_usdc_base_units: 1_000_000,
368                twap5min: 1_000_000,
369                spot_market: &market,
370            },
371        ];
372        let state = MarginState::calculate(&positions).unwrap();
373        assert_eq!(state.total_weighted_collateral, 1_000_000);
374        assert_eq!(state.total_weighted_liabilities, 500_000);
375        assert_eq!(state.total_collateral, 1_000_000);
376        assert_eq!(state.total_liabilities, 500_000);
377        assert_eq!(state.free_collateral(), 500_000);
378        assert_eq!(state.credit_usage_bps().unwrap(), 5_000); // 50%
379    }
380
381    #[test]
382    fn multi_market_positions() {
383        let usdc = usdc_market();
384        let sol = sol_market(); // 80% asset weight
385        let positions = [
386            PositionData {
387                token_balance: 10_000_000, // 10 USDC
388                price_usdc_base_units: 1_000_000,
389                twap5min: 1_000_000,
390                spot_market: &usdc,
391            },
392            PositionData {
393                token_balance: 1_000_000_000,       // 1 SOL
394                price_usdc_base_units: 100_000_000, // $100
395                twap5min: 100_000_000,
396                spot_market: &sol,
397            },
398        ];
399        let state = MarginState::calculate(&positions).unwrap();
400        // 10 USDC * 100% = 10 USDC weighted
401        // 1 SOL * $100 * 80% = $80 weighted
402        assert_eq!(state.total_weighted_collateral, 10_000_000 + 80_000_000);
403        assert_eq!(state.total_weighted_liabilities, 0);
404        // Unweighted: 10 USDC + $100 SOL = $110
405        assert_eq!(state.total_collateral, 10_000_000 + 100_000_000);
406        assert_eq!(state.total_liabilities, 0);
407    }
408
409    #[test]
410    fn strict_pricing_for_assets() {
411        let market = usdc_market();
412        // TWAP is lower than oracle -> use TWAP for asset
413        let positions = [PositionData {
414            token_balance: 1_000_000,
415            price_usdc_base_units: 1_100_000,
416            twap5min: 1_000_000,
417            spot_market: &market,
418        }];
419        let state = MarginState::calculate(&positions).unwrap();
420        // Should use min(1_100_000, 1_000_000) = 1_000_000
421        assert_eq!(state.total_weighted_collateral, 1_000_000);
422    }
423
424    #[test]
425    fn zero_balance_skipped() {
426        let market = usdc_market();
427        let positions = [PositionData {
428            token_balance: 0,
429            price_usdc_base_units: 1_000_000,
430            twap5min: 1_000_000,
431            spot_market: &market,
432        }];
433        let state = MarginState::calculate(&positions).unwrap();
434        assert_eq!(state.total_weighted_collateral, 0);
435    }
436
437    // --- free_collateral clamped to zero ---
438
439    #[test]
440    fn free_collateral_clamped_to_zero() {
441        let state = MarginState {
442            total_weighted_collateral: 5_000_000,
443            total_weighted_liabilities: 10_000_000,
444            total_collateral: 5_000_000,
445            total_liabilities: 10_000_000,
446        };
447        assert_eq!(state.free_collateral(), 0);
448    }
449
450    // --- credit_usage capped at 10_000 ---
451
452    #[test]
453    fn credit_usage_capped_at_10000() {
454        let state = MarginState {
455            total_weighted_collateral: 5_000_000,
456            total_weighted_liabilities: 10_000_000,
457            total_collateral: 5_000_000,
458            total_liabilities: 10_000_000,
459        };
460        assert_eq!(state.credit_usage_bps().unwrap(), 10_000);
461    }
462
463    // --- Position limits tests ---
464
465    #[test]
466    fn withdraw_limit_no_liabilities() {
467        let market = usdc_market();
468        let state = MarginState {
469            total_weighted_collateral: 10_000_000,
470            total_weighted_liabilities: 0,
471            total_collateral: 10_000_000,
472            total_liabilities: 0,
473        };
474        let limits =
475            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
476        // No liabilities -> full deposit
477        assert_eq!(limits.withdraw_limit, 10_000_000);
478    }
479
480    #[test]
481    fn withdraw_limit_zero_asset_weight() {
482        let market = make_spot_market(0, 6, 0, 10_000);
483        let state = MarginState {
484            total_weighted_collateral: 10_000_000,
485            total_weighted_liabilities: 5_000_000,
486            total_collateral: 10_000_000,
487            total_liabilities: 5_000_000,
488        };
489        let limits =
490            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
491        assert_eq!(limits.withdraw_limit, 10_000_000);
492    }
493
494    #[test]
495    fn withdraw_limit_with_liabilities() {
496        let market = usdc_market(); // 100% asset weight
497        let state = MarginState {
498            total_weighted_collateral: 10_000_000,
499            total_weighted_liabilities: 5_000_000,
500            total_collateral: 10_000_000,
501            total_liabilities: 5_000_000,
502        };
503        // free_collateral = 5_000_000
504        // withdraw_limit = 5M * 10k / 10k * 1M / 1M * 1 / 1 = 5_000_000
505        let limits =
506            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
507        assert_eq!(limits.withdraw_limit, 5_000_000);
508    }
509
510    #[test]
511    fn withdraw_limit_capped_at_deposit() {
512        let market = usdc_market();
513        let state = MarginState {
514            total_weighted_collateral: 100_000_000,
515            total_weighted_liabilities: 1_000_000,
516            total_collateral: 100_000_000,
517            total_liabilities: 1_000_000,
518        };
519        let deposit = 2_000_000i128;
520        let limits = calculate_position_limits(&state, &market, 1_000_000, deposit, false).unwrap();
521        assert_eq!(limits.withdraw_limit, deposit as u64);
522    }
523
524    #[test]
525    fn withdraw_limit_zero_price() {
526        let market = usdc_market();
527        let state = MarginState {
528            total_weighted_collateral: 10_000_000,
529            total_weighted_liabilities: 5_000_000,
530            total_collateral: 10_000_000,
531            total_liabilities: 5_000_000,
532        };
533        let limits = calculate_position_limits(&state, &market, 0, 10_000_000, false).unwrap();
534        assert_eq!(limits.withdraw_limit, 0);
535        assert_eq!(limits.borrow_limit, 0);
536    }
537
538    #[test]
539    fn withdraw_limit_sol_decimals() {
540        let market = sol_market(); // 9 decimals, 80% weight
541        let state = MarginState {
542            total_weighted_collateral: 100_000_000, // $100
543            total_weighted_liabilities: 20_000_000, // $20
544            total_collateral: 125_000_000,
545            total_liabilities: 20_000_000,
546        };
547        // free_collateral = $80
548        // withdraw_limit = 80M * 10k / 8k * 1M / 100M * 1000 / 1
549        //                = 100M * 0.01 * 1000 = 1_000_000_000 (1 SOL)
550        let limits = calculate_position_limits(
551            &state,
552            &market,
553            100_000_000,   // $100
554            2_000_000_000, // 2 SOL deposit
555            false,
556        )
557        .unwrap();
558        assert_eq!(limits.withdraw_limit, 1_000_000_000); // 1 SOL
559    }
560
561    // --- Borrow limit tests ---
562
563    #[test]
564    fn borrow_limit_basic() {
565        let market = usdc_market(); // 100% asset, 100% liability weight
566        let state = MarginState {
567            total_weighted_collateral: 10_000_000,
568            total_weighted_liabilities: 0,
569            total_collateral: 10_000_000,
570            total_liabilities: 0,
571        };
572        let limits =
573            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
574        // After subtracting full weighted position value (10M * 100% = 10M):
575        // free_collateral_after = 10M - 10M = 0
576        // max_liability = 0 -> borrow = withdraw_limit + 0 = 10M
577        assert_eq!(limits.withdraw_limit, 10_000_000);
578        assert_eq!(limits.borrow_limit, 10_000_000);
579    }
580
581    #[test]
582    fn borrow_limit_with_collateral_headroom() {
583        // SOL collateral backing USDC borrows
584        let usdc = usdc_market();
585        let state = MarginState {
586            total_weighted_collateral: 80_000_000, // $80 weighted (1 SOL at $100, 80% weight)
587            total_weighted_liabilities: 0,
588            total_collateral: 100_000_000,
589            total_liabilities: 0,
590        };
591        // No USDC deposits -> token_balance = 0, so free_collateral_after = free_collateral
592        let limits = calculate_position_limits(&state, &usdc, 1_000_000, 0, false).unwrap();
593        // free_collateral_after = $80 (no position to subtract)
594        // max_liability = 80M * 10k / 10k * 1M / 1M = 80M
595        // borrow = 0 + 80M = $80
596        assert_eq!(limits.withdraw_limit, 0);
597        assert_eq!(limits.borrow_limit, 80_000_000);
598    }
599
600    #[test]
601    fn borrow_limit_zero_asset_weight() {
602        let market = make_spot_market(0, 6, 0, 10_000);
603        let state = MarginState {
604            total_weighted_collateral: 10_000_000,
605            total_weighted_liabilities: 0,
606            total_collateral: 10_000_000,
607            total_liabilities: 0,
608        };
609        let limits =
610            calculate_position_limits(&state, &market, 1_000_000, 5_000_000, false).unwrap();
611        // asset_weight=0 -> can withdraw full deposit
612        assert_eq!(limits.withdraw_limit, 5_000_000);
613        // Weighted position value is 0 (asset_weight=0), so full $10M collateral
614        // (from other positions) remains to back borrowing:
615        // borrow = 5M withdraw + 10M max_liability = 15M
616        assert_eq!(limits.borrow_limit, 15_000_000);
617    }
618
619    // --- reduce_only tests ---
620
621    #[test]
622    fn usdc_reduce_only() {
623        let market = usdc_market();
624        let state = MarginState {
625            total_weighted_collateral: 100_000_000,
626            total_weighted_liabilities: 0,
627            total_collateral: 100_000_000,
628            total_liabilities: 0,
629        };
630        let limits =
631            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, true).unwrap();
632        assert_eq!(limits.borrow_limit, limits.withdraw_limit);
633    }
634
635    // --- decimal_scale tests ---
636
637    #[test]
638    fn decimal_scale_usdc() {
639        let (n, d) = decimal_scale(6).unwrap();
640        assert_eq!((n, d), (1, 1));
641    }
642
643    #[test]
644    fn decimal_scale_sol() {
645        let (n, d) = decimal_scale(9).unwrap();
646        assert_eq!((n, d), (1_000, 1));
647    }
648
649    #[test]
650    fn decimal_scale_small() {
651        let (n, d) = decimal_scale(4).unwrap();
652        assert_eq!((n, d), (1, 100));
653    }
654
655    // --- clamp_to_u64 tests ---
656
657    #[test]
658    fn clamp_negative() {
659        assert_eq!(clamp_to_u64(-100), 0);
660    }
661
662    #[test]
663    fn clamp_overflow() {
664        assert_eq!(clamp_to_u64(i128::from(u64::MAX) + 1), u64::MAX);
665    }
666
667    #[test]
668    fn clamp_normal() {
669        assert_eq!(clamp_to_u64(42), 42);
670    }
671}
672
673#[cfg(test)]
674#[allow(
675    clippy::unwrap_used,
676    clippy::expect_used,
677    clippy::panic,
678    clippy::arithmetic_side_effects
679)]
680mod proptests {
681    use super::*;
682    use proptest::prelude::*;
683    use pyra_types::{HistoricalOracleData, InsuranceFund};
684
685    fn arb_usdc_market() -> SpotMarket {
686        SpotMarket {
687            pubkey: vec![],
688            market_index: 0,
689            initial_asset_weight: 10_000,
690            initial_liability_weight: 10_000,
691            imf_factor: 0,
692            scale_initial_asset_weight_start: 0,
693            decimals: 6,
694            cumulative_deposit_interest: 10_000_000_000_000,
695            cumulative_borrow_interest: 10_000_000_000_000,
696            deposit_balance: 0,
697            borrow_balance: 0,
698            optimal_utilization: 0,
699            optimal_borrow_rate: 0,
700            max_borrow_rate: 0,
701            min_borrow_rate: 0,
702            insurance_fund: InsuranceFund::default(),
703            historical_oracle_data: HistoricalOracleData {
704                last_oracle_price_twap5min: 1_000_000,
705            },
706            oracle: None,
707        }
708    }
709
710    proptest! {
711        #[test]
712        fn withdraw_limit_le_deposit(
713            collateral in 0i128..=1_000_000_000_000i128,
714            liabilities in 0i128..=1_000_000_000_000i128,
715            deposit in 0i128..=1_000_000_000_000i128,
716        ) {
717            let market = arb_usdc_market();
718            let state = MarginState {
719                total_weighted_collateral: collateral,
720                total_weighted_liabilities: liabilities,
721                total_collateral: collateral,
722                total_liabilities: liabilities,
723            };
724            let limits = calculate_position_limits(&state, &market, 1_000_000, deposit, false).unwrap();
725            let deposit_u64 = clamp_to_u64(std::cmp::max(0, deposit));
726            prop_assert!(limits.withdraw_limit <= deposit_u64, "withdraw {} > deposit {}", limits.withdraw_limit, deposit_u64);
727        }
728
729        #[test]
730        fn borrow_limit_ge_withdraw_limit(
731            collateral in 1i128..=1_000_000_000_000i128,
732            liabilities in 0i128..=500_000_000_000i128,
733            deposit in 0i128..=1_000_000_000_000i128,
734        ) {
735            let market = arb_usdc_market();
736            let state = MarginState {
737                total_weighted_collateral: collateral,
738                total_weighted_liabilities: liabilities,
739                total_collateral: collateral,
740                total_liabilities: liabilities,
741            };
742            let limits = calculate_position_limits(&state, &market, 1_000_000, deposit, false).unwrap();
743            prop_assert!(limits.borrow_limit >= limits.withdraw_limit, "borrow {} < withdraw {}", limits.borrow_limit, limits.withdraw_limit);
744        }
745
746        #[test]
747        fn credit_usage_bounded(
748            collateral in 1i128..=1_000_000_000_000i128,
749            liabilities in 0i128..=1_000_000_000_000i128,
750        ) {
751            let state = MarginState {
752                total_weighted_collateral: collateral,
753                total_weighted_liabilities: liabilities,
754                total_collateral: collateral,
755                total_liabilities: liabilities,
756            };
757            let usage = state.credit_usage_bps().unwrap();
758            prop_assert!(usage <= 10_000, "usage {} > 10_000", usage);
759        }
760
761        #[test]
762        fn free_collateral_matches_components(
763            collateral in 0i128..=i128::MAX / 2,
764            liabilities in 0i128..=i128::MAX / 2,
765        ) {
766            let state = MarginState {
767                total_weighted_collateral: collateral,
768                total_weighted_liabilities: liabilities,
769                total_collateral: collateral,
770                total_liabilities: liabilities,
771            };
772            let fc = state.free_collateral();
773            let expected = collateral.saturating_sub(liabilities);
774            let expected_u64 = if expected < 0 { 0u64 } else { clamp_to_u64(expected) };
775            prop_assert_eq!(fc, expected_u64);
776        }
777    }
778}