kabu_defi_pools/virtual_impl/
uniswapv3.rs

1use crate::db_reader::UniswapV3DbReader;
2use crate::virtual_impl::tick_provider::TickProviderEVMDB;
3use crate::UniswapV3Pool;
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_entities::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_entities::PoolId::Address(addr) => addr,
101            kabu_types_entities::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 {
146                    sqrt_price_limit_x_96
147                } else {
148                    step.sqrt_price_next_x96
149                }
150            } else if step.sqrt_price_next_x96 > sqrt_price_limit_x_96 {
151                sqrt_price_limit_x_96
152            } else {
153                step.sqrt_price_next_x96
154            };
155
156            // Compute swap step and update the current state
157            (current_state.sqrt_price_x_96, step.amount_in, step.amount_out, step.fee_amount) =
158                kabu_defi_uniswap_v3_math::swap_math::compute_swap_step(
159                    current_state.sqrt_price_x_96,
160                    swap_target_sqrt_ratio,
161                    current_state.liquidity,
162                    current_state.amount_specified_remaining,
163                    fee,
164                )?;
165
166            // Decrement the amount remaining to be swapped and amount received from the step
167            current_state.amount_specified_remaining = current_state
168                .amount_specified_remaining
169                .overflowing_sub(I256::from_raw(step.amount_in.overflowing_add(step.fee_amount).0))
170                .0;
171
172            current_state.amount_calculated -= I256::from_raw(step.amount_out);
173
174            // If the price moved all the way to the next price, recompute the liquidity change for the next iteration
175            if current_state.sqrt_price_x_96 == step.sqrt_price_next_x96 {
176                if step.initialized {
177                    let mut liquidity_net: i128 =
178                        UniswapV3DbReader::ticks_liquidity_net(&db, pool_address, step.tick_next).unwrap_or_default();
179
180                    // we are on a tick boundary, and the next tick is initialized, so we must charge a protocol fee
181                    if zero_for_one {
182                        liquidity_net = -liquidity_net;
183                    }
184
185                    current_state.liquidity = if liquidity_net < 0 {
186                        if current_state.liquidity < (-liquidity_net as u128) {
187                            return Err(eyre!("LIQUIDITY_UNDERFLOW"));
188                        } else {
189                            current_state.liquidity - (-liquidity_net as u128)
190                        }
191                    } else {
192                        current_state.liquidity + (liquidity_net as u128)
193                    };
194                }
195                // Increment the current tick
196                current_state.tick = if zero_for_one { step.tick_next.wrapping_sub(1) } else { step.tick_next }
197                // If the current_state sqrt price is not equal to the step sqrt price, then we are not on the same tick.
198                // Update the current_state.tick to the tick at the current_state.sqrt_price_x_96
199            } else if current_state.sqrt_price_x_96 != step.sqrt_price_start_x_96 {
200                current_state.tick = kabu_defi_uniswap_v3_math::tick_math::get_tick_at_sqrt_ratio(current_state.sqrt_price_x_96)?;
201            }
202        }
203
204        if current_state.amount_specified_remaining.is_zero() {
205            let amount_out = (-current_state.amount_calculated).into_raw();
206            tracing::trace!("AmountOut : {amount_out}");
207            Ok(amount_out)
208        } else {
209            Err(eyre!("NOT_ENOUGH_LIQUIDITY"))
210        }
211    }
212
213    pub fn simulate_swap_out_amount_provided<DB: DatabaseRef<Error = KabuDBError> + ?Sized>(
214        db: &DB,
215        pool: &UniswapV3Pool,
216        token_in: &Address,
217        amount_out: U256,
218    ) -> eyre::Result<U256> {
219        if amount_out.is_zero() {
220            return Ok(U256::ZERO);
221        }
222
223        let zero_for_one = *token_in == pool.get_tokens()[0];
224
225        // Set sqrt_price_limit_x_96 to the max or min sqrt price in the pool depending on zero_for_one
226        let sqrt_price_limit_x_96 = if zero_for_one { MIN_SQRT_RATIO + U256_1 } else { MAX_SQRT_RATIO - U256_1 };
227
228        let pool_address = match pool.get_address() {
229            kabu_types_entities::PoolId::Address(addr) => addr,
230            kabu_types_entities::PoolId::B256(_) => return Err(eyre!("B256 pool ID variant not supported")),
231        };
232
233        let slot0 = UniswapV3DbReader::slot0(&db, pool_address)?;
234        let liquidity = UniswapV3DbReader::liquidity(&db, pool_address)?;
235        let tick_spacing = pool.tick_spacing();
236        let fee = pool.fee;
237
238        // Initialize a mutable state struct to hold the dynamic simulated state of the pool
239        let mut current_state = CurrentState {
240            sqrt_price_x_96: slot0.sqrtPriceX96.to(),                //Active price on the pool
241            amount_calculated: I256::ZERO,                           //Amount of token_out that has been calculated
242            amount_specified_remaining: -I256::from_raw(amount_out), //Amount of token_in that has not been swapped
243            tick: slot0.tick.as_i32(),                               //Current i24 tick of the pool
244            liquidity,                                               //Current available liquidity in the tick range
245        };
246
247        while current_state.amount_specified_remaining != I256::ZERO && current_state.sqrt_price_x_96 != sqrt_price_limit_x_96 {
248            // Initialize a new step struct to hold the dynamic state of the pool at each step
249            let mut step = StepComputations {
250                // Set the sqrt_price_start_x_96 to the current sqrt_price_x_96
251                sqrt_price_start_x_96: current_state.sqrt_price_x_96,
252                ..Default::default()
253            };
254
255            let tick_provider = TickProviderEVMDB::new(&db, pool_address);
256
257            // Get the next tick from the current tick
258            (step.tick_next, step.initialized) = kabu_defi_uniswap_v3_math::tick_bitmap::next_initialized_tick_within_one_word(
259                &tick_provider,
260                current_state.tick,
261                tick_spacing as i32,
262                zero_for_one,
263            )?;
264
265            // ensure that we do not overshoot the min/max tick, as the tick bitmap is not aware of these bounds
266            // Note: this could be removed as we are clamping in the batch contract
267            step.tick_next = step.tick_next.clamp(MIN_TICK, MAX_TICK);
268
269            // Get the next sqrt price from the input amount
270            step.sqrt_price_next_x96 = kabu_defi_uniswap_v3_math::tick_math::get_sqrt_ratio_at_tick(step.tick_next)?;
271
272            // Target spot price
273            let swap_target_sqrt_ratio = if zero_for_one {
274                if step.sqrt_price_next_x96 < sqrt_price_limit_x_96 {
275                    sqrt_price_limit_x_96
276                } else {
277                    step.sqrt_price_next_x96
278                }
279            } else if step.sqrt_price_next_x96 > sqrt_price_limit_x_96 {
280                sqrt_price_limit_x_96
281            } else {
282                step.sqrt_price_next_x96
283            };
284
285            // Compute swap step and update the current state
286            (current_state.sqrt_price_x_96, step.amount_in, step.amount_out, step.fee_amount) =
287                kabu_defi_uniswap_v3_math::swap_math::compute_swap_step(
288                    current_state.sqrt_price_x_96,
289                    swap_target_sqrt_ratio,
290                    current_state.liquidity,
291                    current_state.amount_specified_remaining,
292                    fee,
293                )?;
294
295            // Decrement the amount remaining to be swapped and amount received from the step
296            current_state.amount_specified_remaining =
297                current_state.amount_specified_remaining.overflowing_add(I256::from_raw(step.amount_out)).0;
298
299            current_state.amount_calculated =
300                current_state.amount_calculated.overflowing_add(I256::from_raw(step.amount_in.overflowing_add(step.fee_amount).0)).0;
301
302            // If the price moved all the way to the next price, recompute the liquidity change for the next iteration
303            if current_state.sqrt_price_x_96 == step.sqrt_price_next_x96 {
304                if step.initialized {
305                    let mut liquidity_net: i128 =
306                        UniswapV3DbReader::ticks_liquidity_net(&db, pool_address, step.tick_next).unwrap_or_default();
307
308                    // we are on a tick boundary, and the next tick is initialized, so we must charge a protocol fee
309                    if zero_for_one {
310                        liquidity_net = -liquidity_net;
311                    }
312
313                    current_state.liquidity = if liquidity_net < 0 {
314                        if current_state.liquidity < (-liquidity_net as u128) {
315                            return Err(eyre!("LIQUIDITY_UNDERFLOW"));
316                        } else {
317                            current_state.liquidity - (-liquidity_net as u128)
318                        }
319                    } else {
320                        current_state.liquidity + (liquidity_net as u128)
321                    };
322                }
323                // Increment the current tick
324                current_state.tick = if zero_for_one { step.tick_next.wrapping_sub(1) } else { step.tick_next }
325                // If the current_state sqrt price is not equal to the step sqrt price, then we are not on the same tick.
326                // Update the current_state.tick to the tick at the current_state.sqrt_price_x_96
327            } else if current_state.sqrt_price_x_96 != step.sqrt_price_start_x_96 {
328                current_state.tick = kabu_defi_uniswap_v3_math::tick_math::get_tick_at_sqrt_ratio(current_state.sqrt_price_x_96)?;
329            }
330        }
331
332        if current_state.amount_specified_remaining.is_zero() {
333            let amount_in = current_state.amount_calculated.into_raw();
334
335            tracing::trace!("Amount In : {amount_in}");
336
337            Ok(amount_in)
338        } else {
339            Err(eyre!("NOT_ENOUGH_LIQUIDITY"))
340        }
341    }
342}
343
344#[cfg(test)]
345mod test {
346    use alloy::primitives::U256;
347    use kabu_defi_uniswap_v3_math::full_math::mul_div_rounding_up;
348
349    #[test]
350    fn test_mul_rounding_up() {
351        let amount = U256::from_limbs([1230267133767, 0, 0, 0]);
352        let ret = mul_div_rounding_up(amount, U256::from(500), U256::from(1e6)).unwrap();
353        assert_eq!(ret, U256::from(615133567u128));
354    }
355}
kabu_defi_pools/virtual_impl/uniswapv3.rs

kabu_defi_pools/virtual_impl/
uniswapv3.rs
