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
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::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 #[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, 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, price_usdc_base_units: 1_000_000,
365 twap5min: 1_000_000,
366 spot_market: &market,
367 },
368 PositionData {
369 token_balance: -500_000, 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); }
383
384 #[test]
385 fn multi_market_positions() {
386 let usdc = usdc_market();
387 let sol = sol_market(); let positions = [
389 PositionData {
390 token_balance: 10_000_000, 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, price_usdc_base_units: 100_000_000, twap5min: 100_000_000,
399 spot_market: &sol,
400 },
401 ];
402 let state = MarginState::calculate(&positions).unwrap();
403 assert_eq!(state.total_weighted_collateral, 10_000_000 + 80_000_000);
406 assert_eq!(state.total_weighted_liabilities, 0);
407 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 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 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 #[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 #[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 #[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 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(); 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 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(); let state = MarginState {
545 total_weighted_collateral: 100_000_000, total_weighted_liabilities: 20_000_000, total_collateral: 125_000_000,
548 total_liabilities: 20_000_000,
549 };
550 let limits = calculate_position_limits(
554 &state,
555 &market,
556 100_000_000, 2_000_000_000, false,
559 )
560 .unwrap();
561 assert_eq!(limits.withdraw_limit, 1_000_000_000); }
563
564 #[test]
567 fn borrow_limit_basic() {
568 let market = usdc_market(); 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 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 let usdc = usdc_market();
588 let state = MarginState {
589 total_weighted_collateral: 80_000_000, total_weighted_liabilities: 0,
591 total_collateral: 100_000_000,
592 total_liabilities: 0,
593 };
594 let limits = calculate_position_limits(&state, &usdc, 1_000_000, 0, false).unwrap();
596 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 assert_eq!(limits.withdraw_limit, 5_000_000);
616 assert_eq!(limits.borrow_limit, 15_000_000);
620 }
621
622 #[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 #[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 #[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}