kabu_defi_uniswap_v3_math/
swap_math.rs

1use alloy::primitives::{I256, U256};
2
3use crate::{
4    error::UniswapV3MathError,
5    full_math::{mul_div, mul_div_rounding_up},
6    sqrt_price_math::{_get_amount_0_delta, _get_amount_1_delta, get_next_sqrt_price_from_input, get_next_sqrt_price_from_output},
7};
8
9// //returns (
10//         uint160 sqrtRatioNextX96,
11//         uint256 amountIn,
12//         uint256 amountOut,
13//         uint256 feeAmount
14//     )
15pub fn compute_swap_step(
16    sqrt_ratio_current_x_96: U256,
17    sqrt_ratio_target_x_96: U256,
18    liquidity: u128,
19    amount_remaining: I256,
20    fee_pips: u32,
21) -> Result<(U256, U256, U256, U256), UniswapV3MathError> {
22    let zero_for_one = sqrt_ratio_current_x_96 >= sqrt_ratio_target_x_96;
23    let exact_in = amount_remaining >= I256::ZERO;
24
25    let sqrt_ratio_next_x_96: U256;
26    let mut amount_in = U256::ZERO;
27    let mut amount_out = U256::ZERO;
28
29    if exact_in {
30        let amount_remaining_less_fee = mul_div(
31            amount_remaining.into_raw(),
32            U256::from(1e6 as u32 - fee_pips),    //1e6 - fee_pips
33            U256::from_limbs([1000000, 0, 0, 0]), //1e6
34        )?;
35
36        amount_in = if zero_for_one {
37            _get_amount_0_delta(sqrt_ratio_target_x_96, sqrt_ratio_current_x_96, liquidity, true)?
38        } else {
39            _get_amount_1_delta(sqrt_ratio_current_x_96, sqrt_ratio_target_x_96, liquidity, true)?
40        };
41
42        if amount_remaining_less_fee >= amount_in {
43            sqrt_ratio_next_x_96 = sqrt_ratio_target_x_96;
44        } else {
45            sqrt_ratio_next_x_96 =
46                get_next_sqrt_price_from_input(sqrt_ratio_current_x_96, liquidity, amount_remaining_less_fee, zero_for_one)?;
47        }
48    } else {
49        amount_out = if zero_for_one {
50            _get_amount_1_delta(sqrt_ratio_target_x_96, sqrt_ratio_current_x_96, liquidity, false)?
51        } else {
52            _get_amount_0_delta(sqrt_ratio_current_x_96, sqrt_ratio_target_x_96, liquidity, false)?
53        };
54
55        sqrt_ratio_next_x_96 = if (-amount_remaining).into_raw() >= amount_out {
56            sqrt_ratio_target_x_96
57        } else {
58            get_next_sqrt_price_from_output(sqrt_ratio_current_x_96, liquidity, (-amount_remaining).into_raw(), zero_for_one)?
59        };
60    }
61
62    let max = sqrt_ratio_target_x_96 == sqrt_ratio_next_x_96;
63
64    if zero_for_one {
65        if !max || !exact_in {
66            amount_in = _get_amount_0_delta(sqrt_ratio_next_x_96, sqrt_ratio_current_x_96, liquidity, true)?
67        }
68
69        if !max || exact_in {
70            amount_out = _get_amount_1_delta(sqrt_ratio_next_x_96, sqrt_ratio_current_x_96, liquidity, false)?
71        }
72    } else {
73        if !max || !exact_in {
74            amount_in = _get_amount_1_delta(sqrt_ratio_current_x_96, sqrt_ratio_next_x_96, liquidity, true)?
75        }
76
77        if !max || exact_in {
78            amount_out = _get_amount_0_delta(sqrt_ratio_current_x_96, sqrt_ratio_next_x_96, liquidity, false)?
79        }
80    }
81
82    if !exact_in && amount_out > (-amount_remaining).into_raw() {
83        amount_out = (-amount_remaining).into_raw();
84    }
85
86    if exact_in && sqrt_ratio_next_x_96 != sqrt_ratio_target_x_96 {
87        let fee_amount = amount_remaining.into_raw() - amount_in;
88        Ok((sqrt_ratio_next_x_96, amount_in, amount_out, fee_amount))
89    } else {
90        let fee_amount = mul_div_rounding_up(amount_in, U256::from(fee_pips), U256::from(1e6 as u32 - fee_pips))?;
91
92        Ok((sqrt_ratio_next_x_96, amount_in, amount_out, fee_amount))
93    }
94}
95
96#[cfg(test)]
97mod test {
98
99    use crate::sqrt_price_math::{get_next_sqrt_price_from_input, get_next_sqrt_price_from_output};
100    use crate::swap_math::compute_swap_step;
101    use crate::U256_1;
102    use alloy::primitives::{I256, U256};
103    use std::str::FromStr;
104
105    #[allow(unused)]
106    #[test]
107    fn test_compute_swap_step() {
108        //------------------------------------------------------------
109
110        //exact amount in that gets capped at price target in one for zero
111        let price = U256::from_str("79228162514264337593543950336").unwrap();
112        let price_target = U256::from_str("79623317895830914510639640423").unwrap();
113        let liquidity = 2e18 as u128;
114        let amount = I256::from_str("1000000000000000000").unwrap();
115        let fee = 600;
116        let zero_for_one = false;
117
118        let (sqrt_p, amount_in, amount_out, fee_amount) = compute_swap_step(price, price_target, liquidity, amount, fee).unwrap();
119
120        assert_eq!(sqrt_p, U256::from_str("79623317895830914510639640423").unwrap());
121
122        assert_eq!(amount_in, U256::from_str("9975124224178055").unwrap());
123        assert_eq!(fee_amount, U256::from_str("5988667735148").unwrap());
124        assert_eq!(amount_out, U256::from_str("9925619580021728").unwrap());
125
126        assert!(amount_in + fee_amount < U256::from_limbs(*amount.as_limbs()));
127
128        let price_after_whole_input_amount = get_next_sqrt_price_from_input(price, liquidity, amount_in, zero_for_one).unwrap();
129
130        assert_eq!(sqrt_p, price_target);
131        assert!(sqrt_p < price_after_whole_input_amount);
132
133        //------------------------------------------------------------
134
135        //exact amount out that gets capped at price target in one for zero
136        let price = U256::from_str("79228162514264337593543950336").unwrap();
137        let price_target = U256::from_str("79623317895830914510639640423").unwrap();
138        let liquidity = 2e18 as u128;
139        let amount = I256::from_str("-1000000000000000000").unwrap();
140        let fee = 600;
141        let zero_for_one = false;
142
143        let (sqrt_p, amount_in, amount_out, fee_amount) = compute_swap_step(price, price_target, liquidity, amount, fee).unwrap();
144
145        assert_eq!(amount_in, U256::from_str("9975124224178055").unwrap());
146        assert_eq!(fee_amount, U256::from_str("5988667735148").unwrap());
147        assert_eq!(amount_out, U256::from_str("9925619580021728").unwrap());
148        assert!(amount_out < (amount * -I256::ONE).into_raw());
149
150        assert!(amount_in + fee_amount < U256::from_limbs(*amount.as_limbs()));
151
152        let price_after_whole_output_amount =
153            get_next_sqrt_price_from_output(price, liquidity, (amount * -I256::ONE).into_raw(), zero_for_one).unwrap();
154
155        assert_eq!(sqrt_p, price_target);
156        assert!(sqrt_p < price_after_whole_output_amount);
157
158        //------------------------------------------------------------
159
160        //exact amount in that is fully spent in one for zero
161        let price = U256::from_str("79228162514264337593543950336").unwrap();
162        let price_target = U256::from_str("0xe6666666666666666666666666").unwrap();
163        let liquidity = 2e18 as u128;
164        let amount = I256::from_str("1000000000000000000").unwrap();
165        let fee = 600;
166        let zero_for_one = false;
167
168        let (sqrt_p, amount_in, amount_out, fee_amount) = compute_swap_step(price, price_target, liquidity, amount, fee).unwrap();
169
170        assert_eq!(amount_in, U256::from_str("999400000000000000").unwrap());
171        assert_eq!(fee_amount, U256::from_str("600000000000000").unwrap());
172        assert_eq!(amount_out, U256::from_str("666399946655997866").unwrap());
173        assert_eq!(amount_in + fee_amount, amount.into_raw());
174
175        let price_after_whole_input_amount_less_fee =
176            get_next_sqrt_price_from_input(price, liquidity, (amount - I256::from_raw(fee_amount)).into_raw(), zero_for_one).unwrap();
177
178        assert!(sqrt_p < price_target);
179        assert_eq!(sqrt_p, price_after_whole_input_amount_less_fee);
180
181        //------------------------------------------------------------
182
183        //exact amount out that is fully received in one for zero
184        let price = U256::from_str("79228162514264337593543950336").unwrap();
185        let price_target = U256::from_str("792281625142643375935439503360").unwrap();
186        let liquidity = 2e18 as u128;
187        let amount = I256::from_str("1000000000000000000").unwrap() * -I256::ONE;
188        let fee = 600;
189        let zero_for_one = false;
190
191        let (sqrt_p, amount_in, amount_out, fee_amount) = compute_swap_step(price, price_target, liquidity, amount, fee).unwrap();
192
193        assert_eq!(amount_in, U256::from_str("2000000000000000000").unwrap());
194        assert_eq!(fee_amount, U256::from_str("1200720432259356").unwrap());
195        assert_eq!(amount_out, (amount * -I256::ONE).into_raw());
196
197        let price_after_whole_output_amount =
198            get_next_sqrt_price_from_output(price, liquidity, (amount * -I256::ONE).into_raw(), zero_for_one).unwrap();
199        //sqrtPrice 158456325028528675187087900672
200        //price_after_whole_output_amount Should be: 158456325028528675187087900672
201        // sqrtp: 158456325028528675187087900672, price_after_whole output amount: 118842243771396506390315925504
202
203        assert!(sqrt_p < price_target);
204        //TODO:FIXME: failing
205        println!("sqrtp: {sqrt_p:?}, price_after_whole output amount: {price_after_whole_output_amount:?}");
206        assert_eq!(sqrt_p, price_after_whole_output_amount);
207
208        //------------------------------------------------------------
209
210        //amount out is capped at the desired amount out
211        let (sqrt_p, amount_in, amount_out, fee_amount) = compute_swap_step(
212            U256::from_str("417332158212080721273783715441582").unwrap(),
213            U256::from_str("1452870262520218020823638996").unwrap(),
214            159344665391607089467575320103_u128,
215            I256::from_str("-1").unwrap(),
216            1,
217        )
218        .unwrap();
219
220        assert_eq!(amount_in, U256::from_str("1").unwrap());
221        assert_eq!(fee_amount, U256::from_str("1").unwrap());
222        assert_eq!(amount_out, U256::from_str("1").unwrap());
223        assert_eq!(sqrt_p, U256::from_str("417332158212080721273783715441581").unwrap());
224
225        //------------------------------------------------------------
226
227        //target price of 1 uses partial input amount
228        let (sqrt_p, amount_in, amount_out, fee_amount) = compute_swap_step(
229            U256::from_str("2").unwrap(),
230            U256::from_str("1").unwrap(),
231            1_u128,
232            I256::from_str("3915081100057732413702495386755767").unwrap(),
233            1,
234        )
235        .unwrap();
236
237        assert_eq!(amount_in, U256::from_str("39614081257132168796771975168").unwrap());
238        assert_eq!(fee_amount, U256::from_str("39614120871253040049813").unwrap());
239        assert!(amount_in + fee_amount < U256::from_str("3915081100057732413702495386755767").unwrap());
240        assert_eq!(amount_out, U256::from_str("0").unwrap());
241
242        assert_eq!(sqrt_p, U256::from_str("1").unwrap());
243
244        //------------------------------------------------------------
245
246        //entire input amount taken as fee
247        let (sqrt_p, amount_in, amount_out, fee_amount) = compute_swap_step(
248            U256::from_str("2413").unwrap(),
249            U256::from_str("79887613182836312").unwrap(),
250            1985041575832132834610021537970_u128,
251            I256::from_str("10").unwrap(),
252            1872,
253        )
254        .unwrap();
255
256        assert_eq!(amount_in, U256::from_str("0").unwrap());
257        assert_eq!(fee_amount, U256::from_str("10").unwrap());
258        assert_eq!(amount_out, U256::from_str("0").unwrap());
259        assert_eq!(sqrt_p, U256::from_str("2413").unwrap());
260
261        //------------------------------------------------------------
262
263        //handles intermediate insufficient liquidity in zero for one exact output case
264
265        let price = U256::from_str("20282409603651670423947251286016").unwrap();
266        let price_target = price * U256::from(11) / U256::from(10);
267        let liquidity = 1024;
268        // virtual reserves of one are only 4
269        // https://www.wolframalpha.com/input/?i=1024+%2F+%2820282409603651670423947251286016+%2F+2**96%29
270        let amount_remaining = -I256::from_limbs(*U256::from(4).as_limbs());
271        let fee = 3000;
272
273        let (sqrt_p, amount_in, amount_out, fee_amount) = compute_swap_step(price, price_target, liquidity, amount_remaining, fee).unwrap();
274
275        assert_eq!(amount_out, U256::ZERO);
276        assert_eq!(sqrt_p, price_target);
277        assert_eq!(amount_in, U256::from(26215));
278        assert_eq!(fee_amount, U256::from(79));
279
280        //------------------------------------------------------------
281
282        //handles intermediate insufficient liquidity in one for zero exact output case
283
284        let price = U256::from_str("20282409603651670423947251286016").unwrap();
285
286        let price_target = price * U256::from(9) / U256::from(10);
287        let liquidity = 1024;
288        // virtual reserves of zero are only 262144
289        // https://www.wolframalpha.com/input/?i=1024+*+%2820282409603651670423947251286016+%2F+2**96%29
290        let amount_remaining = -I256::from_limbs(*U256::from(263000).as_limbs());
291        let fee = 3000;
292
293        let (sqrt_p, amount_in, amount_out, fee_amount) = compute_swap_step(price, price_target, liquidity, amount_remaining, fee).unwrap();
294
295        assert_eq!(amount_out, U256::from(26214));
296        assert_eq!(sqrt_p, price_target);
297        assert_eq!(amount_in, U256_1);
298        assert_eq!(fee_amount, U256_1);
299    }
300}
kabu_defi_uniswap_v3_math/swap_math.rs

kabu_defi_uniswap_v3_math/
swap_math.rs
