kabu_defi_uniswap_v3_math/
sqrt_price_math.rs

1use alloy::primitives::{I256, U256};
2
3use crate::{
4    error::UniswapV3MathError,
5    full_math::{mul_div, mul_div_rounding_up},
6    unsafe_math::div_rounding_up,
7};
8
9pub const MAX_U160: U256 = U256::from_limbs([18446744073709551615, 18446744073709551615, 4294967295, 0]);
10pub const Q96: U256 = U256::from_limbs([0, 4294967296, 0, 0]);
11pub const FIXED_POINT_96_RESOLUTION: U256 = U256::from_limbs([96, 0, 0, 0]);
12
13// returns (sqrtQX96)
14pub fn get_next_sqrt_price_from_input(
15    sqrt_price: U256,
16    liquidity: u128,
17    amount_in: U256,
18    zero_for_one: bool,
19) -> Result<U256, UniswapV3MathError> {
20    if sqrt_price.is_zero() {
21        return Err(UniswapV3MathError::SqrtPriceIsZero);
22    } else if liquidity == 0 {
23        return Err(UniswapV3MathError::LiquidityIsZero);
24    }
25
26    if zero_for_one {
27        get_next_sqrt_price_from_amount_0_rounding_up(sqrt_price, liquidity, amount_in, true)
28    } else {
29        get_next_sqrt_price_from_amount_1_rounding_down(sqrt_price, liquidity, amount_in, true)
30    }
31}
32
33// returns (sqrtQX96)
34pub fn get_next_sqrt_price_from_output(
35    sqrt_price: U256,
36    liquidity: u128,
37    amount_out: U256,
38    zero_for_one: bool,
39) -> Result<U256, UniswapV3MathError> {
40    if sqrt_price.is_zero() {
41        return Err(UniswapV3MathError::SqrtPriceIsZero);
42    } else if liquidity == 0 {
43        return Err(UniswapV3MathError::LiquidityIsZero);
44    }
45
46    if zero_for_one {
47        get_next_sqrt_price_from_amount_1_rounding_down(sqrt_price, liquidity, amount_out, false)
48    } else {
49        get_next_sqrt_price_from_amount_0_rounding_up(sqrt_price, liquidity, amount_out, false)
50    }
51}
52
53// returns (uint160 sqrtQX96)
54pub fn get_next_sqrt_price_from_amount_0_rounding_up(
55    sqrt_price_x_96: U256,
56    liquidity: u128,
57    amount: U256,
58    add: bool,
59) -> Result<U256, UniswapV3MathError> {
60    if amount.is_zero() {
61        return Ok(sqrt_price_x_96);
62    }
63
64    let numerator_1: U256 = U256::from(liquidity) << 96;
65
66    if add {
67        let product = amount.wrapping_mul(sqrt_price_x_96);
68
69        if product.wrapping_div(amount) == sqrt_price_x_96 {
70            let denominator = numerator_1.wrapping_add(product);
71
72            if denominator >= numerator_1 {
73                return mul_div_rounding_up(numerator_1, sqrt_price_x_96, denominator);
74            }
75        }
76
77        Ok(div_rounding_up(numerator_1, (numerator_1.wrapping_div(sqrt_price_x_96)).wrapping_add(amount)))
78    } else {
79        let product = amount.wrapping_mul(sqrt_price_x_96);
80        if product.wrapping_div(amount) == sqrt_price_x_96 && numerator_1 > product {
81            let denominator = numerator_1.wrapping_sub(product);
82
83            mul_div_rounding_up(numerator_1, sqrt_price_x_96, denominator)
84        } else {
85            Err(UniswapV3MathError::ProductDivAmount)
86        }
87    }
88}
89
90// returns (uint160 sqrtQX96)
91pub fn get_next_sqrt_price_from_amount_1_rounding_down(
92    sqrt_price_x_96: U256,
93    liquidity: u128,
94    amount: U256,
95    add: bool,
96) -> Result<U256, UniswapV3MathError> {
97    let liquidity = U256::from(liquidity);
98
99    if add {
100        let quotient =
101            if amount <= MAX_U160 { (amount << FIXED_POINT_96_RESOLUTION) / liquidity } else { mul_div(amount, Q96, liquidity)? };
102
103        let next_sqrt_price = sqrt_price_x_96 + quotient;
104
105        if next_sqrt_price > MAX_U160 { Err(UniswapV3MathError::SafeCastToU160Overflow) } else { Ok(next_sqrt_price) }
106    } else {
107        let quotient = if amount <= MAX_U160 {
108            div_rounding_up(amount << FIXED_POINT_96_RESOLUTION, liquidity)
109        } else {
110            mul_div_rounding_up(amount, Q96, liquidity)?
111        };
112
113        //require(sqrtPX96 > quotient);
114        if sqrt_price_x_96 <= quotient {
115            return Err(UniswapV3MathError::SqrtPriceIsLteQuotient);
116        }
117
118        Ok(sqrt_price_x_96.overflowing_sub(quotient).0)
119    }
120}
121
122// returns (uint256 amount0)
123pub fn _get_amount_0_delta(
124    mut sqrt_ratio_a_x_96: U256,
125    mut sqrt_ratio_b_x_96: U256,
126    liquidity: u128,
127    round_up: bool,
128) -> Result<U256, UniswapV3MathError> {
129    if sqrt_ratio_a_x_96 > sqrt_ratio_b_x_96 {
130        (sqrt_ratio_a_x_96, sqrt_ratio_b_x_96) = (sqrt_ratio_b_x_96, sqrt_ratio_a_x_96)
131    };
132
133    let numerator_1 = U256::from(liquidity) << 96;
134    let numerator_2 = sqrt_ratio_b_x_96 - sqrt_ratio_a_x_96;
135
136    if sqrt_ratio_a_x_96.is_zero() {
137        return Err(UniswapV3MathError::SqrtPriceIsZero);
138    }
139
140    if round_up {
141        let numerator_partial = mul_div_rounding_up(numerator_1, numerator_2, sqrt_ratio_b_x_96)?;
142        Ok(div_rounding_up(numerator_partial, sqrt_ratio_a_x_96))
143    } else {
144        Ok(mul_div(numerator_1, numerator_2, sqrt_ratio_b_x_96)? / sqrt_ratio_a_x_96)
145    }
146}
147
148// returns (uint256 amount1)
149pub fn _get_amount_1_delta(
150    mut sqrt_ratio_a_x_96: U256,
151    mut sqrt_ratio_b_x_96: U256,
152    liquidity: u128,
153    round_up: bool,
154) -> Result<U256, UniswapV3MathError> {
155    if sqrt_ratio_a_x_96 > sqrt_ratio_b_x_96 {
156        (sqrt_ratio_a_x_96, sqrt_ratio_b_x_96) = (sqrt_ratio_b_x_96, sqrt_ratio_a_x_96)
157    };
158
159    let denominator = U256::from_limbs([0, 4294967296, 0, 0]);
160
161    if round_up {
162        mul_div_rounding_up(U256::from(liquidity), sqrt_ratio_b_x_96 - sqrt_ratio_a_x_96, denominator)
163    } else {
164        mul_div(U256::from(liquidity), sqrt_ratio_b_x_96 - sqrt_ratio_a_x_96, denominator)
165    }
166}
167
168pub fn get_amount_0_delta(sqrt_ratio_a_x_96: U256, sqrt_ratio_b_x_96: U256, liquidity: i128) -> Result<I256, UniswapV3MathError> {
169    if liquidity < 0 {
170        Ok(-I256::from_raw(_get_amount_0_delta(sqrt_ratio_a_x_96, sqrt_ratio_b_x_96, -liquidity as u128, false)?))
171    } else {
172        Ok(I256::from_raw(_get_amount_0_delta(sqrt_ratio_a_x_96, sqrt_ratio_b_x_96, liquidity as u128, true)?))
173    }
174}
175
176pub fn get_amount_1_delta(sqrt_ratio_a_x_96: U256, sqrt_ratio_b_x_96: U256, liquidity: i128) -> Result<I256, UniswapV3MathError> {
177    if liquidity < 0 {
178        Ok(-I256::from_raw(_get_amount_1_delta(sqrt_ratio_a_x_96, sqrt_ratio_b_x_96, -liquidity as u128, false)?))
179    } else {
180        Ok(I256::from_raw(_get_amount_1_delta(sqrt_ratio_a_x_96, sqrt_ratio_b_x_96, liquidity as u128, true)?))
181    }
182}
183
184#[cfg(test)]
185mod test {
186    use std::{
187        ops::{Add, Sub},
188        str::FromStr,
189    };
190
191    use alloy::primitives::U256;
192
193    use crate::{
194        U256_1, U256_2,
195        sqrt_price_math::{_get_amount_1_delta, MAX_U160, get_next_sqrt_price_from_output},
196    };
197
198    use super::{_get_amount_0_delta, get_next_sqrt_price_from_input};
199
200    #[test]
201    fn test_get_next_sqrt_price_from_input() {
202        //Fails if price is zero
203        let result = get_next_sqrt_price_from_input(U256::ZERO, 0, U256::from(100000000000000000_u128), false);
204        assert_eq!(result.unwrap_err().to_string(), "Sqrt price is 0");
205
206        //Fails if liquidity is zero
207        let result = get_next_sqrt_price_from_input(U256_1, 0, U256::from(100000000000000000_u128), true);
208        assert_eq!(result.unwrap_err().to_string(), "Liquidity is 0");
209
210        //fails if input amount overflows the price
211        let result = get_next_sqrt_price_from_input(MAX_U160, 1024, U256::from(1024), false);
212        assert_eq!(result.unwrap_err().to_string(), "Overflow when casting to U160");
213
214        //any input amount cannot underflow the price
215        let result = get_next_sqrt_price_from_input(
216            U256_1,
217            1,
218            U256::from_str("57896044618658097711785492504343953926634992332820282019728792003956564819968").unwrap(),
219            true,
220        );
221
222        assert_eq!(result.unwrap(), U256_1);
223
224        //returns input price if amount in is zero and zeroForOne = true
225        let result =
226            get_next_sqrt_price_from_input(U256::from_str("79228162514264337593543950336").unwrap(), 1e17 as u128, U256::ZERO, true);
227
228        assert_eq!(result.unwrap(), U256::from_str("79228162514264337593543950336").unwrap());
229
230        //returns input price if amount in is zero and zeroForOne = false
231        let result =
232            get_next_sqrt_price_from_input(U256::from_str("79228162514264337593543950336").unwrap(), 1e17 as u128, U256::ZERO, true);
233
234        assert_eq!(result.unwrap(), U256::from_str("79228162514264337593543950336").unwrap());
235
236        //returns the minimum price for max inputs
237
238        let sqrt_price = MAX_U160;
239        let liquidity = u128::MAX;
240        let max_amount_no_overflow = U256::MAX - ((U256::from(liquidity) << 96) / sqrt_price);
241        let result = get_next_sqrt_price_from_input(sqrt_price, liquidity, max_amount_no_overflow, true);
242        assert_eq!(result.unwrap(), U256_1);
243
244        //input amount of 0.1 token1
245        let result = get_next_sqrt_price_from_input(
246            U256::from_str("79228162514264337593543950336").unwrap(),
247            1e18 as u128,
248            U256::from_str("100000000000000000").unwrap(),
249            false,
250        );
251
252        assert_eq!(result.unwrap(), U256::from_str("87150978765690771352898345369").unwrap());
253
254        //input amount of 0.1 token0
255        let result = get_next_sqrt_price_from_input(
256            U256::from_str("79228162514264337593543950336").unwrap(),
257            1e18 as u128,
258            U256::from_str("100000000000000000").unwrap(),
259            true,
260        );
261
262        assert_eq!(result.unwrap(), U256::from_str("72025602285694852357767227579").unwrap());
263
264        //amountIn > type(uint96).max and zeroForOne = true
265        let result = get_next_sqrt_price_from_input(
266            U256::from_str("79228162514264337593543950336").unwrap(),
267            1e19 as u128,
268            U256::from_str("1267650600228229401496703205376").unwrap(),
269            true,
270        );
271        // perfect answer:
272        // https://www.wolframalpha.com/input/?i=624999999995069620+-+%28%281e19+*+1+%2F+%281e19+%2B+2%5E100+*+1%29%29+*+2%5E96%29
273        assert_eq!(result.unwrap(), U256::from_str("624999999995069620").unwrap());
274
275        //can return 1 with enough amountIn and zeroForOne = true
276        let result = get_next_sqrt_price_from_input(U256::from_str("79228162514264337593543950336").unwrap(), 1, U256::MAX / U256_2, true);
277
278        assert_eq!(result.unwrap(), U256_1);
279    }
280
281    #[test]
282    fn test_get_next_sqrt_price_from_output() {
283        //fails if price is zero
284        let result = get_next_sqrt_price_from_output(U256::ZERO, 0, U256::from(1000000000), false);
285        assert_eq!(result.unwrap_err().to_string(), "Sqrt price is 0");
286
287        //fails if liquidity is zero
288        let result = get_next_sqrt_price_from_output(U256_1, 0, U256::from(1000000000), false);
289        assert_eq!(result.unwrap_err().to_string(), "Liquidity is 0");
290
291        //fails if output amount is exactly the virtual reserves of token0
292        let result =
293            get_next_sqrt_price_from_output(U256::from_str("20282409603651670423947251286016").unwrap(), 1024, U256::from(4), false);
294        assert_eq!(result.unwrap_err().to_string(), "require((product = amount * sqrtPX96) / amount == sqrtPX96 && numerator1 > product);");
295
296        //fails if output amount is greater than virtual reserves of token0
297        let result =
298            get_next_sqrt_price_from_output(U256::from_str("20282409603651670423947251286016").unwrap(), 1024, U256::from(5), false);
299        assert_eq!(result.unwrap_err().to_string(), "require((product = amount * sqrtPX96) / amount == sqrtPX96 && numerator1 > product);");
300
301        //fails if output amount is greater than virtual reserves of token1
302        let result =
303            get_next_sqrt_price_from_output(U256::from_str("20282409603651670423947251286016").unwrap(), 1024, U256::from(262145), true);
304        assert_eq!(result.unwrap_err().to_string(), "Sqrt price is less than or equal to quotient");
305
306        //fails if output amount is exactly the virtual reserves of token1
307        let result =
308            get_next_sqrt_price_from_output(U256::from_str("20282409603651670423947251286016").unwrap(), 1024, U256::from(262144), true);
309        assert_eq!(result.unwrap_err().to_string(), "Sqrt price is less than or equal to quotient");
310
311        //succeeds if output amount is just less than the virtual
312        let result =
313            get_next_sqrt_price_from_output(U256::from_str("20282409603651670423947251286016").unwrap(), 1024, U256::from(262143), true);
314        assert_eq!(result.unwrap(), U256::from_str("77371252455336267181195264").unwrap());
315
316        //puzzling echidna test
317        let result =
318            get_next_sqrt_price_from_output(U256::from_str("20282409603651670423947251286016").unwrap(), 1024, U256::from(4), false);
319        assert_eq!(result.unwrap_err().to_string(), "require((product = amount * sqrtPX96) / amount == sqrtPX96 && numerator1 > product);");
320
321        //returns input price if amount in is zero and zeroForOne = true
322        let result =
323            get_next_sqrt_price_from_output(U256::from_str("79228162514264337593543950336").unwrap(), 1e17 as u128, U256::ZERO, true);
324        assert_eq!(result.unwrap(), U256::from_str("79228162514264337593543950336").unwrap());
325
326        //returns input price if amount in is zero and zeroForOne = false
327        let result =
328            get_next_sqrt_price_from_output(U256::from_str("79228162514264337593543950336").unwrap(), 1e17 as u128, U256::ZERO, false);
329        assert_eq!(result.unwrap(), U256::from_str("79228162514264337593543950336").unwrap());
330
331        //output amount of 0.1 token1
332        let result = get_next_sqrt_price_from_output(
333            U256::from_str("79228162514264337593543950336").unwrap(),
334            1e18 as u128,
335            U256::from(1e17 as u128),
336            false,
337        );
338        assert_eq!(result.unwrap(), U256::from_str("88031291682515930659493278152").unwrap());
339
340        //output amount of 0.1 token1
341        let result = get_next_sqrt_price_from_output(
342            U256::from_str("79228162514264337593543950336").unwrap(),
343            1e18 as u128,
344            U256::from(1e17 as u128),
345            true,
346        );
347        assert_eq!(result.unwrap(), U256::from_str("71305346262837903834189555302").unwrap());
348
349        //reverts if amountOut is impossible in zero for one direction
350        let result = get_next_sqrt_price_from_output(U256::from_str("79228162514264337593543950336").unwrap(), 1, U256::MAX, true);
351        assert_eq!(result.unwrap_err().to_string(), "Denominator is less than or equal to prod_1");
352
353        //reverts if amountOut is impossible in one for zero direction
354        let result = get_next_sqrt_price_from_output(U256::from_str("79228162514264337593543950336").unwrap(), 1, U256::MAX, false);
355        assert_eq!(result.unwrap_err().to_string(), "require((product = amount * sqrtPX96) / amount == sqrtPX96 && numerator1 > product);");
356    }
357
358    #[test]
359    fn test_get_amount_0_delta() {
360        // returns 0 if liquidity is 0
361        let amount_0 = _get_amount_0_delta(
362            U256::from_str("79228162514264337593543950336").unwrap(),
363            U256::from_str("79228162514264337593543950336").unwrap(),
364            0,
365            true,
366        );
367
368        assert_eq!(amount_0.unwrap(), U256::ZERO);
369
370        // returns 0 if prices are equal
371        let amount_0 = _get_amount_0_delta(
372            U256::from_str("79228162514264337593543950336").unwrap(),
373            U256::from_str("87150978765690771352898345369").unwrap(),
374            0,
375            true,
376        );
377
378        assert_eq!(amount_0.unwrap(), U256::ZERO);
379
380        // returns 0.1 amount1 for price of 1 to 1.21
381        let amount_0 = _get_amount_0_delta(
382            U256::from_str("79228162514264337593543950336").unwrap(),
383            U256::from_str("87150978765690771352898345369").unwrap(),
384            1e18 as u128,
385            true,
386        )
387        .unwrap();
388
389        assert_eq!(amount_0.clone(), U256::from_str("90909090909090910").unwrap());
390
391        let amount_0_rounded_down = _get_amount_0_delta(
392            U256::from_str("79228162514264337593543950336").unwrap(),
393            U256::from_str("87150978765690771352898345369").unwrap(),
394            1e18 as u128,
395            false,
396        );
397
398        assert_eq!(amount_0_rounded_down.unwrap(), amount_0.sub(U256_1));
399
400        // works for prices that overflow
401        let amount_0_up = _get_amount_0_delta(
402            U256::from_str("2787593149816327892691964784081045188247552").unwrap(),
403            U256::from_str("22300745198530623141535718272648361505980416").unwrap(),
404            1e18 as u128,
405            true,
406        )
407        .unwrap();
408
409        let amount_0_down = _get_amount_0_delta(
410            U256::from_str("2787593149816327892691964784081045188247552").unwrap(),
411            U256::from_str("22300745198530623141535718272648361505980416").unwrap(),
412            1e18 as u128,
413            false,
414        )
415        .unwrap();
416
417        assert_eq!(amount_0_up, amount_0_down.add(U256_1));
418    }
419
420    #[test]
421    fn test_get_amount_1_delta() {
422        // returns 0 if liquidity is 0
423        let amount_1 = _get_amount_1_delta(
424            U256::from_str("79228162514264337593543950336").unwrap(),
425            U256::from_str("79228162514264337593543950336").unwrap(),
426            0,
427            true,
428        );
429
430        assert_eq!(amount_1.unwrap(), U256::ZERO);
431
432        // returns 0 if prices are equal
433        let amount_1 = _get_amount_1_delta(
434            U256::from_str("79228162514264337593543950336").unwrap(),
435            U256::from_str("87150978765690771352898345369").unwrap(),
436            0,
437            true,
438        );
439
440        assert_eq!(amount_1.unwrap(), U256::ZERO);
441
442        // returns 0.1 amount1 for price of 1 to 1.21
443        let amount_1 = _get_amount_1_delta(
444            U256::from_str("79228162514264337593543950336").unwrap(),
445            U256::from_str("87150978765690771352898345369").unwrap(),
446            1e18 as u128,
447            true,
448        )
449        .unwrap();
450
451        assert_eq!(amount_1.clone(), U256::from_str("100000000000000000").unwrap());
452
453        let amount_1_rounded_down = _get_amount_1_delta(
454            U256::from_str("79228162514264337593543950336").unwrap(),
455            U256::from_str("87150978765690771352898345369").unwrap(),
456            1e18 as u128,
457            false,
458        );
459
460        assert_eq!(amount_1_rounded_down.unwrap(), amount_1.sub(U256_1));
461    }
462
463    #[test]
464    fn test_swap_computation() {
465        let sqrt_price = U256::from_str("1025574284609383690408304870162715216695788925244").unwrap();
466        let liquidity = 50015962439936049619261659728067971248;
467        let zero_for_one = true;
468        let amount_in = U256::from(406);
469
470        let sqrt_q = get_next_sqrt_price_from_input(sqrt_price, liquidity, amount_in, zero_for_one).unwrap();
471
472        assert_eq!(sqrt_q, U256::from_str("1025574284609383582644711336373707553698163132913").unwrap());
473
474        let amount_0_delta = _get_amount_0_delta(sqrt_q, sqrt_price, liquidity, true).unwrap();
475
476        assert_eq!(amount_0_delta, U256::from(406));
477    }
478}
kabu_defi_uniswap_v3_math/sqrt_price_math.rs

kabu_defi_uniswap_v3_math/
sqrt_price_math.rs
