Skip to main content

pyra_margin/drift/
limits.rs

1use std::cmp;
2
3use pyra_types::SpotMarket;
4
5use super::balance::calculate_value_usdc_base_units;
6use super::weights::{calculate_asset_weight, calculate_liability_weight, get_strict_price};
7use crate::error::{MathError, MathResult};
8use crate::math::CheckedDivCeil;
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::allow_attributes,
279    clippy::allow_attributes_without_reason,
280    clippy::unwrap_used,
281    clippy::expect_used,
282    clippy::panic,
283    clippy::arithmetic_side_effects,
284    reason = "test code"
285)]
286mod tests {
287    use super::*;
288    use pyra_types::{HistoricalOracleData, InsuranceFund};
289
290    fn make_spot_market(
291        market_index: u16,
292        decimals: u32,
293        initial_asset_weight: u32,
294        initial_liability_weight: u32,
295    ) -> SpotMarket {
296        let precision_decrease = 10u128.pow(19u32.saturating_sub(decimals));
297        SpotMarket {
298            pubkey: vec![],
299            market_index,
300            initial_asset_weight,
301            initial_liability_weight,
302            imf_factor: 0,
303            scale_initial_asset_weight_start: 0,
304            decimals,
305            cumulative_deposit_interest: precision_decrease,
306            cumulative_borrow_interest: precision_decrease,
307            deposit_balance: 0,
308            borrow_balance: 0,
309            optimal_utilization: 0,
310            optimal_borrow_rate: 0,
311            max_borrow_rate: 0,
312            min_borrow_rate: 0,
313            insurance_fund: InsuranceFund::default(),
314            historical_oracle_data: HistoricalOracleData {
315                last_oracle_price_twap5min: 1_000_000,
316            },
317            oracle: None,
318        }
319    }
320
321    fn usdc_market() -> SpotMarket {
322        make_spot_market(0, 6, 10_000, 10_000)
323    }
324
325    fn sol_market() -> SpotMarket {
326        make_spot_market(1, 9, 8_000, 12_000)
327    }
328
329    // --- MarginState tests ---
330
331    #[test]
332    fn empty_positions() {
333        let state = MarginState::calculate(&[]).unwrap();
334        assert_eq!(state.total_weighted_collateral, 0);
335        assert_eq!(state.total_weighted_liabilities, 0);
336        assert_eq!(state.free_collateral(), 0);
337        assert_eq!(state.credit_usage_bps().unwrap(), 0);
338    }
339
340    #[test]
341    fn single_deposit() {
342        let market = usdc_market();
343        let positions = [PositionData {
344            token_balance: 1_000_000, // 1 USDC
345            price_usdc_base_units: 1_000_000,
346            twap5min: 1_000_000,
347            spot_market: &market,
348        }];
349        let state = MarginState::calculate(&positions).unwrap();
350        assert_eq!(state.total_weighted_collateral, 1_000_000);
351        assert_eq!(state.total_weighted_liabilities, 0);
352        assert_eq!(state.total_collateral, 1_000_000);
353        assert_eq!(state.total_liabilities, 0);
354        assert_eq!(state.free_collateral(), 1_000_000);
355        assert_eq!(state.credit_usage_bps().unwrap(), 0);
356    }
357
358    #[test]
359    fn deposit_and_borrow() {
360        let market = usdc_market();
361        let positions = [
362            PositionData {
363                token_balance: 1_000_000, // 1 USDC deposit
364                price_usdc_base_units: 1_000_000,
365                twap5min: 1_000_000,
366                spot_market: &market,
367            },
368            PositionData {
369                token_balance: -500_000, // 0.5 USDC borrow
370                price_usdc_base_units: 1_000_000,
371                twap5min: 1_000_000,
372                spot_market: &market,
373            },
374        ];
375        let state = MarginState::calculate(&positions).unwrap();
376        assert_eq!(state.total_weighted_collateral, 1_000_000);
377        assert_eq!(state.total_weighted_liabilities, 500_000);
378        assert_eq!(state.total_collateral, 1_000_000);
379        assert_eq!(state.total_liabilities, 500_000);
380        assert_eq!(state.free_collateral(), 500_000);
381        assert_eq!(state.credit_usage_bps().unwrap(), 5_000); // 50%
382    }
383
384    #[test]
385    fn multi_market_positions() {
386        let usdc = usdc_market();
387        let sol = sol_market(); // 80% asset weight
388        let positions = [
389            PositionData {
390                token_balance: 10_000_000, // 10 USDC
391                price_usdc_base_units: 1_000_000,
392                twap5min: 1_000_000,
393                spot_market: &usdc,
394            },
395            PositionData {
396                token_balance: 1_000_000_000,       // 1 SOL
397                price_usdc_base_units: 100_000_000, // $100
398                twap5min: 100_000_000,
399                spot_market: &sol,
400            },
401        ];
402        let state = MarginState::calculate(&positions).unwrap();
403        // 10 USDC * 100% = 10 USDC weighted
404        // 1 SOL * $100 * 80% = $80 weighted
405        assert_eq!(state.total_weighted_collateral, 10_000_000 + 80_000_000);
406        assert_eq!(state.total_weighted_liabilities, 0);
407        // Unweighted: 10 USDC + $100 SOL = $110
408        assert_eq!(state.total_collateral, 10_000_000 + 100_000_000);
409        assert_eq!(state.total_liabilities, 0);
410    }
411
412    #[test]
413    fn strict_pricing_for_assets() {
414        let market = usdc_market();
415        // TWAP is lower than oracle -> use TWAP for asset
416        let positions = [PositionData {
417            token_balance: 1_000_000,
418            price_usdc_base_units: 1_100_000,
419            twap5min: 1_000_000,
420            spot_market: &market,
421        }];
422        let state = MarginState::calculate(&positions).unwrap();
423        // Should use min(1_100_000, 1_000_000) = 1_000_000
424        assert_eq!(state.total_weighted_collateral, 1_000_000);
425    }
426
427    #[test]
428    fn zero_balance_skipped() {
429        let market = usdc_market();
430        let positions = [PositionData {
431            token_balance: 0,
432            price_usdc_base_units: 1_000_000,
433            twap5min: 1_000_000,
434            spot_market: &market,
435        }];
436        let state = MarginState::calculate(&positions).unwrap();
437        assert_eq!(state.total_weighted_collateral, 0);
438    }
439
440    // --- free_collateral clamped to zero ---
441
442    #[test]
443    fn free_collateral_clamped_to_zero() {
444        let state = MarginState {
445            total_weighted_collateral: 5_000_000,
446            total_weighted_liabilities: 10_000_000,
447            total_collateral: 5_000_000,
448            total_liabilities: 10_000_000,
449        };
450        assert_eq!(state.free_collateral(), 0);
451    }
452
453    // --- credit_usage capped at 10_000 ---
454
455    #[test]
456    fn credit_usage_capped_at_10000() {
457        let state = MarginState {
458            total_weighted_collateral: 5_000_000,
459            total_weighted_liabilities: 10_000_000,
460            total_collateral: 5_000_000,
461            total_liabilities: 10_000_000,
462        };
463        assert_eq!(state.credit_usage_bps().unwrap(), 10_000);
464    }
465
466    // --- Position limits tests ---
467
468    #[test]
469    fn withdraw_limit_no_liabilities() {
470        let market = usdc_market();
471        let state = MarginState {
472            total_weighted_collateral: 10_000_000,
473            total_weighted_liabilities: 0,
474            total_collateral: 10_000_000,
475            total_liabilities: 0,
476        };
477        let limits =
478            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
479        // No liabilities -> full deposit
480        assert_eq!(limits.withdraw_limit, 10_000_000);
481    }
482
483    #[test]
484    fn withdraw_limit_zero_asset_weight() {
485        let market = make_spot_market(0, 6, 0, 10_000);
486        let state = MarginState {
487            total_weighted_collateral: 10_000_000,
488            total_weighted_liabilities: 5_000_000,
489            total_collateral: 10_000_000,
490            total_liabilities: 5_000_000,
491        };
492        let limits =
493            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
494        assert_eq!(limits.withdraw_limit, 10_000_000);
495    }
496
497    #[test]
498    fn withdraw_limit_with_liabilities() {
499        let market = usdc_market(); // 100% asset weight
500        let state = MarginState {
501            total_weighted_collateral: 10_000_000,
502            total_weighted_liabilities: 5_000_000,
503            total_collateral: 10_000_000,
504            total_liabilities: 5_000_000,
505        };
506        // free_collateral = 5_000_000
507        // withdraw_limit = 5M * 10k / 10k * 1M / 1M * 1 / 1 = 5_000_000
508        let limits =
509            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
510        assert_eq!(limits.withdraw_limit, 5_000_000);
511    }
512
513    #[test]
514    fn withdraw_limit_capped_at_deposit() {
515        let market = usdc_market();
516        let state = MarginState {
517            total_weighted_collateral: 100_000_000,
518            total_weighted_liabilities: 1_000_000,
519            total_collateral: 100_000_000,
520            total_liabilities: 1_000_000,
521        };
522        let deposit = 2_000_000i128;
523        let limits = calculate_position_limits(&state, &market, 1_000_000, deposit, false).unwrap();
524        assert_eq!(limits.withdraw_limit, deposit as u64);
525    }
526
527    #[test]
528    fn withdraw_limit_zero_price() {
529        let market = usdc_market();
530        let state = MarginState {
531            total_weighted_collateral: 10_000_000,
532            total_weighted_liabilities: 5_000_000,
533            total_collateral: 10_000_000,
534            total_liabilities: 5_000_000,
535        };
536        let limits = calculate_position_limits(&state, &market, 0, 10_000_000, false).unwrap();
537        assert_eq!(limits.withdraw_limit, 0);
538        assert_eq!(limits.borrow_limit, 0);
539    }
540
541    #[test]
542    fn withdraw_limit_sol_decimals() {
543        let market = sol_market(); // 9 decimals, 80% weight
544        let state = MarginState {
545            total_weighted_collateral: 100_000_000, // $100
546            total_weighted_liabilities: 20_000_000, // $20
547            total_collateral: 125_000_000,
548            total_liabilities: 20_000_000,
549        };
550        // free_collateral = $80
551        // withdraw_limit = 80M * 10k / 8k * 1M / 100M * 1000 / 1
552        //                = 100M * 0.01 * 1000 = 1_000_000_000 (1 SOL)
553        let limits = calculate_position_limits(
554            &state,
555            &market,
556            100_000_000,   // $100
557            2_000_000_000, // 2 SOL deposit
558            false,
559        )
560        .unwrap();
561        assert_eq!(limits.withdraw_limit, 1_000_000_000); // 1 SOL
562    }
563
564    // --- Borrow limit tests ---
565
566    #[test]
567    fn borrow_limit_basic() {
568        let market = usdc_market(); // 100% asset, 100% liability weight
569        let state = MarginState {
570            total_weighted_collateral: 10_000_000,
571            total_weighted_liabilities: 0,
572            total_collateral: 10_000_000,
573            total_liabilities: 0,
574        };
575        let limits =
576            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, false).unwrap();
577        // After subtracting full weighted position value (10M * 100% = 10M):
578        // free_collateral_after = 10M - 10M = 0
579        // max_liability = 0 -> borrow = withdraw_limit + 0 = 10M
580        assert_eq!(limits.withdraw_limit, 10_000_000);
581        assert_eq!(limits.borrow_limit, 10_000_000);
582    }
583
584    #[test]
585    fn borrow_limit_with_collateral_headroom() {
586        // SOL collateral backing USDC borrows
587        let usdc = usdc_market();
588        let state = MarginState {
589            total_weighted_collateral: 80_000_000, // $80 weighted (1 SOL at $100, 80% weight)
590            total_weighted_liabilities: 0,
591            total_collateral: 100_000_000,
592            total_liabilities: 0,
593        };
594        // No USDC deposits -> token_balance = 0, so free_collateral_after = free_collateral
595        let limits = calculate_position_limits(&state, &usdc, 1_000_000, 0, false).unwrap();
596        // free_collateral_after = $80 (no position to subtract)
597        // max_liability = 80M * 10k / 10k * 1M / 1M = 80M
598        // borrow = 0 + 80M = $80
599        assert_eq!(limits.withdraw_limit, 0);
600        assert_eq!(limits.borrow_limit, 80_000_000);
601    }
602
603    #[test]
604    fn borrow_limit_zero_asset_weight() {
605        let market = make_spot_market(0, 6, 0, 10_000);
606        let state = MarginState {
607            total_weighted_collateral: 10_000_000,
608            total_weighted_liabilities: 0,
609            total_collateral: 10_000_000,
610            total_liabilities: 0,
611        };
612        let limits =
613            calculate_position_limits(&state, &market, 1_000_000, 5_000_000, false).unwrap();
614        // asset_weight=0 -> can withdraw full deposit
615        assert_eq!(limits.withdraw_limit, 5_000_000);
616        // Weighted position value is 0 (asset_weight=0), so full $10M collateral
617        // (from other positions) remains to back borrowing:
618        // borrow = 5M withdraw + 10M max_liability = 15M
619        assert_eq!(limits.borrow_limit, 15_000_000);
620    }
621
622    // --- reduce_only tests ---
623
624    #[test]
625    fn usdc_reduce_only() {
626        let market = usdc_market();
627        let state = MarginState {
628            total_weighted_collateral: 100_000_000,
629            total_weighted_liabilities: 0,
630            total_collateral: 100_000_000,
631            total_liabilities: 0,
632        };
633        let limits =
634            calculate_position_limits(&state, &market, 1_000_000, 10_000_000, true).unwrap();
635        assert_eq!(limits.borrow_limit, limits.withdraw_limit);
636    }
637
638    // --- decimal_scale tests ---
639
640    #[test]
641    fn decimal_scale_usdc() {
642        let (n, d) = decimal_scale(6).unwrap();
643        assert_eq!((n, d), (1, 1));
644    }
645
646    #[test]
647    fn decimal_scale_sol() {
648        let (n, d) = decimal_scale(9).unwrap();
649        assert_eq!((n, d), (1_000, 1));
650    }
651
652    #[test]
653    fn decimal_scale_small() {
654        let (n, d) = decimal_scale(4).unwrap();
655        assert_eq!((n, d), (1, 100));
656    }
657
658    // --- clamp_to_u64 tests ---
659
660    #[test]
661    fn clamp_negative() {
662        assert_eq!(clamp_to_u64(-100), 0);
663    }
664
665    #[test]
666    fn clamp_overflow() {
667        assert_eq!(clamp_to_u64(i128::from(u64::MAX) + 1), u64::MAX);
668    }
669
670    #[test]
671    fn clamp_normal() {
672        assert_eq!(clamp_to_u64(42), 42);
673    }
674}
675
676#[cfg(test)]
677#[allow(
678    clippy::allow_attributes,
679    clippy::allow_attributes_without_reason,
680    clippy::unwrap_used,
681    clippy::expect_used,
682    clippy::panic,
683    clippy::arithmetic_side_effects,
684    reason = "test code"
685)]
686mod proptests {
687    use super::*;
688    use proptest::prelude::*;
689    use pyra_types::{HistoricalOracleData, InsuranceFund};
690
691    fn arb_usdc_market() -> SpotMarket {
692        SpotMarket {
693            pubkey: vec![],
694            market_index: 0,
695            initial_asset_weight: 10_000,
696            initial_liability_weight: 10_000,
697            imf_factor: 0,
698            scale_initial_asset_weight_start: 0,
699            decimals: 6,
700            cumulative_deposit_interest: 10_000_000_000_000,
701            cumulative_borrow_interest: 10_000_000_000_000,
702            deposit_balance: 0,
703            borrow_balance: 0,
704            optimal_utilization: 0,
705            optimal_borrow_rate: 0,
706            max_borrow_rate: 0,
707            min_borrow_rate: 0,
708            insurance_fund: InsuranceFund::default(),
709            historical_oracle_data: HistoricalOracleData {
710                last_oracle_price_twap5min: 1_000_000,
711            },
712            oracle: None,
713        }
714    }
715
716    proptest! {
717        #[test]
718        fn withdraw_limit_le_deposit(
719            collateral in 0i128..=1_000_000_000_000i128,
720            liabilities in 0i128..=1_000_000_000_000i128,
721            deposit in 0i128..=1_000_000_000_000i128,
722        ) {
723            let market = arb_usdc_market();
724            let state = MarginState {
725                total_weighted_collateral: collateral,
726                total_weighted_liabilities: liabilities,
727                total_collateral: collateral,
728                total_liabilities: liabilities,
729            };
730            let limits = calculate_position_limits(&state, &market, 1_000_000, deposit, false).unwrap();
731            let deposit_u64 = clamp_to_u64(std::cmp::max(0, deposit));
732            prop_assert!(limits.withdraw_limit <= deposit_u64, "withdraw {} > deposit {}", limits.withdraw_limit, deposit_u64);
733        }
734
735        #[test]
736        fn borrow_limit_ge_withdraw_limit(
737            collateral in 1i128..=1_000_000_000_000i128,
738            liabilities in 0i128..=500_000_000_000i128,
739            deposit in 0i128..=1_000_000_000_000i128,
740        ) {
741            let market = arb_usdc_market();
742            let state = MarginState {
743                total_weighted_collateral: collateral,
744                total_weighted_liabilities: liabilities,
745                total_collateral: collateral,
746                total_liabilities: liabilities,
747            };
748            let limits = calculate_position_limits(&state, &market, 1_000_000, deposit, false).unwrap();
749            prop_assert!(limits.borrow_limit >= limits.withdraw_limit, "borrow {} < withdraw {}", limits.borrow_limit, limits.withdraw_limit);
750        }
751
752        #[test]
753        fn credit_usage_bounded(
754            collateral in 1i128..=1_000_000_000_000i128,
755            liabilities in 0i128..=1_000_000_000_000i128,
756        ) {
757            let state = MarginState {
758                total_weighted_collateral: collateral,
759                total_weighted_liabilities: liabilities,
760                total_collateral: collateral,
761                total_liabilities: liabilities,
762            };
763            let usage = state.credit_usage_bps().unwrap();
764            prop_assert!(usage <= 10_000, "usage {} > 10_000", usage);
765        }
766
767        #[test]
768        fn free_collateral_matches_components(
769            collateral in 0i128..=i128::MAX / 2,
770            liabilities in 0i128..=i128::MAX / 2,
771        ) {
772            let state = MarginState {
773                total_weighted_collateral: collateral,
774                total_weighted_liabilities: liabilities,
775                total_collateral: collateral,
776                total_liabilities: liabilities,
777            };
778            let fc = state.free_collateral();
779            let expected = collateral.saturating_sub(liabilities);
780            let expected_u64 = if expected < 0 { 0u64 } else { clamp_to_u64(expected) };
781            prop_assert_eq!(fc, expected_u64);
782        }
783    }
784}