kabu_defi_pools/virtual_impl/
uniswapv3.rs

1use crate::UniswapV3Pool;
2use crate::db_reader::UniswapV3DbReader;
3use crate::virtual_impl::tick_provider::TickProviderEVMDB;
4use alloy::primitives::{Address, I256, U256};
5use eyre::eyre;
6use kabu_defi_uniswap_v3_math::tick_math::{MAX_SQRT_RATIO, MAX_TICK, MIN_SQRT_RATIO, MIN_TICK};
7use kabu_evm_db::KabuDBError;
8use kabu_types_market::Pool;
9use revm::DatabaseRef;
10
11pub struct UniswapV3PoolVirtual;
12
13/* Unused useful constants
14pub const U256_0X100000000: U256 = U256::from_limbs([4294967296, 0, 0, 0]);
15pub const U256_0X10000: U256 = U256::from_limbs([65536, 0, 0, 0]);
16pub const U256_0X100: U256 = U256::from_limbs([256, 0, 0, 0]);
17pub const U256_255: U256 = U256::from_limbs([255, 0, 0, 0]);
18pub const U256_192: U256 = U256::from_limbs([192, 0, 0, 0]);
19pub const U256_191: U256 = U256::from_limbs([191, 0, 0, 0]);
20pub const U256_128: U256 = U256::from_limbs([128, 0, 0, 0]);
21pub const U256_64: U256 = U256::from_limbs([64, 0, 0, 0]);
22pub const U256_32: U256 = U256::from_limbs([32, 0, 0, 0]);
23pub const U256_16: U256 = U256::from_limbs([16, 0, 0, 0]);
24pub const U256_8: U256 = U256::from_limbs([8, 0, 0, 0]);
25pub const U256_4: U256 = U256::from_limbs([4, 0, 0, 0]);
26pub const U256_2: U256 = U256::from_limbs([2, 0, 0, 0]);
27
28pub const POPULATE_TICK_DATA_STEP: u64 = 100000;
29
30pub const Q128: U256 = U256::from_limbs([0, 0, 1, 0]);
31pub const Q224: U256 = U256::from_limbs([0, 0, 0, 4294967296]);
32
33pub const U128_0X10000000000000000: u128 = 18446744073709551616;
34pub const U256_0XFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF: U256 = U256::from_limbs([
35    18446744073709551615,
36    18446744073709551615,
37    18446744073709551615,
38    0,
39]);
40pub const U256_0XFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF: U256 =
41    U256::from_limbs([18446744073709551615, 18446744073709551615, 0, 0]);
42
43*/
44
45// commonly used U256s
46pub const U256_1: U256 = U256::from_limbs([1, 0, 0, 0]);
47
48// Uniswap V3 specific
49
50// Others
51
52pub struct CurrentState {
53    amount_specified_remaining: I256,
54    amount_calculated: I256,
55    sqrt_price_x_96: U256,
56    tick: i32,
57    liquidity: u128,
58}
59
60#[derive(Default)]
61pub struct StepComputations {
62    pub sqrt_price_start_x_96: U256,
63    pub tick_next: i32,
64    pub initialized: bool,
65    pub sqrt_price_next_x96: U256,
66    pub amount_in: U256,
67    pub amount_out: U256,
68    pub fee_amount: U256,
69}
70
71#[allow(dead_code)]
72pub struct Tick {
73    pub liquidity_gross: u128,
74    pub liquidity_net: i128,
75    pub fee_growth_outside_0_x_128: U256,
76    pub fee_growth_outside_1_x_128: U256,
77    pub tick_cumulative_outside: U256,
78    pub seconds_per_liquidity_outside_x_128: U256,
79    pub seconds_outside: u32,
80    pub initialized: bool,
81}
82
83impl UniswapV3PoolVirtual {
84    pub fn simulate_swap_in_amount_provider<DB: DatabaseRef<Error = KabuDBError> + ?Sized>(
85        db: &DB,
86        pool: &UniswapV3Pool,
87        token_in: &Address,
88        amount_in: U256,
89    ) -> eyre::Result<U256> {
90        if amount_in.is_zero() {
91            return Ok(U256::ZERO);
92        }
93
94        let zero_for_one = *token_in == pool.get_tokens()[0];
95
96        // Set sqrt_price_limit_x_96 to the max or min sqrt price in the pool depending on zero_for_one
97        let sqrt_price_limit_x_96 = if zero_for_one { MIN_SQRT_RATIO + U256_1 } else { MAX_SQRT_RATIO - U256_1 };
98
99        let pool_address = match pool.get_address() {
100            kabu_types_market::PoolId::Address(addr) => addr,
101            kabu_types_market::PoolId::B256(_) => return Err(eyre!("B256 pool ID variant not supported")),
102        };
103
104        let slot0 = UniswapV3DbReader::slot0(&db, pool_address)?;
105        let liquidity = UniswapV3DbReader::liquidity(&db, pool_address)?;
106        let tick_spacing = pool.tick_spacing();
107        let fee = pool.fee;
108
109        // Initialize a mutable state struct to hold the dynamic simulated state of the pool
110        let mut current_state = CurrentState {
111            sqrt_price_x_96: slot0.sqrtPriceX96.to(),              //Active price on the pool
112            amount_calculated: I256::ZERO,                         //Amount of token_out that has been calculated
113            amount_specified_remaining: I256::from_raw(amount_in), //Amount of token_in that has not been swapped
114            tick: slot0.tick.as_i32(),                             //Current i24 tick of the pool
115            liquidity,                                             //Current available liquidity in the tick range
116        };
117
118        let tick_provider = TickProviderEVMDB::new(db, pool_address);
119
120        while current_state.amount_specified_remaining != I256::ZERO && current_state.sqrt_price_x_96 != sqrt_price_limit_x_96 {
121            // Initialize a new step struct to hold the dynamic state of the pool at each step
122            let mut step = StepComputations {
123                // Set the sqrt_price_start_x_96 to the current sqrt_price_x_96
124                sqrt_price_start_x_96: current_state.sqrt_price_x_96,
125                ..Default::default()
126            };
127
128            // Get the next tick from the current tick
129            (step.tick_next, step.initialized) = kabu_defi_uniswap_v3_math::tick_bitmap::next_initialized_tick_within_one_word(
130                &tick_provider,
131                current_state.tick,
132                tick_spacing as i32,
133                zero_for_one,
134            )?;
135
136            // ensure that we do not overshoot the min/max tick, as the tick bitmap is not aware of these bounds
137            // Note: this could be removed as we are clamping in the batch contract
138            step.tick_next = step.tick_next.clamp(MIN_TICK, MAX_TICK);
139
140            // Get the next sqrt price from the input amount
141            step.sqrt_price_next_x96 = kabu_defi_uniswap_v3_math::tick_math::get_sqrt_ratio_at_tick(step.tick_next)?;
142
143            // Target spot price
144            let swap_target_sqrt_ratio = if zero_for_one {
145                if step.sqrt_price_next_x96 < sqrt_price_limit_x_96 { sqrt_price_limit_x_96 } else { step.sqrt_price_next_x96 }
146            } else if step.sqrt_price_next_x96 > sqrt_price_limit_x_96 {
147                sqrt_price_limit_x_96
148            } else {
149                step.sqrt_price_next_x96
150            };
151
152            // Compute swap step and update the current state
153            (current_state.sqrt_price_x_96, step.amount_in, step.amount_out, step.fee_amount) =
154                kabu_defi_uniswap_v3_math::swap_math::compute_swap_step(
155                    current_state.sqrt_price_x_96,
156                    swap_target_sqrt_ratio,
157                    current_state.liquidity,
158                    current_state.amount_specified_remaining,
159                    fee,
160                )?;
161
162            // Decrement the amount remaining to be swapped and amount received from the step
163            current_state.amount_specified_remaining = current_state
164                .amount_specified_remaining
165                .overflowing_sub(I256::from_raw(step.amount_in.overflowing_add(step.fee_amount).0))
166                .0;
167
168            current_state.amount_calculated -= I256::from_raw(step.amount_out);
169
170            // If the price moved all the way to the next price, recompute the liquidity change for the next iteration
171            if current_state.sqrt_price_x_96 == step.sqrt_price_next_x96 {
172                if step.initialized {
173                    let mut liquidity_net: i128 =
174                        UniswapV3DbReader::ticks_liquidity_net(&db, pool_address, step.tick_next).unwrap_or_default();
175
176                    // we are on a tick boundary, and the next tick is initialized, so we must charge a protocol fee
177                    if zero_for_one {
178                        liquidity_net = -liquidity_net;
179                    }
180
181                    current_state.liquidity = if liquidity_net < 0 {
182                        if current_state.liquidity < (-liquidity_net as u128) {
183                            return Err(eyre!("LIQUIDITY_UNDERFLOW"));
184                        } else {
185                            current_state.liquidity - (-liquidity_net as u128)
186                        }
187                    } else {
188                        current_state.liquidity + (liquidity_net as u128)
189                    };
190                }
191                // Increment the current tick
192                current_state.tick = if zero_for_one { step.tick_next.wrapping_sub(1) } else { step.tick_next }
193                // If the current_state sqrt price is not equal to the step sqrt price, then we are not on the same tick.
194                // Update the current_state.tick to the tick at the current_state.sqrt_price_x_96
195            } else if current_state.sqrt_price_x_96 != step.sqrt_price_start_x_96 {
196                current_state.tick = kabu_defi_uniswap_v3_math::tick_math::get_tick_at_sqrt_ratio(current_state.sqrt_price_x_96)?;
197            }
198        }
199
200        if current_state.amount_specified_remaining.is_zero() {
201            let amount_out = (-current_state.amount_calculated).into_raw();
202            tracing::trace!("AmountOut : {amount_out}");
203            Ok(amount_out)
204        } else {
205            Err(eyre!("NOT_ENOUGH_LIQUIDITY"))
206        }
207    }
208
209    pub fn simulate_swap_out_amount_provided<DB: DatabaseRef<Error = KabuDBError> + ?Sized>(
210        db: &DB,
211        pool: &UniswapV3Pool,
212        token_in: &Address,
213        amount_out: U256,
214    ) -> eyre::Result<U256> {
215        if amount_out.is_zero() {
216            return Ok(U256::ZERO);
217        }
218
219        let zero_for_one = *token_in == pool.get_tokens()[0];
220
221        // Set sqrt_price_limit_x_96 to the max or min sqrt price in the pool depending on zero_for_one
222        let sqrt_price_limit_x_96 = if zero_for_one { MIN_SQRT_RATIO + U256_1 } else { MAX_SQRT_RATIO - U256_1 };
223
224        let pool_address = match pool.get_address() {
225            kabu_types_market::PoolId::Address(addr) => addr,
226            kabu_types_market::PoolId::B256(_) => return Err(eyre!("B256 pool ID variant not supported")),
227        };
228
229        let slot0 = UniswapV3DbReader::slot0(&db, pool_address)?;
230        let liquidity = UniswapV3DbReader::liquidity(&db, pool_address)?;
231        let tick_spacing = pool.tick_spacing();
232        let fee = pool.fee;
233
234        // Initialize a mutable state struct to hold the dynamic simulated state of the pool
235        let mut current_state = CurrentState {
236            sqrt_price_x_96: slot0.sqrtPriceX96.to(),                //Active price on the pool
237            amount_calculated: I256::ZERO,                           //Amount of token_out that has been calculated
238            amount_specified_remaining: -I256::from_raw(amount_out), //Amount of token_in that has not been swapped
239            tick: slot0.tick.as_i32(),                               //Current i24 tick of the pool
240            liquidity,                                               //Current available liquidity in the tick range
241        };
242
243        while current_state.amount_specified_remaining != I256::ZERO && current_state.sqrt_price_x_96 != sqrt_price_limit_x_96 {
244            // Initialize a new step struct to hold the dynamic state of the pool at each step
245            let mut step = StepComputations {
246                // Set the sqrt_price_start_x_96 to the current sqrt_price_x_96
247                sqrt_price_start_x_96: current_state.sqrt_price_x_96,
248                ..Default::default()
249            };
250
251            let tick_provider = TickProviderEVMDB::new(&db, pool_address);
252
253            // Get the next tick from the current tick
254            (step.tick_next, step.initialized) = kabu_defi_uniswap_v3_math::tick_bitmap::next_initialized_tick_within_one_word(
255                &tick_provider,
256                current_state.tick,
257                tick_spacing as i32,
258                zero_for_one,
259            )?;
260
261            // ensure that we do not overshoot the min/max tick, as the tick bitmap is not aware of these bounds
262            // Note: this could be removed as we are clamping in the batch contract
263            step.tick_next = step.tick_next.clamp(MIN_TICK, MAX_TICK);
264
265            // Get the next sqrt price from the input amount
266            step.sqrt_price_next_x96 = kabu_defi_uniswap_v3_math::tick_math::get_sqrt_ratio_at_tick(step.tick_next)?;
267
268            // Target spot price
269            let swap_target_sqrt_ratio = if zero_for_one {
270                if step.sqrt_price_next_x96 < sqrt_price_limit_x_96 { sqrt_price_limit_x_96 } else { step.sqrt_price_next_x96 }
271            } else if step.sqrt_price_next_x96 > sqrt_price_limit_x_96 {
272                sqrt_price_limit_x_96
273            } else {
274                step.sqrt_price_next_x96
275            };
276
277            // Compute swap step and update the current state
278            (current_state.sqrt_price_x_96, step.amount_in, step.amount_out, step.fee_amount) =
279                kabu_defi_uniswap_v3_math::swap_math::compute_swap_step(
280                    current_state.sqrt_price_x_96,
281                    swap_target_sqrt_ratio,
282                    current_state.liquidity,
283                    current_state.amount_specified_remaining,
284                    fee,
285                )?;
286
287            // Decrement the amount remaining to be swapped and amount received from the step
288            current_state.amount_specified_remaining =
289                current_state.amount_specified_remaining.overflowing_add(I256::from_raw(step.amount_out)).0;
290
291            current_state.amount_calculated =
292                current_state.amount_calculated.overflowing_add(I256::from_raw(step.amount_in.overflowing_add(step.fee_amount).0)).0;
293
294            // If the price moved all the way to the next price, recompute the liquidity change for the next iteration
295            if current_state.sqrt_price_x_96 == step.sqrt_price_next_x96 {
296                if step.initialized {
297                    let mut liquidity_net: i128 =
298                        UniswapV3DbReader::ticks_liquidity_net(&db, pool_address, step.tick_next).unwrap_or_default();
299
300                    // we are on a tick boundary, and the next tick is initialized, so we must charge a protocol fee
301                    if zero_for_one {
302                        liquidity_net = -liquidity_net;
303                    }
304
305                    current_state.liquidity = if liquidity_net < 0 {
306                        if current_state.liquidity < (-liquidity_net as u128) {
307                            return Err(eyre!("LIQUIDITY_UNDERFLOW"));
308                        } else {
309                            current_state.liquidity - (-liquidity_net as u128)
310                        }
311                    } else {
312                        current_state.liquidity + (liquidity_net as u128)
313                    };
314                }
315                // Increment the current tick
316                current_state.tick = if zero_for_one { step.tick_next.wrapping_sub(1) } else { step.tick_next }
317                // If the current_state sqrt price is not equal to the step sqrt price, then we are not on the same tick.
318                // Update the current_state.tick to the tick at the current_state.sqrt_price_x_96
319            } else if current_state.sqrt_price_x_96 != step.sqrt_price_start_x_96 {
320                current_state.tick = kabu_defi_uniswap_v3_math::tick_math::get_tick_at_sqrt_ratio(current_state.sqrt_price_x_96)?;
321            }
322        }
323
324        if current_state.amount_specified_remaining.is_zero() {
325            let amount_in = current_state.amount_calculated.into_raw();
326
327            tracing::trace!("Amount In : {amount_in}");
328
329            Ok(amount_in)
330        } else {
331            Err(eyre!("NOT_ENOUGH_LIQUIDITY"))
332        }
333    }
334}
335
336#[cfg(test)]
337mod test {
338    use alloy::primitives::U256;
339    use kabu_defi_uniswap_v3_math::full_math::mul_div_rounding_up;
340
341    #[test]
342    fn test_mul_rounding_up() {
343        let amount = U256::from_limbs([1230267133767, 0, 0, 0]);
344        let ret = mul_div_rounding_up(amount, U256::from(500), U256::from(1e6)).unwrap();
345        assert_eq!(ret, U256::from(615133567u128));
346    }
347}
kabu_defi_pools/virtual_impl/uniswapv3.rs

kabu_defi_pools/virtual_impl/
uniswapv3.rs
