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
13pub struct PositionData<'a> {
19 pub token_balance: i128,
23 pub price_usdc_base_units: u64,
25 pub twap5min: i64,
27 pub spot_market: &'a SpotMarket,
29}
30
31#[derive(Debug, Clone, Copy)]
36pub struct MarginState {
37 pub total_weighted_collateral: i128,
39 pub total_weighted_liabilities: i128,
41 pub total_collateral: i128,
43 pub total_liabilities: i128,
45}
46
47impl MarginState {
48 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 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 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 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 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 Ok(cmp::min(clamp_to_u64(cmp::max(0, usage)), 10_000))
143 }
144}
145
146#[derive(Debug, Clone, Copy, PartialEq, Eq)]
148pub struct PositionLimits {
149 pub withdraw_limit: u64,
150 pub borrow_limit: u64,
151}
152
153pub 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 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 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 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
247fn 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
265fn 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 #[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, 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, price_usdc_base_units: 1_000_000,
362 twap5min: 1_000_000,
363 spot_market: &market,
364 },
365 PositionData {
366 token_balance: -500_000, 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); }
380
381 #[test]
382 fn multi_market_positions() {
383 let usdc = usdc_market();
384 let sol = sol_market(); let positions = [
386 PositionData {
387 token_balance: 10_000_000, 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, price_usdc_base_units: 100_000_000, twap5min: 100_000_000,
396 spot_market: &sol,
397 },
398 ];
399 let state = MarginState::calculate(&positions).unwrap();
400 assert_eq!(state.total_weighted_collateral, 10_000_000 + 80_000_000);
403 assert_eq!(state.total_weighted_liabilities, 0);
404 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 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 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 #[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 #[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 #[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 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(); 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 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(); let state = MarginState {
542 total_weighted_collateral: 100_000_000, total_weighted_liabilities: 20_000_000, total_collateral: 125_000_000,
545 total_liabilities: 20_000_000,
546 };
547 let limits = calculate_position_limits(
551 &state,
552 &market,
553 100_000_000, 2_000_000_000, false,
556 )
557 .unwrap();
558 assert_eq!(limits.withdraw_limit, 1_000_000_000); }
560
561 #[test]
564 fn borrow_limit_basic() {
565 let market = usdc_market(); 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 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 let usdc = usdc_market();
585 let state = MarginState {
586 total_weighted_collateral: 80_000_000, total_weighted_liabilities: 0,
588 total_collateral: 100_000_000,
589 total_liabilities: 0,
590 };
591 let limits = calculate_position_limits(&state, &usdc, 1_000_000, 0, false).unwrap();
593 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 assert_eq!(limits.withdraw_limit, 5_000_000);
613 assert_eq!(limits.borrow_limit, 15_000_000);
617 }
618
619 #[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 #[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 #[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}