1use alloy_primitives::{U160, U256};
17
18use super::full_math::FullMath;
19use crate::defi::tick_map::tick_math::get_sqrt_ratio_at_tick;
20
21pub fn encode_sqrt_ratio_x96(amount0: u128, amount1: u128) -> U160 {
33 let amount0_u256 = U256::from(amount0);
34 let amount1_u256 = U256::from(amount1);
35
36 if amount1_u256.is_zero() {
37 panic!("Division by zero");
38 }
39 if amount0_u256.is_zero() {
40 return U160::ZERO;
41 }
42
43 let q192 = U256::from(1u128) << 192;
49
50 if amount0_u256 > U256::MAX / q192 {
52 let sqrt_amount0 = FullMath::sqrt(amount0_u256);
55 let sqrt_amount1 = FullMath::sqrt(amount1_u256);
56
57 if sqrt_amount1.is_zero() {
58 panic!("Division by zero in sqrt");
59 }
60
61 let q96 = U256::from(1u128) << 96;
62
63 let result = FullMath::mul_div(sqrt_amount0, q96, sqrt_amount1).expect("mul_div overflow");
65
66 return if result > U256::from(U160::MAX) {
68 U160::MAX
69 } else {
70 U160::from(result)
71 };
72 }
73
74 let ratio_q192 = FullMath::mul_div(amount0_u256, q192, amount1_u256).expect("mul_div overflow");
76
77 let sqrt_result = FullMath::sqrt(ratio_q192);
79
80 if sqrt_result > U256::from(U160::MAX) {
82 U160::MAX
83 } else {
84 U160::from(sqrt_result)
85 }
86}
87
88fn get_next_sqrt_price_from_amount0_rounding_up(
90 sqrt_price_x96: U160,
91 liquidity: u128,
92 amount: U256,
93 add: bool,
94) -> U160 {
95 if amount.is_zero() {
96 return sqrt_price_x96;
97 }
98 let numerator = U256::from(liquidity) << 96;
99 let sqrt_price_x96 = U256::from(sqrt_price_x96);
100 let product = amount * sqrt_price_x96;
101
102 if add {
103 if product / amount == sqrt_price_x96 {
104 let denominator = numerator + product;
105 if denominator >= numerator {
106 let result = FullMath::mul_div_rounding_up(numerator, sqrt_price_x96, denominator)
108 .expect("mul_div_rounding_up failed");
109 return U160::from(result);
110 }
111 }
112
113 let fallback_denominator = (numerator / sqrt_price_x96) + amount;
115 let result = FullMath::div_rounding_up(numerator, fallback_denominator)
116 .expect("div_rounding_up failed");
117
118 if result > U256::from(U160::MAX) {
120 panic!("Result overflows U160");
121 }
122 U160::from(result)
123 } else {
124 if !((product / amount) == sqrt_price_x96 && numerator > product) {
126 panic!("Invalid conditions for amount0 removal: overflow or underflow detected")
127 }
128
129 let denominator = numerator - product;
130 let result = FullMath::mul_div_rounding_up(numerator, sqrt_price_x96, denominator)
131 .expect("mul_div_rounding_up failed");
132 U160::from(result)
133 }
134}
135
136fn get_next_sqrt_price_from_amount1_rounding_down(
138 sqrt_price_x96: U160,
139 liquidity: u128,
140 amount: U256,
141 add: bool,
142) -> U160 {
143 if add {
146 let quotient = if amount <= U256::from(U160::MAX) {
147 (amount << 96) / U256::from(liquidity)
149 } else {
150 FullMath::mul_div(amount, U256::from(1u128) << 96, U256::from(liquidity))
152 .unwrap_or(U256::ZERO)
153 };
154
155 U160::from(U256::from(sqrt_price_x96) + quotient)
157 } else {
158 let quotient = if amount <= U256::from(U160::MAX) {
159 FullMath::div_rounding_up(amount << 96, U256::from(liquidity)).unwrap_or(U256::ZERO)
161 } else {
162 FullMath::mul_div_rounding_up(amount, U256::from(1u128) << 96, U256::from(liquidity))
164 .unwrap_or(U256::ZERO)
165 };
166
167 if U256::from(sqrt_price_x96) <= quotient {
169 panic!("sqrt_price_x96 must be greater than quotient");
170 }
171
172 U160::from(U256::from(sqrt_price_x96) - quotient)
174 }
175}
176
177pub fn get_next_sqrt_price_from_input(
182 sqrt_price_x96: U160,
183 liquidity: u128,
184 amount_in: U256,
185 zero_for_one: bool,
186) -> U160 {
187 assert!(
188 sqrt_price_x96 > U160::ZERO,
189 "sqrt_price_x96 must be greater than zero"
190 );
191 assert!(liquidity > 0, "Liquidity must be greater than zero");
192
193 if zero_for_one {
194 get_next_sqrt_price_from_amount0_rounding_up(sqrt_price_x96, liquidity, amount_in, true)
195 } else {
196 get_next_sqrt_price_from_amount1_rounding_down(sqrt_price_x96, liquidity, amount_in, true)
197 }
198}
199
200pub fn get_next_sqrt_price_from_output(
205 sqrt_price_x96: U160,
206 liquidity: u128,
207 amount_out: U256,
208 zero_for_one: bool,
209) -> U160 {
210 assert!(
211 sqrt_price_x96 > U160::ZERO,
212 "sqrt_price_x96 must be greater than zero"
213 );
214 assert!(liquidity > 0, "Liquidity must be greater than zero");
215
216 if zero_for_one {
217 get_next_sqrt_price_from_amount1_rounding_down(sqrt_price_x96, liquidity, amount_out, false)
218 } else {
219 get_next_sqrt_price_from_amount0_rounding_up(sqrt_price_x96, liquidity, amount_out, false)
220 }
221}
222
223#[must_use]
225pub fn get_amount0_delta(
226 sqrt_ratio_ax96: U160,
227 sqrt_ratio_bx96: U160,
228 liquidity: u128,
229 round_up: bool,
230) -> U256 {
231 let (sqrt_ratio_a, sqrt_ratio_b) = if sqrt_ratio_ax96 > sqrt_ratio_bx96 {
232 (sqrt_ratio_bx96, sqrt_ratio_ax96)
233 } else {
234 (sqrt_ratio_ax96, sqrt_ratio_bx96)
235 };
236
237 let numerator1 = U256::from(liquidity) << 96;
238 let numerator2 = U256::from(sqrt_ratio_b - sqrt_ratio_a);
239
240 if round_up {
241 let result =
243 FullMath::mul_div_rounding_up(numerator1, numerator2, U256::from(sqrt_ratio_b))
244 .unwrap_or(U256::ZERO);
245
246 FullMath::div_rounding_up(result, U256::from(sqrt_ratio_a)).unwrap_or(U256::ZERO)
248 } else {
249 let result = FullMath::mul_div(numerator1, numerator2, U256::from(sqrt_ratio_b))
250 .unwrap_or(U256::ZERO);
251 result / U256::from(sqrt_ratio_a)
252 }
253}
254#[must_use]
256pub fn get_amount1_delta(
257 sqrt_ratio_ax96: U160,
258 sqrt_ratio_bx96: U160,
259 liquidity: u128,
260 round_up: bool,
261) -> U256 {
262 let (sqrt_ratio_a, sqrt_ratio_b) = if sqrt_ratio_ax96 > sqrt_ratio_bx96 {
263 (sqrt_ratio_bx96, sqrt_ratio_ax96)
264 } else {
265 (sqrt_ratio_ax96, sqrt_ratio_bx96)
266 };
267
268 let liquidity_u256 = U256::from(liquidity);
269 let sqrt_ratio_diff = U256::from(sqrt_ratio_b - sqrt_ratio_a);
270 let q96 = U256::from(1u128) << 96;
271
272 if round_up {
273 FullMath::mul_div_rounding_up(liquidity_u256, sqrt_ratio_diff, q96).unwrap_or(U256::ZERO)
274 } else {
275 FullMath::mul_div(liquidity_u256, sqrt_ratio_diff, q96).unwrap_or(U256::ZERO)
276 }
277}
278
279#[must_use]
281pub fn get_amounts_for_liquidity(
282 sqrt_ratio_x96: U160,
283 tick_lower: i32,
284 tick_upper: i32,
285 liquidity: u128,
286 round_up: bool,
287) -> (U256, U256) {
288 let sqrt_ratio_lower_x96 = get_sqrt_ratio_at_tick(tick_lower);
289 let sqrt_ratio_upper_x96 = get_sqrt_ratio_at_tick(tick_upper);
290
291 let (sqrt_ratio_a, sqrt_ratio_b) = if sqrt_ratio_lower_x96 > sqrt_ratio_upper_x96 {
293 (sqrt_ratio_upper_x96, sqrt_ratio_lower_x96)
294 } else {
295 (sqrt_ratio_lower_x96, sqrt_ratio_upper_x96)
296 };
297
298 let amount0 = if sqrt_ratio_x96 <= sqrt_ratio_a {
299 get_amount0_delta(sqrt_ratio_a, sqrt_ratio_b, liquidity, round_up)
301 } else if sqrt_ratio_x96 < sqrt_ratio_b {
302 get_amount0_delta(sqrt_ratio_x96, sqrt_ratio_b, liquidity, round_up)
304 } else {
305 U256::ZERO
307 };
308
309 let amount1 = if sqrt_ratio_x96 < sqrt_ratio_a {
310 U256::ZERO
312 } else if sqrt_ratio_x96 < sqrt_ratio_b {
313 get_amount1_delta(sqrt_ratio_a, sqrt_ratio_x96, liquidity, round_up)
315 } else {
316 get_amount1_delta(sqrt_ratio_a, sqrt_ratio_b, liquidity, round_up)
318 };
319
320 (amount0, amount1)
321}
322
323pub fn expand_to_18_decimals(amount: u64) -> u128 {
325 amount as u128 * 10u128.pow(18)
326}
327
328#[cfg(test)]
333mod tests {
334 use rstest::*;
336
337 use super::*;
338 use crate::defi::tick_map::full_math::Q96_U160;
339
340 #[rstest]
341 #[should_panic(expected = "sqrt_price_x96 must be greater than zero")]
342 fn test_if_get_next_sqrt_price_from_input_panic_if_price_zero() {
343 let _ = get_next_sqrt_price_from_input(U160::ZERO, 1, U256::ZERO, true);
344 }
345
346 #[rstest]
347 #[should_panic(expected = "Liquidity must be greater than zero")]
348 fn test_if_get_next_sqrt_price_from_input_panic_if_liquidity_zero() {
349 let _ = get_next_sqrt_price_from_input(U160::from(1), 0, U256::ZERO, true);
350 }
351
352 #[rstest]
353 #[should_panic(expected = "Uint conversion error: Value is too large for Uint<160>")]
354 fn test_if_get_next_sqrt_price_from_input_panics_from_big_price() {
355 let price = U160::MAX - U160::from(1);
356 let _ = get_next_sqrt_price_from_input(price, 1024, U256::from(1024), false);
357 }
358
359 #[rstest]
360 fn test_any_input_amount_cannot_underflow_the_price() {
361 let price = U160::from(1);
364 let liquidity = 1;
365 let amount_in = U256::from(2).pow(U256::from(255));
366 let result = get_next_sqrt_price_from_input(price, liquidity, amount_in, true);
367 assert_eq!(result, U160::from(1));
368 }
369
370 #[rstest]
371 fn test_returns_input_price_if_amount_in_is_zero_and_zero_for_one_true() {
372 let price = encode_sqrt_ratio_x96(1, 1);
373 let liquidity = expand_to_18_decimals(1) / 10;
374 let result = get_next_sqrt_price_from_input(price, liquidity, U256::ZERO, true);
375 assert_eq!(result, price);
376 }
377
378 #[rstest]
379 fn test_returns_input_price_if_amount_in_is_zero_and_zero_for_one_false() {
380 let price = encode_sqrt_ratio_x96(1, 1);
381 let liquidity = expand_to_18_decimals(1) / 10;
382 let result = get_next_sqrt_price_from_input(price, liquidity, U256::ZERO, false);
383 assert_eq!(result, price);
384 }
385
386 #[rstest]
387 fn test_returns_the_minimum_price_for_max_inputs() {
388 let sqrt_p = U160::MAX;
389 let liquidity = u128::MAX;
390 let max_amount_no_overflow = U256::MAX - (U256::from(liquidity) << 96) / U256::from(sqrt_p);
391 let result =
392 get_next_sqrt_price_from_input(sqrt_p, liquidity, max_amount_no_overflow, true);
393 assert_eq!(result, U160::from(1));
394 }
395
396 #[rstest]
397 fn test_input_amount_of_0_1_token1() {
398 let sqrt_q = get_next_sqrt_price_from_input(
399 encode_sqrt_ratio_x96(1, 1),
400 expand_to_18_decimals(1),
401 U256::from(expand_to_18_decimals(1)) / U256::from(10),
402 false,
403 );
404 assert_eq!(
405 sqrt_q,
406 U160::from_str_radix("87150978765690771352898345369", 10).unwrap()
407 );
408 }
409
410 #[rstest]
411 fn test_input_amount_of_0_1_token0() {
412 let sqrt_q = get_next_sqrt_price_from_input(
413 encode_sqrt_ratio_x96(1, 1),
414 expand_to_18_decimals(1),
415 U256::from(expand_to_18_decimals(1)) / U256::from(10),
416 true,
417 );
418 assert_eq!(
419 sqrt_q,
420 U160::from_str_radix("72025602285694852357767227579", 10).unwrap()
421 );
422 }
423
424 #[rstest]
425 fn test_amount_in_greater_than_uint96_max_and_zero_for_one_true() {
426 let result = get_next_sqrt_price_from_input(
427 encode_sqrt_ratio_x96(1, 1),
428 expand_to_18_decimals(10),
429 U256::from(2).pow(U256::from(100)),
430 true,
431 );
432 assert_eq!(
433 result,
434 U160::from_str_radix("624999999995069620", 10).unwrap()
435 );
436 }
437
438 #[rstest]
439 fn test_can_return_1_with_enough_amount_in_and_zero_for_one_true() {
440 let result = get_next_sqrt_price_from_input(
441 encode_sqrt_ratio_x96(1, 1),
442 1,
443 U256::MAX / U256::from(2),
444 true,
445 );
446 assert_eq!(result, U160::from(1));
447 }
448
449 #[rstest]
450 #[should_panic(
451 expected = "Invalid conditions for amount0 removal: overflow or underflow detected"
452 )]
453 fn test_fails_if_output_amount_is_exactly_virtual_reserves_of_token0() {
454 let price = U160::from_str_radix("20282409603651670423947251286016", 10).unwrap();
455 let liquidity = 1024;
456 let amount_out = U256::from(4);
457 let _ = get_next_sqrt_price_from_output(price, liquidity, amount_out, false);
458 }
459
460 #[rstest]
461 #[should_panic(
462 expected = "Invalid conditions for amount0 removal: overflow or underflow detected"
463 )]
464 fn test_fails_if_output_amount_is_greater_than_virtual_reserves_of_token0() {
465 let price = U160::from_str_radix("20282409603651670423947251286016", 10).unwrap();
466 let liquidity = 1024;
467 let amount_out = U256::from(5);
468 let _ = get_next_sqrt_price_from_output(price, liquidity, amount_out, false);
469 }
470
471 #[rstest]
472 #[should_panic(expected = "sqrt_price_x96 must be greater than quotient")]
473 fn test_fails_if_output_amount_is_greater_than_virtual_reserves_of_token1() {
474 let price = U160::from_str_radix("20282409603651670423947251286016", 10).unwrap();
475 let liquidity = 1024;
476 let amount_out = U256::from(262145);
477 let _ = get_next_sqrt_price_from_output(price, liquidity, amount_out, true);
478 }
479
480 #[rstest]
481 #[should_panic(expected = "sqrt_price_x96 must be greater than quotient")]
482 fn test_fails_if_output_amount_is_exactly_virtual_reserves_of_token1() {
483 let price = U160::from_str_radix("20282409603651670423947251286016", 10).unwrap();
484 let liquidity = 1024;
485 let amount_out = U256::from(262144);
486 let _ = get_next_sqrt_price_from_output(price, liquidity, amount_out, true);
487 }
488
489 #[rstest]
490 fn test_succeeds_if_output_amount_is_just_less_than_virtual_reserves_of_token1() {
491 let price = U160::from_str_radix("20282409603651670423947251286016", 10).unwrap();
492 let liquidity = 1024;
493 let amount_out = U256::from(262143);
494 let result = get_next_sqrt_price_from_output(price, liquidity, amount_out, true);
495 assert_eq!(
496 result,
497 U160::from_str_radix("77371252455336267181195264", 10).unwrap()
498 );
499 }
500
501 #[rstest]
502 fn test_returns_input_price_if_amount_out_is_zero_and_zero_for_one_true() {
503 let price = encode_sqrt_ratio_x96(1, 1);
504 let liquidity = expand_to_18_decimals(1) / 10;
505 let result = get_next_sqrt_price_from_output(price, liquidity, U256::ZERO, true);
506 assert_eq!(result, price);
507 }
508
509 #[rstest]
510 fn test_returns_input_price_if_amount_out_is_zero_and_zero_for_one_false() {
511 let price = encode_sqrt_ratio_x96(1, 1);
512 let liquidity = expand_to_18_decimals(1) / 10;
513 let result = get_next_sqrt_price_from_output(price, liquidity, U256::ZERO, false);
514 assert_eq!(result, price);
515 }
516
517 #[rstest]
518 fn test_output_amount_of_0_1_token1_zero_for_one_false() {
519 let sqrt_q = get_next_sqrt_price_from_output(
520 encode_sqrt_ratio_x96(1, 1),
521 expand_to_18_decimals(1),
522 U256::from(expand_to_18_decimals(1)) / U256::from(10),
523 false,
524 );
525 assert_eq!(
526 sqrt_q,
527 U160::from_str_radix("88031291682515930659493278152", 10).unwrap()
528 );
529 }
530
531 #[rstest]
532 fn test_output_amount_of_0_1_token1_zero_for_one_true() {
533 let sqrt_q = get_next_sqrt_price_from_output(
534 encode_sqrt_ratio_x96(1, 1),
535 expand_to_18_decimals(1),
536 U256::from(expand_to_18_decimals(1)) / U256::from(10),
537 true,
538 );
539 assert_eq!(
540 sqrt_q,
541 U160::from_str_radix("71305346262837903834189555302", 10).unwrap()
542 );
543 }
544
545 #[rstest]
546 #[should_panic(expected = "sqrt_price_x96 must be greater than zero")]
547 fn test_if_get_next_sqrt_price_from_output_panic_if_price_zero() {
548 let _ = get_next_sqrt_price_from_output(U160::ZERO, 1, U256::ZERO, true);
549 }
550
551 #[rstest]
552 #[should_panic(expected = "Liquidity must be greater than zero")]
553 fn test_if_get_next_sqrt_price_from_output_panic_if_liquidity_zero() {
554 let _ = get_next_sqrt_price_from_output(U160::from(1), 0, U256::ZERO, true);
555 }
556
557 #[rstest]
558 fn test_encode_sqrt_ratio_x98_some_values() {
559 assert_eq!(encode_sqrt_ratio_x96(1, 1), Q96_U160);
560 assert_eq!(
561 encode_sqrt_ratio_x96(100, 1),
562 U160::from(792281625142643375935439503360_u128)
563 );
564 assert_eq!(
565 encode_sqrt_ratio_x96(1, 100),
566 U160::from(7922816251426433759354395033_u128)
567 );
568 assert_eq!(
569 encode_sqrt_ratio_x96(111, 333),
570 U160::from(45742400955009932534161870629_u128)
571 );
572 assert_eq!(
573 encode_sqrt_ratio_x96(333, 111),
574 U160::from(137227202865029797602485611888_u128)
575 );
576 }
577
578 #[rstest]
579 fn test_get_amount0_delta_returns_0_if_liquidity_is_0() {
580 let amount0 = get_amount0_delta(
581 encode_sqrt_ratio_x96(1, 1),
582 encode_sqrt_ratio_x96(2, 1),
583 0,
584 true,
585 );
586 assert_eq!(amount0, U256::ZERO);
587 }
588
589 #[rstest]
590 fn test_get_amount0_delta_returns_0_if_prices_are_equal() {
591 let amount0 = get_amount0_delta(
592 encode_sqrt_ratio_x96(1, 1),
593 encode_sqrt_ratio_x96(1, 1),
594 0,
595 true,
596 );
597 assert_eq!(amount0, U256::ZERO);
598 }
599
600 #[rstest]
601 fn test_get_amount0_delta_returns_0_1_amount1_for_price_of_1_to_1_21() {
602 let amount0 = get_amount0_delta(
603 encode_sqrt_ratio_x96(1, 1),
604 encode_sqrt_ratio_x96(121, 100),
605 expand_to_18_decimals(1),
606 true,
607 );
608 assert_eq!(
609 amount0,
610 U256::from_str_radix("90909090909090910", 10).unwrap()
611 );
612
613 let amount0_rounded_down = get_amount0_delta(
614 encode_sqrt_ratio_x96(1, 1),
615 encode_sqrt_ratio_x96(121, 100),
616 expand_to_18_decimals(1),
617 false,
618 );
619
620 assert_eq!(amount0_rounded_down, amount0 - U256::from(1));
621 }
622
623 #[rstest]
624 fn test_get_amount0_delta_works_for_prices_that_overflow() {
625 let price_low =
627 encode_sqrt_ratio_x96(U256::from(2).pow(U256::from(90)).try_into().unwrap(), 1);
628 let price_high =
629 encode_sqrt_ratio_x96(U256::from(2).pow(U256::from(96)).try_into().unwrap(), 1);
630
631 let amount0_up = get_amount0_delta(price_low, price_high, expand_to_18_decimals(1), true);
632
633 let amount0_down =
634 get_amount0_delta(price_low, price_high, expand_to_18_decimals(1), false);
635
636 assert_eq!(amount0_up, amount0_down + U256::from(1));
637 }
638
639 #[rstest]
640 fn test_get_amount1_delta_returns_0_if_liquidity_is_0() {
641 let amount1 = get_amount1_delta(
642 encode_sqrt_ratio_x96(1, 1),
643 encode_sqrt_ratio_x96(2, 1),
644 0,
645 true,
646 );
647 assert_eq!(amount1, U256::ZERO);
648 }
649
650 #[rstest]
651 fn test_get_amount1_delta_returns_0_if_prices_are_equal() {
652 let amount1 = get_amount1_delta(
653 encode_sqrt_ratio_x96(1, 1),
654 encode_sqrt_ratio_x96(1, 1),
655 0,
656 true,
657 );
658 assert_eq!(amount1, U256::ZERO);
659 }
660
661 #[rstest]
662 fn test_get_amount1_delta_returns_0_1_amount1_for_price_of_1_to_1_21() {
663 let amount1 = get_amount1_delta(
664 encode_sqrt_ratio_x96(1, 1),
665 encode_sqrt_ratio_x96(121, 100),
666 expand_to_18_decimals(1),
667 true,
668 );
669 assert_eq!(
670 amount1,
671 U256::from_str_radix("100000000000000000", 10).unwrap()
672 );
673
674 let amount1_rounded_down = get_amount1_delta(
675 encode_sqrt_ratio_x96(1, 1),
676 encode_sqrt_ratio_x96(121, 100),
677 expand_to_18_decimals(1),
678 false,
679 );
680
681 assert_eq!(amount1_rounded_down, amount1 - U256::from(1));
682 }
683}