Skip to main content

tycho_simulation/evm/protocol/uniswap_v4/
state.rs

1use std::{any::Any, collections::HashMap, fmt};
2
3use alloy::primitives::{Address, Sign, I256, U256};
4use num_bigint::BigUint;
5use num_traits::{CheckedSub, ToPrimitive, Zero};
6use revm::primitives::I128;
7use tracing::trace;
8use tycho_common::{
9    dto::ProtocolStateDelta,
10    models::token::Token,
11    simulation::{
12        errors::{SimulationError, TransitionError},
13        protocol_sim::{
14            Balances, GetAmountOutResult, PoolSwap, ProtocolSim, QueryPoolSwapParams,
15            SwapConstraint,
16        },
17    },
18    Bytes,
19};
20
21use super::hooks::utils::{has_permission, HookOptions};
22use crate::{
23    evm::protocol::{
24        clmm::clmm_swap_to_price,
25        safe_math::{safe_add_u256, safe_sub_u256},
26        u256_num::u256_to_biguint,
27        uniswap_v4::hooks::{
28            hook_handler::HookHandler,
29            models::{
30                AfterSwapParameters, BalanceDelta, BeforeSwapDelta, BeforeSwapParameters,
31                StateContext, SwapParams,
32            },
33        },
34        utils::{
35            add_fee_markup,
36            uniswap::{
37                i24_be_bytes_to_i32, liquidity_math,
38                lp_fee::{self, is_dynamic},
39                sqrt_price_math::{get_amount0_delta, get_amount1_delta, sqrt_price_q96_to_f64},
40                swap_math,
41                tick_list::{TickInfo, TickList, TickListErrorKind},
42                tick_math::{
43                    get_sqrt_ratio_at_tick, get_tick_at_sqrt_ratio, MAX_SQRT_RATIO, MAX_TICK,
44                    MIN_SQRT_RATIO, MIN_TICK,
45                },
46                StepComputation, SwapResults, SwapState,
47            },
48        },
49        vm::constants::EXTERNAL_ACCOUNT,
50    },
51    impl_non_serializable_protocol,
52};
53
54// Fixed overhead per swap: covers router overhead, executor preamble (decode, sync,
55// unlock/callback pattern), and token transfer-in.
56const SWAP_BASE_GAS: u64 = 185_000;
57// Per loop: PoolManager.swap bitmap lookup + sqrt math + computeSwapStep.
58// V4's singleton PoolManager hits warmer storage than V3 standalone pools: ~3,500/loop.
59const GAS_PER_BITMAP_LOOKUP: u64 = 3_500;
60// Initialized tick crossing: _updateTick() updates feeGrowthOutside0/1 (2 SSTOREs).
61// Warm ≈ 10–17k, cold ≈ 40–52k. 17,540 reflects warm scenario.
62const GAS_PER_TICK: u64 = 17_540;
63// Settlement overhead within swapExactInputSingle: _settle() + _getFullCredit() + misc.
64const V4_CALLBACK_SETTLEMENT_GAS: u64 = 30_000;
65// Conservative max gas budget for a single swap (Ethereum transaction gas limit)
66const MAX_SWAP_GAS: u64 = 16_700_000;
67const MAX_TICKS_CROSSED: u64 = (MAX_SWAP_GAS - SWAP_BASE_GAS) / GAS_PER_TICK;
68
69#[derive(Clone)]
70pub struct UniswapV4State {
71    liquidity: u128,
72    sqrt_price: U256,
73    fees: UniswapV4Fees,
74    tick: i32,
75    ticks: TickList,
76    tick_spacing: i32,
77    pub hook: Option<Box<dyn HookHandler>>,
78}
79
80impl_non_serializable_protocol!(UniswapV4State, "not supported due vm state deps");
81
82impl fmt::Debug for UniswapV4State {
83    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
84        f.debug_struct("UniswapV4State")
85            .field("liquidity", &self.liquidity)
86            .field("sqrt_price", &self.sqrt_price)
87            .field("fees", &self.fees)
88            .field("tick", &self.tick)
89            .field("tick_spacing", &self.tick_spacing)
90            .finish_non_exhaustive()
91    }
92}
93
94impl PartialEq for UniswapV4State {
95    fn eq(&self, other: &Self) -> bool {
96        match (&self.hook, &other.hook) {
97            (Some(a), Some(b)) => a.is_equal(&**b),
98            (None, None) => true,
99            _ => false,
100        }
101    }
102}
103
104impl Eq for UniswapV4State {}
105
106#[derive(Clone, Debug, PartialEq, Eq)]
107pub struct UniswapV4Fees {
108    // Protocol fees in the zero for one direction
109    pub zero_for_one: u32,
110    // Protocol fees in the one for zero direction
111    pub one_for_zero: u32,
112    // Liquidity providers fees
113    pub lp_fee: u32,
114}
115
116impl UniswapV4Fees {
117    pub fn new(zero_for_one: u32, one_for_zero: u32, lp_fee: u32) -> Self {
118        Self { zero_for_one, one_for_zero, lp_fee }
119    }
120
121    fn calculate_swap_fees_pips(&self, zero_for_one: bool, lp_fee_override: Option<u32>) -> u32 {
122        let protocol_fee = if zero_for_one { self.zero_for_one } else { self.one_for_zero };
123        let lp_fee = lp_fee_override.unwrap_or_else(|| {
124            // If a protocol has dynamic fees,
125            if is_dynamic(self.lp_fee) {
126                0
127            } else {
128                self.lp_fee
129            }
130        });
131
132        // UniswapV4 formula: protocolFee + lpFee - (protocolFee * lpFee / 1_000_000)
133        // Source: https://raw.githubusercontent.com/Uniswap/v4-core/main/src/libraries/ProtocolFeeLibrary.sol
134        // This accounts for the fact that protocol fee is taken first, then LP fee applies to
135        // remainder
136        protocol_fee + lp_fee - ((protocol_fee as u64 * lp_fee as u64 / 1_000_000) as u32)
137    }
138}
139
140impl UniswapV4State {
141    /// Creates a new `UniswapV4State` with specified values.
142    pub fn new(
143        liquidity: u128,
144        sqrt_price: U256,
145        fees: UniswapV4Fees,
146        tick: i32,
147        tick_spacing: i32,
148        ticks: Vec<TickInfo>,
149    ) -> Result<Self, SimulationError> {
150        let tick_spacing_u16 = tick_spacing.try_into().map_err(|_| {
151            // even though it's given as int24, tick_spacing must be positive, see here:
152            // https://github.com/Uniswap/v4-core/blob/a22414e4d7c0d0b0765827fe0a6c20dfd7f96291/src/libraries/TickMath.sol#L25-L28
153            SimulationError::FatalError(format!(
154                "tick_spacing {} must be positive (int24 -> u16 conversion failed)",
155                tick_spacing
156            ))
157        })?;
158        let tick_list = TickList::from(tick_spacing_u16, ticks)?;
159        Ok(UniswapV4State {
160            liquidity,
161            sqrt_price,
162            fees,
163            tick,
164            ticks: tick_list,
165            tick_spacing,
166            hook: None,
167        })
168    }
169
170    fn swap(
171        &self,
172        zero_for_one: bool,
173        amount_specified: I256,
174        sqrt_price_limit: Option<U256>,
175        lp_fee_override: Option<u32>,
176    ) -> Result<SwapResults, SimulationError> {
177        if amount_specified == I256::ZERO {
178            return Ok(SwapResults {
179                amount_calculated: I256::ZERO,
180                amount_specified: I256::ZERO,
181                amount_remaining: I256::ZERO,
182                sqrt_price: self.sqrt_price,
183                liquidity: self.liquidity,
184                tick: self.tick,
185                gas_used: U256::from(3_000), // baseline gas cost for no-op swap
186            });
187        }
188
189        if self.liquidity == 0 {
190            return Err(SimulationError::RecoverableError("No liquidity".to_string()));
191        }
192        let price_limit = if let Some(limit) = sqrt_price_limit {
193            limit
194        } else if zero_for_one {
195            safe_add_u256(MIN_SQRT_RATIO, U256::from(1u64))?
196        } else {
197            safe_sub_u256(MAX_SQRT_RATIO, U256::from(1u64))?
198        };
199
200        let price_limit_valid = if zero_for_one {
201            price_limit > MIN_SQRT_RATIO && price_limit < self.sqrt_price
202        } else {
203            price_limit < MAX_SQRT_RATIO && price_limit > self.sqrt_price
204        };
205        if !price_limit_valid {
206            return Err(SimulationError::InvalidInput("Price limit out of range".into(), None));
207        }
208
209        let exact_input = amount_specified < I256::ZERO;
210
211        let mut state = SwapState {
212            amount_remaining: amount_specified,
213            amount_calculated: I256::ZERO,
214            sqrt_price: self.sqrt_price,
215            tick: self.tick,
216            liquidity: self.liquidity,
217        };
218        let mut gas_used = U256::from(SWAP_BASE_GAS);
219
220        while state.amount_remaining != I256::ZERO && state.sqrt_price != price_limit {
221            let (mut next_tick, initialized) = match self
222                .ticks
223                .next_initialized_tick_within_one_word(state.tick, zero_for_one)
224            {
225                Ok((tick, init)) => {
226                    gas_used = safe_add_u256(gas_used, U256::from(GAS_PER_BITMAP_LOOKUP))?;
227                    (tick, init)
228                }
229                Err(tick_err) => match tick_err.kind {
230                    TickListErrorKind::TicksExeeded => {
231                        let mut new_state = self.clone();
232                        new_state.liquidity = state.liquidity;
233                        new_state.tick = state.tick;
234                        new_state.sqrt_price = state.sqrt_price;
235                        return Err(SimulationError::InvalidInput(
236                            "Ticks exceeded".into(),
237                            Some(GetAmountOutResult::new(
238                                u256_to_biguint(state.amount_calculated.abs().into_raw()),
239                                u256_to_biguint(gas_used),
240                                Box::new(new_state),
241                            )),
242                        ));
243                    }
244                    _ => return Err(SimulationError::FatalError("Unknown error".to_string())),
245                },
246            };
247
248            next_tick = next_tick.clamp(MIN_TICK, MAX_TICK);
249
250            let sqrt_price_next = get_sqrt_ratio_at_tick(next_tick)?;
251            let fee_pips = self
252                .fees
253                .calculate_swap_fees_pips(zero_for_one, lp_fee_override);
254
255            let sqrt_price_start = state.sqrt_price;
256            let (sqrt_price, amount_in, amount_out, fee_amount) = swap_math::compute_swap_step(
257                state.sqrt_price,
258                UniswapV4State::get_sqrt_ratio_target(sqrt_price_next, price_limit, zero_for_one),
259                state.liquidity,
260                // The core univ4 swap logic assumes that if the amount is > 0 it's exact in, and
261                // if it's < 0 it's exact out. The compute_swap_step assumes the
262                // opposite (it's like that for univ3).
263                -state.amount_remaining,
264                fee_pips,
265            )?;
266            state.sqrt_price = sqrt_price;
267
268            let step = StepComputation {
269                sqrt_price_start,
270                tick_next: next_tick,
271                initialized,
272                sqrt_price_next,
273                amount_in,
274                amount_out,
275                fee_amount,
276            };
277            if exact_input {
278                state.amount_remaining += I256::checked_from_sign_and_abs(
279                    Sign::Positive,
280                    safe_add_u256(step.amount_in, step.fee_amount)?,
281                )
282                .unwrap();
283                state.amount_calculated -=
284                    I256::checked_from_sign_and_abs(Sign::Positive, step.amount_out).unwrap();
285            } else {
286                state.amount_remaining -=
287                    I256::checked_from_sign_and_abs(Sign::Positive, step.amount_out).unwrap();
288                state.amount_calculated += I256::checked_from_sign_and_abs(
289                    Sign::Positive,
290                    safe_add_u256(step.amount_in, step.fee_amount)?,
291                )
292                .unwrap();
293            }
294            if state.sqrt_price == step.sqrt_price_next {
295                if step.initialized {
296                    let liquidity_raw = self
297                        .ticks
298                        .get_tick(step.tick_next)
299                        .unwrap()
300                        .net_liquidity;
301                    let liquidity_net = if zero_for_one { -liquidity_raw } else { liquidity_raw };
302                    state.liquidity =
303                        liquidity_math::add_liquidity_delta(state.liquidity, liquidity_net)?;
304                    gas_used = safe_add_u256(gas_used, U256::from(GAS_PER_TICK))?;
305                }
306                state.tick = if zero_for_one { step.tick_next - 1 } else { step.tick_next };
307            } else if state.sqrt_price != step.sqrt_price_start {
308                state.tick = get_tick_at_sqrt_ratio(state.sqrt_price)?;
309            }
310        }
311
312        Ok(SwapResults {
313            amount_calculated: state.amount_calculated,
314            amount_specified,
315            amount_remaining: state.amount_remaining,
316            sqrt_price: state.sqrt_price,
317            liquidity: state.liquidity,
318            tick: state.tick,
319            gas_used: safe_add_u256(gas_used, U256::from(V4_CALLBACK_SETTLEMENT_GAS))?,
320        })
321    }
322
323    pub fn set_hook_handler(&mut self, handler: Box<dyn HookHandler>) {
324        self.hook = Some(handler);
325    }
326
327    fn get_sqrt_ratio_target(
328        sqrt_price_next: U256,
329        sqrt_price_limit: U256,
330        zero_for_one: bool,
331    ) -> U256 {
332        let cond1 = if zero_for_one {
333            sqrt_price_next < sqrt_price_limit
334        } else {
335            sqrt_price_next > sqrt_price_limit
336        };
337
338        if cond1 {
339            sqrt_price_limit
340        } else {
341            sqrt_price_next
342        }
343    }
344
345    fn find_limits_experimentally(
346        &self,
347        token_in: Bytes,
348        token_out: Bytes,
349    ) -> Result<(BigUint, BigUint), SimulationError> {
350        // Create dummy token objects with proper addresses. This is fine since `get_amount_out`
351        // only uses the token addresses.
352        let token_in_obj =
353            Token::new(&token_in, "TOKEN_IN", 18, 0, &[Some(10_000)], Default::default(), 100);
354        let token_out_obj =
355            Token::new(&token_out, "TOKEN_OUT", 18, 0, &[Some(10_000)], Default::default(), 100);
356
357        self.find_max_amount(&token_in_obj, &token_out_obj)
358    }
359
360    /// Finds max amount by performing exponential search.
361    ///
362    /// Reasoning:
363    /// - get_amount_out(I256::MAX) will almost always fail, so this will waste time checking values
364    ///   unrealistically high.
365    /// - If you were to start binary search from 1 to 10^76, you'd need hundreds of iterations.
366    ///
367    /// More about exponential search: https://en.wikipedia.org/wiki/Exponential_search
368    ///
369    /// # Returns
370    ///
371    /// Returns a tuple containing the max amount in and max amount out respectively.
372    fn find_max_amount(
373        &self,
374        token_in: &Token,
375        token_out: &Token,
376    ) -> Result<(BigUint, BigUint), SimulationError> {
377        let mut low = BigUint::from(1u64);
378
379        // The max you can swap on a USV4 is I256::MAX is 5.7e76, since input amount is I256.
380        // So start with something much smaller to search for a reasonable upper bound.
381        let mut high = BigUint::from(10u64).pow(18); // 1 ether in wei
382        let mut last_successful_amount_in = BigUint::from(1u64);
383        let mut last_successful_amount_out = BigUint::from(0u64);
384
385        // First, find an upper bound where the swap fails using exponential search.
386        // Save and return both the amount in and amount out.
387        while let Ok(result) = self.get_amount_out(high.clone(), token_in, token_out) {
388            // We haven't found the upper bound yet, increase the attempted upper bound
389            // by order of magnitude and store the last success as the lower bound.
390            low = last_successful_amount_in.clone();
391            last_successful_amount_in = high.clone();
392            last_successful_amount_out = result.amount;
393            high *= BigUint::from(10u64);
394
395            // Stop if we're getting too large for I256 (about 10^75)
396            if high > BigUint::from(10u64).pow(75) {
397                return Ok((last_successful_amount_in, last_successful_amount_out));
398            }
399        }
400
401        // Use binary search to narrow down value between low and high
402        while &high - &low > BigUint::from(1u64) {
403            let mid = (&low + &high) / BigUint::from(2u64);
404
405            match self.get_amount_out(mid.clone(), token_in, token_out) {
406                Ok(result) => {
407                    last_successful_amount_in = mid.clone();
408                    last_successful_amount_out = result.amount;
409                    low = mid;
410                }
411                Err(_) => {
412                    high = mid;
413                }
414            }
415        }
416
417        Ok((last_successful_amount_in, last_successful_amount_out))
418    }
419
420    /// Helper method to check if there are no initialized ticks in either direction
421    fn has_no_initialized_ticks(&self) -> bool {
422        !self.ticks.has_initialized_ticks()
423    }
424}
425
426#[typetag::serde]
427impl ProtocolSim for UniswapV4State {
428    // Not possible to implement correctly with the current interface because we need to know the
429    // swap direction.
430    fn fee(&self) -> f64 {
431        todo!()
432    }
433
434    fn spot_price(&self, base: &Token, quote: &Token) -> Result<f64, SimulationError> {
435        if let Some(hook) = &self.hook {
436            match hook.spot_price(base, quote) {
437                Ok(price) => return Ok(price),
438                Err(SimulationError::RecoverableError(_)) => {
439                    // Calculate spot price by swapping two amounts and use the approximation
440                    // to get the derivative, following the pattern from vm/state.rs
441
442                    // Calculate the first sell amount (x1) as a small amount
443                    let x1 = BigUint::from(10u64).pow(base.decimals) / BigUint::from(100u64); // 0.01 token
444
445                    // Calculate the second sell amount (x2) as x1 + 1% of x1
446                    let x2 = &x1 + (&x1 / BigUint::from(100u64));
447
448                    // Perform swaps to get the received amounts
449                    let y1 = self.get_amount_out(x1.clone(), base, quote)?;
450                    let y2 = self.get_amount_out(x2.clone(), base, quote)?;
451
452                    // Calculate the marginal price
453                    let num = y2
454                        .amount
455                        .checked_sub(&y1.amount)
456                        .ok_or_else(|| {
457                            SimulationError::FatalError(
458                                "Cannot calculate spot price: y2 < y1".to_string(),
459                            )
460                        })?;
461                    let den = x2.checked_sub(&x1).ok_or_else(|| {
462                        SimulationError::FatalError(
463                            "Cannot calculate spot price: x2 < x1".to_string(),
464                        )
465                    })?;
466
467                    if den == BigUint::from(0u64) {
468                        return Err(SimulationError::FatalError(
469                            "Cannot calculate spot price: denominator is zero".to_string(),
470                        ));
471                    }
472
473                    // Convert to f64 and adjust for decimals
474                    let num_f64 = num.to_f64().ok_or_else(|| {
475                        SimulationError::FatalError(
476                            "Failed to convert numerator to f64".to_string(),
477                        )
478                    })?;
479                    let den_f64 = den.to_f64().ok_or_else(|| {
480                        SimulationError::FatalError(
481                            "Failed to convert denominator to f64".to_string(),
482                        )
483                    })?;
484
485                    let token_correction = 10f64.powi(base.decimals as i32 - quote.decimals as i32);
486
487                    return Ok(num_f64 / den_f64 * token_correction);
488                }
489                Err(e) => return Err(e),
490            }
491        }
492
493        let zero_for_one = base < quote;
494        let fee_pips = self
495            .fees
496            .calculate_swap_fees_pips(zero_for_one, None);
497        let fee = fee_pips as f64 / 1_000_000.0;
498
499        let price = if zero_for_one {
500            sqrt_price_q96_to_f64(self.sqrt_price, base.decimals, quote.decimals)?
501        } else {
502            1.0f64 / sqrt_price_q96_to_f64(self.sqrt_price, quote.decimals, base.decimals)?
503        };
504
505        Ok(add_fee_markup(price, fee))
506    }
507
508    fn get_amount_out(
509        &self,
510        amount_in: BigUint,
511        token_in: &Token,
512        token_out: &Token,
513    ) -> Result<GetAmountOutResult, SimulationError> {
514        let zero_for_one = token_in < token_out;
515        let amount_specified = I256::checked_from_sign_and_abs(
516            Sign::Negative,
517            U256::from_be_slice(&amount_in.to_bytes_be()),
518        )
519        .ok_or_else(|| {
520            SimulationError::InvalidInput("I256 overflow: amount_in".to_string(), None)
521        })?;
522
523        let mut amount_to_swap = amount_specified;
524        let mut lp_fee_override: Option<u32> = None;
525        let mut before_swap_gas = 0u64;
526        let mut after_swap_gas = 0u64;
527        let mut before_swap_delta = BeforeSwapDelta(I256::ZERO);
528        let mut storage_overwrites = None;
529
530        let token_in_address = Address::from_slice(&token_in.address);
531        let token_out_address = Address::from_slice(&token_out.address);
532
533        let state_context = StateContext {
534            currency_0: if zero_for_one { token_in_address } else { token_out_address },
535            currency_1: if zero_for_one { token_out_address } else { token_in_address },
536            fees: self.fees.clone(),
537            tick_spacing: self.tick_spacing,
538        };
539
540        let swap_params = SwapParams {
541            zero_for_one,
542            amount_specified: amount_to_swap,
543            sqrt_price_limit: self.sqrt_price,
544        };
545
546        // Check if hook is set and has before_swap permissions
547        if let Some(ref hook) = self.hook {
548            if has_permission(hook.address(), HookOptions::BeforeSwap) {
549                let before_swap_params = BeforeSwapParameters {
550                    context: state_context.clone(),
551                    sender: *EXTERNAL_ACCOUNT,
552                    swap_params: swap_params.clone(),
553                    hook_data: Bytes::new(),
554                };
555
556                let before_swap_result = hook
557                    .before_swap(before_swap_params, None, None)
558                    .map_err(|e| {
559                        SimulationError::FatalError(format!(
560                            "BeforeSwap hook simulation failed: {e:?}"
561                        ))
562                    })?;
563
564                before_swap_gas = before_swap_result.gas_estimate;
565                before_swap_delta = before_swap_result.result.amount_delta;
566                storage_overwrites = Some(before_swap_result.result.overwrites);
567
568                // Convert amountDelta to amountToSwap as per Uniswap V4 spec
569                // See: https://github.com/Uniswap/v4-core/blob/main/src/libraries/Hooks.sol#L270
570                if before_swap_delta.as_i256() != I256::ZERO {
571                    amount_to_swap += I256::from(before_swap_delta.get_specified_delta());
572                    if amount_to_swap > I256::ZERO {
573                        return Err(SimulationError::FatalError(
574                            "Hook delta exceeds swap amount".into(),
575                        ));
576                    }
577                }
578
579                // Set LP fee override if provided by hook
580                // The fee returned by beforeSwap may have the override flag (bit 22) set,
581                // which needs to be removed before using the fee value.
582                // See: https://github.com/Uniswap/v4-core/blob/main/src/libraries/LPFeeLibrary.sol
583                let hook_fee = before_swap_result
584                    .result
585                    .fee
586                    .to::<u32>();
587                if hook_fee != 0 {
588                    // Remove the override flag (bit 22) as per LPFeeLibrary.sol
589                    let cleaned_fee = lp_fee::remove_override_flag(hook_fee);
590
591                    // Validate the fee doesn't exceed MAX_LP_FEE (1,000,000 pips = 100%)
592                    if !lp_fee::is_valid(cleaned_fee) {
593                        return Err(SimulationError::FatalError(format!(
594                            "LP fee override {} exceeds maximum {} pips",
595                            cleaned_fee,
596                            lp_fee::MAX_LP_FEE
597                        )));
598                    }
599
600                    lp_fee_override = Some(cleaned_fee);
601                }
602            }
603        }
604
605        // Perform the swap with potential hook modifications
606        let result = self.swap(zero_for_one, amount_to_swap, None, lp_fee_override)?;
607
608        // Create BalanceDelta from swap result using the proper constructor
609        let mut swap_delta = BalanceDelta::from_swap_result(result.amount_calculated, zero_for_one);
610
611        // Get deltas (change in the specified/given and unspecified/computed token balances after
612        // calling before swap)
613        let hook_delta_specified = before_swap_delta.get_specified_delta();
614        let mut hook_delta_unspecified = before_swap_delta.get_unspecified_delta();
615
616        if let Some(ref hook) = self.hook {
617            if has_permission(hook.address(), HookOptions::AfterSwap) {
618                let after_swap_params = AfterSwapParameters {
619                    context: state_context,
620                    sender: *EXTERNAL_ACCOUNT,
621                    swap_params,
622                    delta: swap_delta,
623                    hook_data: Bytes::new(),
624                };
625
626                let after_swap_result = hook
627                    .after_swap(after_swap_params, storage_overwrites, None)
628                    .map_err(|e| {
629                        SimulationError::FatalError(format!(
630                            "AfterSwap hook simulation failed: {e:?}"
631                        ))
632                    })?;
633                after_swap_gas = after_swap_result.gas_estimate;
634                hook_delta_unspecified += after_swap_result.result;
635            }
636        }
637
638        // Replicates the behaviour of the Hooks library wrapper of the afterSwap method:
639        // https://github.com/Uniswap/v4-core/blob/59d3ecf53afa9264a16bba0e38f4c5d2231f80bc/src/libraries/Hooks.sol
640        if (hook_delta_specified != I128::ZERO) || (hook_delta_unspecified != I128::ZERO) {
641            let hook_delta = if (amount_specified < I256::ZERO) == zero_for_one {
642                BalanceDelta::new(hook_delta_specified, hook_delta_unspecified)
643            } else {
644                BalanceDelta::new(hook_delta_unspecified, hook_delta_specified)
645            };
646            // This is a BalanceDelta subtraction
647            swap_delta = swap_delta - hook_delta
648        }
649
650        let amount_out = if (amount_specified < I256::ZERO) == zero_for_one {
651            swap_delta.amount1()
652        } else {
653            swap_delta.amount0()
654        };
655
656        trace!(?amount_in, ?token_in, ?token_out, ?zero_for_one, ?result, "V4 SWAP");
657        let mut new_state = self.clone();
658        new_state.liquidity = result.liquidity;
659        new_state.tick = result.tick;
660        new_state.sqrt_price = result.sqrt_price;
661
662        // Add hook gas costs to baseline swap cost
663        let total_gas_used = result.gas_used + U256::from(before_swap_gas + after_swap_gas);
664        Ok(GetAmountOutResult::new(
665            u256_to_biguint(U256::from(amount_out.abs())),
666            u256_to_biguint(total_gas_used),
667            Box::new(new_state),
668        ))
669    }
670
671    fn get_limits(
672        &self,
673        token_in: Bytes,
674        token_out: Bytes,
675    ) -> Result<(BigUint, BigUint), SimulationError> {
676        if let Some(hook) = &self.hook {
677            // Check if pool has no liquidity & ticks -> hook manages liquidity
678            if self.liquidity == 0 && self.has_no_initialized_ticks() {
679                // If the hook has a get_amount_ranges entrypoint, call it and return (0, limits[1])
680                match hook.get_amount_ranges(token_in.clone(), token_out.clone()) {
681                    Ok(amount_ranges) => {
682                        return Ok((
683                            u256_to_biguint(amount_ranges.amount_in_range.1),
684                            u256_to_biguint(amount_ranges.amount_out_range.1),
685                        ))
686                    }
687                    // Check if hook get_amount_ranges is not implemented or the limits entrypoint
688                    // is not set for this hook
689                    Err(SimulationError::RecoverableError(msg))
690                        if msg.contains("not implemented") || msg.contains("not set") =>
691                    {
692                        // Hook manages liquidity but doesn't have get_amount_ranges
693                        // Use binary search to find limits by calling swap with increasing amounts
694                        return self.find_limits_experimentally(token_in, token_out);
695                        // Otherwise fall back to default implementation
696                    }
697                    Err(e) => return Err(e),
698                }
699            }
700        }
701
702        // If the pool has no liquidity, return zeros for both limits
703        if self.liquidity == 0 {
704            return Ok((BigUint::zero(), BigUint::zero()));
705        }
706
707        let zero_for_one = token_in < token_out;
708        let mut current_tick = self.tick;
709        let mut current_sqrt_price = self.sqrt_price;
710        let mut current_liquidity = self.liquidity;
711        let mut total_amount_in = U256::ZERO;
712        let mut total_amount_out = U256::ZERO;
713        let mut ticks_crossed: u64 = 0;
714
715        // Iterate through ticks in the direction of the swap
716        // Stops when: no more liquidity, no more ticks, or gas limit would be exceeded
717        while let Ok((tick, initialized)) = self
718            .ticks
719            .next_initialized_tick_within_one_word(current_tick, zero_for_one)
720        {
721            // Cap iteration to prevent exceeding Ethereum's gas limit
722            if ticks_crossed >= MAX_TICKS_CROSSED {
723                break;
724            }
725            ticks_crossed += 1;
726
727            // Clamp the tick value to ensure it's within valid range
728            let next_tick = tick.clamp(MIN_TICK, MAX_TICK);
729
730            // Calculate the sqrt price at the next tick boundary
731            let sqrt_price_next = get_sqrt_ratio_at_tick(next_tick)?;
732
733            // Calculate the amount of tokens swapped when moving from current_sqrt_price to
734            // sqrt_price_next. Direction determines which token is being swapped in vs out
735            let (amount_in, amount_out) = if zero_for_one {
736                let amount0 = get_amount0_delta(
737                    sqrt_price_next,
738                    current_sqrt_price,
739                    current_liquidity,
740                    true,
741                )?;
742                let amount1 = get_amount1_delta(
743                    sqrt_price_next,
744                    current_sqrt_price,
745                    current_liquidity,
746                    false,
747                )?;
748                (amount0, amount1)
749            } else {
750                let amount0 = get_amount0_delta(
751                    sqrt_price_next,
752                    current_sqrt_price,
753                    current_liquidity,
754                    false,
755                )?;
756                let amount1 = get_amount1_delta(
757                    sqrt_price_next,
758                    current_sqrt_price,
759                    current_liquidity,
760                    true,
761                )?;
762                (amount1, amount0)
763            };
764
765            // Accumulate total amounts for this tick range
766            total_amount_in = safe_add_u256(total_amount_in, amount_in)?;
767            total_amount_out = safe_add_u256(total_amount_out, amount_out)?;
768
769            // If this tick is "initialized" (meaning its someone's position boundary), update the
770            // liquidity when crossing it
771            // For zero_for_one, liquidity is removed when crossing a tick
772            // For one_for_zero, liquidity is added when crossing a tick
773            if initialized {
774                let liquidity_raw = self
775                    .ticks
776                    .get_tick(next_tick)
777                    .unwrap()
778                    .net_liquidity;
779                let liquidity_delta = if zero_for_one { -liquidity_raw } else { liquidity_raw };
780
781                // Check if applying this liquidity delta would cause underflow
782                // If so, stop here rather than continuing with invalid state
783                match liquidity_math::add_liquidity_delta(current_liquidity, liquidity_delta) {
784                    Ok(new_liquidity) => {
785                        current_liquidity = new_liquidity;
786                    }
787                    Err(_) => {
788                        // Liquidity would underflow, stop iteration here
789                        // This represents the maximum liquidity we can actually use
790                        break;
791                    }
792                }
793            }
794
795            // Move to the next tick position
796            current_tick = if zero_for_one { next_tick - 1 } else { next_tick };
797            current_sqrt_price = sqrt_price_next;
798
799            // If we've consumed all liquidity, no point continuing the loop
800            if current_liquidity == 0 {
801                break;
802            }
803        }
804
805        Ok((u256_to_biguint(total_amount_in), u256_to_biguint(total_amount_out)))
806    }
807
808    fn delta_transition(
809        &mut self,
810        delta: ProtocolStateDelta,
811        tokens: &HashMap<Bytes, Token>,
812        balances: &Balances,
813    ) -> Result<(), TransitionError> {
814        if let Some(mut hook) = self.hook.clone() {
815            match hook.delta_transition(delta.clone(), tokens, balances) {
816                Ok(()) => self.set_hook_handler(hook),
817                Err(TransitionError::SimulationError(SimulationError::RecoverableError(msg)))
818                    if msg.contains("not implemented") =>
819                {
820                    // Fall back to default implementation
821                }
822                Err(e) => return Err(e),
823            }
824        }
825
826        // Apply attribute changes
827        if let Some(liquidity) = delta
828            .updated_attributes
829            .get("liquidity")
830        {
831            self.liquidity = u128::from(liquidity.clone());
832        }
833        if let Some(sqrt_price) = delta
834            .updated_attributes
835            .get("sqrt_price_x96")
836        {
837            self.sqrt_price = U256::from_be_slice(sqrt_price);
838        }
839        if let Some(tick) = delta.updated_attributes.get("tick") {
840            self.tick = i24_be_bytes_to_i32(tick);
841        }
842        if let Some(lp_fee) = delta.updated_attributes.get("fee") {
843            self.fees.lp_fee = u32::from(lp_fee.clone());
844        }
845        if let Some(zero2one_protocol_fee) = delta
846            .updated_attributes
847            .get("protocol_fees/zero2one")
848        {
849            self.fees.zero_for_one = u32::from(zero2one_protocol_fee.clone());
850        }
851        if let Some(one2zero_protocol_fee) = delta
852            .updated_attributes
853            .get("protocol_fees/one2zero")
854        {
855            self.fees.one_for_zero = u32::from(one2zero_protocol_fee.clone());
856        }
857
858        // apply tick changes
859        for (key, value) in delta.updated_attributes.iter() {
860            // tick liquidity keys are in the format "ticks/{tick_index}/net_liquidity"
861            if key.starts_with("ticks/") {
862                let parts: Vec<&str> = key.split('/').collect();
863                self.ticks
864                    .set_tick_liquidity(
865                        parts[1]
866                            .parse::<i32>()
867                            .map_err(|err| TransitionError::DecodeError(err.to_string()))?,
868                        i128::from(value.clone()),
869                    )
870                    .map_err(|err| TransitionError::DecodeError(err.to_string()))?;
871            }
872        }
873        // delete ticks - ignores deletes for attributes other than tick liquidity
874        for key in delta.deleted_attributes.iter() {
875            // tick liquidity keys are in the format "ticks/{tick_index}/net_liquidity"
876            if key.starts_with("ticks/") {
877                let parts: Vec<&str> = key.split('/').collect();
878                self.ticks
879                    .set_tick_liquidity(
880                        parts[1]
881                            .parse::<i32>()
882                            .map_err(|err| TransitionError::DecodeError(err.to_string()))?,
883                        0,
884                    )
885                    .map_err(|err| TransitionError::DecodeError(err.to_string()))?;
886            }
887        }
888
889        Ok(())
890    }
891
892    /// See [`ProtocolSim::query_pool_swap`] for the trait documentation.
893    ///
894    /// This method uses Uniswap V4 internal swap logic by swapping an infinite amount of token_in
895    /// until the target price is reached. Takes into account V4-specific features like protocol
896    /// fees and dynamic LP fees.
897    ///
898    /// Note: This implementation does not invoke hooks, as it is a query-only operation meant to
899    /// determine available liquidity at a given price without executing an actual swap.
900    fn query_pool_swap(&self, params: &QueryPoolSwapParams) -> Result<PoolSwap, SimulationError> {
901        if self.liquidity == 0 {
902            return Err(SimulationError::FatalError("No liquidity".to_string()));
903        }
904
905        // Calculate total fee (protocol + LP fee) for V4
906        let zero_for_one = params.token_in().address < params.token_out().address;
907        let fee_pips = self
908            .fees
909            .calculate_swap_fees_pips(zero_for_one, None);
910
911        match params.swap_constraint() {
912            SwapConstraint::TradeLimitPrice { .. } => Err(SimulationError::InvalidInput(
913                "Uniswap V4 does not support TradeLimitPrice constraint in query_pool_swap"
914                    .to_string(),
915                None,
916            )),
917            SwapConstraint::PoolTargetPrice {
918                target,
919                tolerance: _,
920                min_amount_in: _,
921                max_amount_in: _,
922            } => {
923                if self.liquidity == 0 {
924                    return Err(SimulationError::FatalError("No liquidity".to_string()));
925                }
926
927                let (amount_in, amount_out, swap_result) = clmm_swap_to_price(
928                    self.sqrt_price,
929                    &params.token_in().address,
930                    &params.token_out().address,
931                    target,
932                    fee_pips,
933                    Sign::Negative, // V4 uses negative for exact input
934                    |zero_for_one, amount_specified, sqrt_price_limit| {
935                        self.swap(zero_for_one, amount_specified, Some(sqrt_price_limit), None)
936                    },
937                )?;
938
939                let mut new_state = self.clone();
940                new_state.liquidity = swap_result.liquidity;
941                new_state.tick = swap_result.tick;
942                new_state.sqrt_price = swap_result.sqrt_price;
943
944                Ok(PoolSwap::new(amount_in, amount_out, Box::new(new_state), None))
945            }
946        }
947    }
948
949    fn clone_box(&self) -> Box<dyn ProtocolSim> {
950        Box::new(self.clone())
951    }
952
953    fn as_any(&self) -> &dyn Any {
954        self
955    }
956
957    fn as_any_mut(&mut self) -> &mut dyn Any {
958        self
959    }
960
961    fn eq(&self, other: &dyn ProtocolSim) -> bool {
962        if let Some(other_state) = other
963            .as_any()
964            .downcast_ref::<UniswapV4State>()
965        {
966            self.liquidity == other_state.liquidity &&
967                self.sqrt_price == other_state.sqrt_price &&
968                self.fees == other_state.fees &&
969                self.tick == other_state.tick &&
970                self.ticks == other_state.ticks
971        } else {
972            false
973        }
974    }
975}
976
977#[cfg(test)]
978mod tests {
979    use std::{collections::HashSet, fs, path::Path, str::FromStr};
980
981    use alloy::primitives::aliases::U24;
982    use num_traits::FromPrimitive;
983    use rstest::rstest;
984    use serde_json::Value;
985    use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader};
986    use tycho_common::{models::Chain, simulation::protocol_sim::Price};
987
988    use super::*;
989    use crate::{
990        evm::{
991            engine_db::{
992                create_engine,
993                simulation_db::SimulationDB,
994                utils::{get_client, get_runtime},
995            },
996            protocol::{
997                uniswap_v4::hooks::{
998                    angstrom::hook_handler::{AngstromFees, AngstromHookHandler},
999                    generic_vm_hook_handler::GenericVMHookHandler,
1000                },
1001                utils::uniswap::{lp_fee, sqrt_price_math::get_sqrt_price_q96},
1002            },
1003        },
1004        protocol::models::{DecoderContext, TryFromWithBlock},
1005    };
1006
1007    // Helper methods to create commonly used tokens
1008    fn usdc() -> Token {
1009        Token::new(
1010            &Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(),
1011            "USDC",
1012            6,
1013            0,
1014            &[Some(10_000)],
1015            Default::default(),
1016            100,
1017        )
1018    }
1019
1020    fn weth() -> Token {
1021        Token::new(
1022            &Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(),
1023            "WETH",
1024            18,
1025            0,
1026            &[Some(10_000)],
1027            Default::default(),
1028            100,
1029        )
1030    }
1031
1032    fn eth() -> Token {
1033        Token::new(
1034            &Bytes::from_str("0x0000000000000000000000000000000000000000").unwrap(),
1035            "ETH",
1036            18,
1037            0,
1038            &[Some(10_000)],
1039            Default::default(),
1040            100,
1041        )
1042    }
1043
1044    fn token_x() -> Token {
1045        Token::new(
1046            &Bytes::from_str("0x0000000000000000000000000000000000000001").unwrap(),
1047            "X",
1048            18,
1049            0,
1050            &[Some(10_000)],
1051            Default::default(),
1052            100,
1053        )
1054    }
1055
1056    fn token_y() -> Token {
1057        Token::new(
1058            &Bytes::from_str("0x0000000000000000000000000000000000000002").unwrap(),
1059            "Y",
1060            18,
1061            0,
1062            &[Some(10_000)],
1063            Default::default(),
1064            100,
1065        )
1066    }
1067
1068    #[test]
1069    fn test_delta_transition() {
1070        let mut pool = UniswapV4State::new(
1071            1000,
1072            U256::from_str("1000").unwrap(),
1073            UniswapV4Fees { zero_for_one: 100, one_for_zero: 90, lp_fee: 700 },
1074            100,
1075            60,
1076            vec![TickInfo::new(120, 10000).unwrap(), TickInfo::new(180, -10000).unwrap()],
1077        )
1078        .unwrap();
1079
1080        let attributes: HashMap<String, Bytes> = [
1081            ("liquidity".to_string(), Bytes::from(2000_u64.to_be_bytes().to_vec())),
1082            ("sqrt_price_x96".to_string(), Bytes::from(1001_u64.to_be_bytes().to_vec())),
1083            ("tick".to_string(), Bytes::from(120_i32.to_be_bytes().to_vec())),
1084            ("protocol_fees/zero2one".to_string(), Bytes::from(50_u32.to_be_bytes().to_vec())),
1085            ("protocol_fees/one2zero".to_string(), Bytes::from(75_u32.to_be_bytes().to_vec())),
1086            ("fee".to_string(), Bytes::from(100_u32.to_be_bytes().to_vec())),
1087            ("ticks/-120/net_liquidity".to_string(), Bytes::from(10200_u64.to_be_bytes().to_vec())),
1088            ("ticks/120/net_liquidity".to_string(), Bytes::from(9800_u64.to_be_bytes().to_vec())),
1089            ("block_number".to_string(), Bytes::from(2000_u64.to_be_bytes().to_vec())),
1090            ("block_timestamp".to_string(), Bytes::from(1758201935_u64.to_be_bytes().to_vec())),
1091        ]
1092        .into_iter()
1093        .collect();
1094
1095        let delta = ProtocolStateDelta {
1096            component_id: "State1".to_owned(),
1097            updated_attributes: attributes,
1098            deleted_attributes: HashSet::new(),
1099        };
1100
1101        pool.delta_transition(delta, &HashMap::new(), &Balances::default())
1102            .unwrap();
1103
1104        assert_eq!(pool.liquidity, 2000);
1105        assert_eq!(pool.sqrt_price, U256::from(1001));
1106        assert_eq!(pool.tick, 120);
1107        assert_eq!(pool.fees.zero_for_one, 50);
1108        assert_eq!(pool.fees.one_for_zero, 75);
1109        assert_eq!(pool.fees.lp_fee, 100);
1110        assert_eq!(
1111            pool.ticks
1112                .get_tick(-120)
1113                .unwrap()
1114                .net_liquidity,
1115            10200
1116        );
1117        assert_eq!(
1118            pool.ticks
1119                .get_tick(120)
1120                .unwrap()
1121                .net_liquidity,
1122            9800
1123        );
1124    }
1125
1126    #[tokio::test]
1127    /// Compares a quote from the UniswapV4 Quoter contract on Sepolia with a simulation.
1128    async fn test_swap_sim() {
1129        use tycho_client::feed::dto;
1130        let project_root = env!("CARGO_MANIFEST_DIR");
1131        let asset_path = Path::new(project_root)
1132            .join("tests/assets/decoder/uniswap_v4_snapshot_sepolia_block_7239119.json");
1133        let json_data = fs::read_to_string(asset_path).expect("Failed to read test asset");
1134        let data: Value = serde_json::from_str(&json_data).expect("Failed to parse JSON");
1135        let state: ComponentWithState = serde_json::from_value::<dto::ComponentWithState>(data)
1136            .expect("Expected json to match ComponentWithState structure")
1137            .into();
1138
1139        let block = BlockHeader {
1140            number: 7239119,
1141            hash: Bytes::from_str(
1142                "0x28d41d40f2ac275a4f5f621a636b9016b527d11d37d610a45ac3a821346ebf8c",
1143            )
1144            .expect("Invalid block hash"),
1145            parent_hash: Bytes::from(vec![0; 32]),
1146            ..Default::default()
1147        };
1148
1149        let t0 = Token::new(
1150            &Bytes::from_str("0x647e32181a64f4ffd4f0b0b4b052ec05b277729c").unwrap(),
1151            "T0",
1152            18,
1153            0,
1154            &[Some(10_000)],
1155            Chain::Ethereum,
1156            100,
1157        );
1158        let t1 = Token::new(
1159            &Bytes::from_str("0xe390a1c311b26f14ed0d55d3b0261c2320d15ca5").unwrap(),
1160            "T1",
1161            18,
1162            0,
1163            &[Some(10_000)],
1164            Chain::Ethereum,
1165            100,
1166        );
1167
1168        let all_tokens = [t0.clone(), t1.clone()]
1169            .iter()
1170            .map(|t| (t.address.clone(), t.clone()))
1171            .collect();
1172
1173        let usv4_state = UniswapV4State::try_from_with_header(
1174            state,
1175            block,
1176            &Default::default(),
1177            &all_tokens,
1178            &DecoderContext::new(),
1179        )
1180        .await
1181        .unwrap();
1182
1183        let res = usv4_state
1184            .get_amount_out(BigUint::from_u64(1000000000000000000).unwrap(), &t0, &t1)
1185            .unwrap();
1186
1187        let expected_amount = BigUint::from(9999909699895_u64);
1188        assert_eq!(res.amount, expected_amount);
1189    }
1190
1191    #[tokio::test]
1192    async fn test_get_limits() {
1193        use tycho_client::feed::dto;
1194        let block = BlockHeader {
1195            number: 22689129,
1196            hash: Bytes::from_str(
1197                "0x7763ea30d11aef68da729b65250c09a88ad00458c041064aad8c9a9dbf17adde",
1198            )
1199            .expect("Invalid block hash"),
1200            parent_hash: Bytes::from(vec![0; 32]),
1201            ..Default::default()
1202        };
1203
1204        let project_root = env!("CARGO_MANIFEST_DIR");
1205        let asset_path =
1206            Path::new(project_root).join("tests/assets/decoder/uniswap_v4_snapshot.json");
1207        let json_data = fs::read_to_string(asset_path).expect("Failed to read test asset");
1208        let data: Value = serde_json::from_str(&json_data).expect("Failed to parse JSON");
1209        let state: ComponentWithState = serde_json::from_value::<dto::ComponentWithState>(data)
1210            .expect("Expected json to match ComponentWithState structure")
1211            .into();
1212
1213        let t0 = Token::new(
1214            &Bytes::from_str("0x2260fac5e5542a773aa44fbcfedf7c193bc2c599").unwrap(),
1215            "WBTC",
1216            8,
1217            0,
1218            &[Some(10_000)],
1219            Chain::Ethereum,
1220            100,
1221        );
1222        let t1 = Token::new(
1223            &Bytes::from_str("0xdac17f958d2ee523a2206206994597c13d831ec7").unwrap(),
1224            "USDT",
1225            6,
1226            0,
1227            &[Some(10_000)],
1228            Chain::Ethereum,
1229            100,
1230        );
1231
1232        let all_tokens = [t0.clone(), t1.clone()]
1233            .iter()
1234            .map(|t| (t.address.clone(), t.clone()))
1235            .collect();
1236
1237        let usv4_state = UniswapV4State::try_from_with_header(
1238            state,
1239            block,
1240            &Default::default(),
1241            &all_tokens,
1242            &DecoderContext::new(),
1243        )
1244        .await
1245        .unwrap();
1246
1247        let res = usv4_state
1248            .get_limits(t0.address.clone(), t1.address.clone())
1249            .unwrap();
1250
1251        assert_eq!(&res.0, &BigUint::from_u128(71698353688830259750744466706).unwrap());
1252
1253        let out = usv4_state
1254            .get_amount_out(res.0, &t0, &t1)
1255            .expect("swap for limit in didn't work");
1256
1257        assert_eq!(&res.1, &out.amount);
1258    }
1259    #[test]
1260    fn test_get_amount_out_no_hook() {
1261        // Test using transaction 0x78ea4bbb7d4405000f33fdf6f3fa08b5e557d50e5e7f826a79766d50bd643b6f
1262
1263        // Pool ID: 0x00b9edc1583bf6ef09ff3a09f6c23ecb57fd7d0bb75625717ec81eed181e22d7
1264        // Information taken from Tenderly simulation / event emitted on Etherscan
1265        let usv4_state = UniswapV4State::new(
1266            541501951282951892,
1267            U256::from_str("5362798333066270795901222").unwrap(), // Sqrt price
1268            UniswapV4Fees { zero_for_one: 0, one_for_zero: 0, lp_fee: 100 },
1269            -192022,
1270            1,
1271            // Ticks taken from indexer logs
1272            vec![
1273                TickInfo {
1274                    index: -887272,
1275                    net_liquidity: 460382969070005,
1276                    sqrt_price: U256::from(4295128739_u64),
1277                },
1278                TickInfo {
1279                    index: -207244,
1280                    net_liquidity: 561268407024557,
1281                    sqrt_price: U256::from_str("2505291706254206075074035").unwrap(),
1282                },
1283                TickInfo {
1284                    index: -196411,
1285                    net_liquidity: 825711941800452,
1286                    sqrt_price: U256::from_str("4306080513146952705853399").unwrap(),
1287                },
1288                TickInfo {
1289                    index: -196257,
1290                    net_liquidity: 64844666874010,
1291                    sqrt_price: U256::from_str("4339363644587371378270009").unwrap(),
1292                },
1293                TickInfo {
1294                    index: -195611,
1295                    net_liquidity: 2344045150766798,
1296                    sqrt_price: U256::from_str("4481806029599743916020126").unwrap(),
1297                },
1298                TickInfo {
1299                    index: -194715,
1300                    net_liquidity: 391037380558274654,
1301                    sqrt_price: U256::from_str("4687145946111116896040494").unwrap(),
1302                },
1303                TickInfo {
1304                    index: -194599,
1305                    net_liquidity: 89032603464508,
1306                    sqrt_price: U256::from_str("4714409015946702405379370").unwrap(),
1307                },
1308                TickInfo {
1309                    index: -194389,
1310                    net_liquidity: 66635600426483168,
1311                    sqrt_price: U256::from_str("4764168603367683402636621").unwrap(),
1312                },
1313                TickInfo {
1314                    index: -194160,
1315                    net_liquidity: 6123093436523361,
1316                    sqrt_price: U256::from_str("4819029067726467394386780").unwrap(),
1317                },
1318                TickInfo {
1319                    index: -194025,
1320                    net_liquidity: 79940813798964,
1321                    sqrt_price: U256::from_str("4851665907541490407930032").unwrap(),
1322                },
1323                TickInfo {
1324                    index: -193922,
1325                    net_liquidity: 415630967437234,
1326                    sqrt_price: U256::from_str("4876715181040466809166531").unwrap(),
1327                },
1328                TickInfo {
1329                    index: -193876,
1330                    net_liquidity: 9664144015186047,
1331                    sqrt_price: U256::from_str("4887943972687250473582419").unwrap(),
1332                },
1333                TickInfo {
1334                    index: -193818,
1335                    net_liquidity: 435344726052344,
1336                    sqrt_price: U256::from_str("4902138873132735049121973").unwrap(),
1337                },
1338                TickInfo {
1339                    index: -193804,
1340                    net_liquidity: 221726179374067,
1341                    sqrt_price: U256::from_str("4905571399964683340605904").unwrap(),
1342                },
1343                TickInfo {
1344                    index: -193719,
1345                    net_liquidity: 101340835774487,
1346                    sqrt_price: U256::from_str("4926463397882393957462188").unwrap(),
1347                },
1348                TickInfo {
1349                    index: -193690,
1350                    net_liquidity: 193367475630077,
1351                    sqrt_price: U256::from_str("4933611593595025190448924").unwrap(),
1352                },
1353                TickInfo {
1354                    index: -193643,
1355                    net_liquidity: 357016631583746,
1356                    sqrt_price: U256::from_str("4945218633428068823432932").unwrap(),
1357                },
1358                TickInfo {
1359                    index: -193520,
1360                    net_liquidity: 917243184365178,
1361                    sqrt_price: U256::from_str("4975723910367862081017120").unwrap(),
1362                },
1363                TickInfo {
1364                    index: -193440,
1365                    net_liquidity: 114125890211958292,
1366                    sqrt_price: U256::from_str("4995665665861492533686137").unwrap(),
1367                },
1368                TickInfo {
1369                    index: -193380,
1370                    net_liquidity: -65980729148766579,
1371                    sqrt_price: U256::from_str("5010674414300823856025303").unwrap(),
1372                },
1373                TickInfo {
1374                    index: -192891,
1375                    net_liquidity: 1687883551433195,
1376                    sqrt_price: U256::from_str("5134689105039642314202223").unwrap(),
1377                },
1378                TickInfo {
1379                    index: -192573,
1380                    net_liquidity: 11108903221360975,
1381                    sqrt_price: U256::from_str("5216979018647067786855495").unwrap(),
1382                },
1383                TickInfo {
1384                    index: -192448,
1385                    net_liquidity: 32888457482352,
1386                    sqrt_price: U256::from_str("5249685603828944002327927").unwrap(),
1387                },
1388                TickInfo {
1389                    index: -191525,
1390                    net_liquidity: -221726179374067,
1391                    sqrt_price: U256::from_str("5497623359964843320146512").unwrap(),
1392                },
1393                TickInfo {
1394                    index: -191447,
1395                    net_liquidity: -32888457482352,
1396                    sqrt_price: U256::from_str("5519104878745833608097296").unwrap(),
1397                },
1398                TickInfo {
1399                    index: -191444,
1400                    net_liquidity: -114125890211958292,
1401                    sqrt_price: U256::from_str("5519932765173943847315221").unwrap(),
1402                },
1403                TickInfo {
1404                    index: -191417,
1405                    net_liquidity: -101340835774487,
1406                    sqrt_price: U256::from_str("5527389333636021285046380").unwrap(),
1407                },
1408                TickInfo {
1409                    index: -191384,
1410                    net_liquidity: -9664144015186047,
1411                    sqrt_price: U256::from_str("5536516597603056457376182").unwrap(),
1412                },
1413                TickInfo {
1414                    index: -191148,
1415                    net_liquidity: -561268407024557,
1416                    sqrt_price: U256::from_str("5602231161238705865493165").unwrap(),
1417                },
1418                TickInfo {
1419                    index: -191147,
1420                    net_liquidity: -1687883551433195,
1421                    sqrt_price: U256::from_str("5602511265794328966803451").unwrap(),
1422                },
1423                TickInfo {
1424                    index: -191091,
1425                    net_liquidity: -89032603464508,
1426                    sqrt_price: U256::from_str("5618219493196441347292357").unwrap(),
1427                },
1428                TickInfo {
1429                    index: -190950,
1430                    net_liquidity: -189177935487638,
1431                    sqrt_price: U256::from_str("5657965894785859782969011").unwrap(),
1432                },
1433                TickInfo {
1434                    index: -190756,
1435                    net_liquidity: -6123093436523361,
1436                    sqrt_price: U256::from_str("5713112435031881967192022").unwrap(),
1437                },
1438                TickInfo {
1439                    index: -190548,
1440                    net_liquidity: -193367475630077,
1441                    sqrt_price: U256::from_str("5772835841671084402427710").unwrap(),
1442                },
1443                TickInfo {
1444                    index: -190430,
1445                    net_liquidity: -11108903221360975,
1446                    sqrt_price: U256::from_str("5806994534290341208820930").unwrap(),
1447                },
1448                TickInfo {
1449                    index: -190195,
1450                    net_liquidity: -391583014714302569,
1451                    sqrt_price: U256::from_str("5875625707132601785181387").unwrap(),
1452                },
1453                TickInfo {
1454                    index: -190043,
1455                    net_liquidity: -357016631583746,
1456                    sqrt_price: U256::from_str("5920448331650864936739481").unwrap(),
1457                },
1458                TickInfo {
1459                    index: -189779,
1460                    net_liquidity: -917243184365178,
1461                    sqrt_price: U256::from_str("5999112356918485175181346").unwrap(),
1462                },
1463                TickInfo {
1464                    index: -189663,
1465                    net_liquidity: -2344045150766798,
1466                    sqrt_price: U256::from_str("6034006559279282606084981").unwrap(),
1467                },
1468                TickInfo {
1469                    index: -189620,
1470                    net_liquidity: -435344726052344,
1471                    sqrt_price: U256::from_str("6046992979471024289177519").unwrap(),
1472                },
1473                TickInfo {
1474                    index: -189409,
1475                    net_liquidity: -825711941800452,
1476                    sqrt_price: U256::from_str("6111123241285165242130911").unwrap(),
1477                },
1478                TickInfo {
1479                    index: -189325,
1480                    net_liquidity: -3947182209207,
1481                    sqrt_price: U256::from_str("6136842645893819031257990").unwrap(),
1482                },
1483                TickInfo {
1484                    index: -189324,
1485                    net_liquidity: -415630967437234,
1486                    sqrt_price: U256::from_str("6137149480355443943537284").unwrap(),
1487                },
1488                TickInfo {
1489                    index: -115136,
1490                    net_liquidity: 462452451821,
1491                    sqrt_price: U256::from_str("250529060232794967902094762").unwrap(),
1492                },
1493                TickInfo {
1494                    index: -92109,
1495                    net_liquidity: -462452451821,
1496                    sqrt_price: U256::from_str("792242363124136400178523925").unwrap(),
1497                },
1498                TickInfo {
1499                    index: 887272,
1500                    net_liquidity: -521280453734808,
1501                    sqrt_price: U256::from_str("1461446703485210103287273052203988822378723970342")
1502                        .unwrap(),
1503                },
1504            ],
1505        )
1506        .unwrap();
1507
1508        let t0 = usdc();
1509        let t1 = eth();
1510
1511        let out = usv4_state
1512            .get_amount_out(BigUint::from_u64(2000000).unwrap(), &t0, &t1)
1513            .unwrap();
1514
1515        assert_eq!(out.amount, BigUint::from_str("436478419853848").unwrap())
1516    }
1517
1518    #[test]
1519    fn test_get_amount_out_euler_hook() {
1520        // Test using transaction 0xb372306a81c6e840f4ec55f006da6b0b097f435802a2e6fd216998dd12fb4aca
1521        //
1522        // Output of beforeSwap:
1523        // "output":{
1524        //      "amountToSwap":"0"
1525        //      "hookReturn":"2520471492123673565794154180707800634502860978735"
1526        //      "lpFeeOverride":"0"
1527        // }
1528        //
1529        // Output of entire swap, including hooks:
1530        // "swapDelta":"-2520471491783391198873215717244426027071092767279"
1531        //
1532        // Get amount out:
1533        // "amountOut":"2681115183499232721"
1534
1535        let block = BlockHeader {
1536            number: 22689128,
1537            hash: Bytes::from_str(
1538                "0xfbfa716523d25d6d5248c18d001ca02b1caf10cabd1ab7321465e2262c41157b",
1539            )
1540            .expect("Invalid block hash"),
1541            timestamp: 1749739055,
1542            ..Default::default()
1543        };
1544
1545        // Pool ID: 0xdd8dd509e58ec98631b800dd6ba86ee569c517ffbd615853ed5ab815bbc48ccb
1546        // Information taken from Tenderly simulation
1547        let mut usv4_state = UniswapV4State::new(
1548            0,
1549            U256::from_str("4295128740").unwrap(),
1550            UniswapV4Fees { zero_for_one: 100, one_for_zero: 90, lp_fee: 500 },
1551            0,
1552            1,
1553            vec![],
1554        )
1555        .unwrap();
1556
1557        let hook_address: Address = Address::from_str("0x69058613588536167ba0aa94f0cc1fe420ef28a8")
1558            .expect("Invalid hook address");
1559
1560        let db = SimulationDB::new(
1561            get_client(None).expect("Failed to create client"),
1562            get_runtime().expect("Failed to get runtime"),
1563            Some(block.clone()),
1564        );
1565        let engine = create_engine(db, true).expect("Failed to create simulation engine");
1566        let pool_manager = Address::from_str("0x000000000004444c5dc75cb358380d2e3de08a90")
1567            .expect("Invalid pool manager address");
1568
1569        let hook_handler = GenericVMHookHandler::new(
1570            hook_address,
1571            engine,
1572            pool_manager,
1573            HashMap::new(),
1574            HashMap::new(),
1575            None,
1576            true, // Euler hook
1577        )
1578        .unwrap();
1579
1580        let t0 = usdc();
1581        let t1 = weth();
1582
1583        usv4_state.set_hook_handler(Box::new(hook_handler));
1584        let out = usv4_state
1585            .get_amount_out(BigUint::from_u64(7407000000).unwrap(), &t0, &t1)
1586            .unwrap();
1587
1588        assert_eq!(out.amount, BigUint::from_str("2681115183499232721").unwrap())
1589    }
1590
1591    #[test]
1592    fn test_get_amount_out_angstrom_hook() {
1593        // Test using transaction 0x671b8e1d0966cee520dc2bb9628de8e22a17b036e70077504796d0a476932d21
1594        let mut usv4_state = UniswapV4State::new(
1595            // Liquidity and tick taken from tycho indexer for same block as transaction
1596            66319800403673162,
1597            U256::from_str("1314588940601923011323000261788004").unwrap(),
1598            // 8388608 (i.e. 0x800000) signifies a dynamic fee.
1599            UniswapV4Fees { zero_for_one: 0, one_for_zero: 0, lp_fee: 8388608 },
1600            194343,
1601            10,
1602            vec![
1603                TickInfo::new(-887270, 198117767801).unwrap(),
1604                TickInfo::new(191990, 24561988698695).unwrap(),
1605                TickInfo::new(192280, 2839631428751224).unwrap(),
1606                TickInfo::new(193130, 318786492813931).unwrap(),
1607                TickInfo::new(194010, 26209207141081).unwrap(),
1608                TickInfo::new(194210, -26209207141081).unwrap(),
1609                TickInfo::new(194220, 63136622375641511).unwrap(),
1610                TickInfo::new(194420, -63136622375641511).unwrap(),
1611                TickInfo::new(195130, -318786492813931).unwrap(),
1612                TickInfo::new(196330, -2839631428751224).unwrap(),
1613                TickInfo::new(197100, -24561988698695).unwrap(),
1614                TickInfo::new(887270, -198117767801).unwrap(),
1615            ],
1616        )
1617        .unwrap();
1618
1619        let fees = AngstromFees {
1620            // To get these values, enable storage access logs on tenderly,
1621            // and look at the hex value retrieved right after calling afterSwap
1622            //
1623            // The value (hex: 0x70000152) contains two packed uint24 values:
1624            // Lower 24 bits (unlockedFee):         0x152   = 338
1625            // Upper 24 bits (protocolUnlockedFee): 0x70    = 112
1626            unlock: U24::from(338),
1627            protocol_unlock: U24::from(112),
1628        };
1629        let hook_handler = AngstromHookHandler::new(
1630            Address::from_str("0x0000000aa232009084bd71a5797d089aa4edfad4").unwrap(),
1631            Address::from_str("0x000000000004444c5dc75cb358380d2e3de08a90").unwrap(),
1632            fees,
1633            false,
1634        );
1635
1636        let t0 = usdc();
1637        let t1 = weth();
1638
1639        usv4_state.set_hook_handler(Box::new(hook_handler));
1640        let out = usv4_state
1641            .get_amount_out(
1642                BigUint::from_u64(
1643                    6645198144, // usdc
1644                )
1645                .unwrap(),
1646                &t0, // usdc IN
1647                &t1, // weth OUT
1648            )
1649            .unwrap();
1650
1651        assert_eq!(out.amount, BigUint::from_str("1825627051870330472").unwrap())
1652    }
1653
1654    #[test]
1655    fn test_spot_price_with_recoverable_error() {
1656        // Test that spot_price correctly falls back to swap-based calculation
1657        // when a RecoverableError (other than "not implemented") is returned
1658
1659        let usv4_state = UniswapV4State::new(
1660            1000000000000000000u128,                                  // 1e18 liquidity
1661            U256::from_str("79228162514264337593543950336").unwrap(), // 1:1 price
1662            UniswapV4Fees { zero_for_one: 100, one_for_zero: 100, lp_fee: 100 },
1663            0,
1664            60,
1665            vec![
1666                TickInfo::new(-600, 500000000000000000i128).unwrap(),
1667                TickInfo::new(600, -500000000000000000i128).unwrap(),
1668            ],
1669        )
1670        .unwrap();
1671
1672        // Test spot price calculation without a hook (should use default implementation)
1673        let spot_price_result = usv4_state.spot_price(&usdc(), &weth());
1674        assert!(spot_price_result.is_ok());
1675
1676        // The price should be approximately 1.0 (since we set sqrt_price for 1:1)
1677        // Adjusting for decimals difference (USDC has 6, WETH has 18)
1678        let price = spot_price_result.unwrap();
1679        assert!(price > 0.0);
1680    }
1681
1682    #[test]
1683    fn test_get_limits_with_hook_managed_liquidity_no_ranges_entrypoint() {
1684        // This test demonstrates the experimental limit finding logic for hooks that:
1685        // 1. Manage liquidity (pool has no liquidity & no ticks)
1686        // 2. Don't have get_amount_ranges entrypoint
1687
1688        let block = BlockHeader {
1689            number: 22689128,
1690            hash: Bytes::from_str(
1691                "0xfbfa716523d25d6d5248c18d001ca02b1caf10cabd1ab7321465e2262c41157b",
1692            )
1693            .expect("Invalid block hash"),
1694            timestamp: 1749739055,
1695            ..Default::default()
1696        };
1697
1698        let hook_address: Address = Address::from_str("0x69058613588536167ba0aa94f0cc1fe420ef28a8")
1699            .expect("Invalid hook address");
1700
1701        let db = SimulationDB::new(
1702            get_client(None).expect("Failed to create client"),
1703            get_runtime().expect("Failed to get runtime"),
1704            Some(block.clone()),
1705        );
1706        let engine = create_engine(db, true).expect("Failed to create simulation engine");
1707        let pool_manager = Address::from_str("0x000000000004444c5dc75cb358380d2e3de08a90")
1708            .expect("Invalid pool manager address");
1709
1710        // Create a GenericVMHookHandler without limits_entrypoint
1711        // This will trigger the "not set" error path and use experimental limit finding
1712        let hook_handler = GenericVMHookHandler::new(
1713            hook_address,
1714            engine,
1715            pool_manager,
1716            HashMap::new(),
1717            HashMap::new(),
1718            None,
1719            true, // Euler hook
1720        )
1721        .unwrap();
1722
1723        // Create a UniswapV4State with NO liquidity and NO ticks (hook manages all liquidity)
1724        let mut usv4_state = UniswapV4State::new(
1725            0, // no liquidity - hook provides it
1726            U256::from_str("4295128740").unwrap(),
1727            UniswapV4Fees { zero_for_one: 100, one_for_zero: 90, lp_fee: 500 },
1728            0,      // current tick
1729            1,      // tick spacing
1730            vec![], // no ticks - hook manages liquidity
1731        )
1732        .unwrap();
1733
1734        usv4_state.set_hook_handler(Box::new(hook_handler));
1735
1736        let token_in = usdc().address;
1737        let token_out = weth().address;
1738
1739        let (amount_in_limit, amount_out_limit) = usv4_state
1740            .get_limits(token_in, token_out)
1741            .expect("Should find limits through experimental swapping");
1742
1743        // Assuming pool supply doesn't change drastically at time of this test
1744        // At least 1 million USDC, not more than 100 million USDC
1745        assert!(amount_in_limit > BigUint::from(10u64).pow(12));
1746        assert!(amount_in_limit < BigUint::from(10u64).pow(14));
1747
1748        // At least 100 ETH, not more than 10 000 ETH
1749        assert!(amount_out_limit > BigUint::from(10u64).pow(20));
1750        assert!(amount_out_limit < BigUint::from(10u64).pow(22));
1751    }
1752
1753    #[rstest]
1754    #[case::high_liquidity(u128::MAX / 2)] // Very large liquidity
1755    #[case::medium_liquidity(10000000000000000000u128)] // Moderate liquidity: 10e18
1756    #[case::minimal_liquidity(1000u128)] // Very small liquidity
1757    fn test_find_max_amount(#[case] liquidity: u128) {
1758        // Use fixed configuration for all test cases
1759        let fees = UniswapV4Fees { zero_for_one: 100, one_for_zero: 100, lp_fee: 100 };
1760        let tick_spacing = 60;
1761        let ticks = vec![
1762            TickInfo::new(-600, (liquidity / 4) as i128).unwrap(),
1763            TickInfo::new(600, -((liquidity / 4) as i128)).unwrap(),
1764        ];
1765
1766        let usv4_state = UniswapV4State::new(
1767            liquidity,
1768            U256::from_str("79228162514264337593543950336").unwrap(),
1769            fees,
1770            0,
1771            tick_spacing,
1772            ticks,
1773        )
1774        .unwrap();
1775
1776        let token_in = usdc();
1777        let token_out = weth();
1778
1779        let (max_amount_in, _max_amount_out) = usv4_state
1780            .find_max_amount(&token_in, &token_out)
1781            .unwrap();
1782
1783        let success = usv4_state
1784            .get_amount_out(max_amount_in.clone(), &token_in, &token_out)
1785            .is_ok();
1786        assert!(success, "Should be able to swap the exact max amount.");
1787
1788        let one_more = &max_amount_in + BigUint::from(1u64);
1789        let should_fail = usv4_state
1790            .get_amount_out(one_more, &token_in, &token_out)
1791            .is_err();
1792        assert!(should_fail, "Swapping max_amount + 1 should fail.");
1793    }
1794
1795    #[test]
1796    fn test_calculate_swap_fees_with_override() {
1797        // Test that calculate_swap_fees_pips works correctly with overridden fees
1798        let fees = UniswapV4Fees::new(100, 90, 500);
1799
1800        // Without override, should use UniswapV4 formula: protocol + lp - (protocol * lp /
1801        // 1_000_000)
1802        let total_zero_for_one = fees.calculate_swap_fees_pips(true, None);
1803        // 100 + 500 - (100 * 500 / 1_000_000) = 600 - 0 = 600 (rounded down)
1804        assert_eq!(total_zero_for_one, 600);
1805
1806        // With override, should use override fee + protocol fee with same formula
1807        let total_with_override = fees.calculate_swap_fees_pips(true, Some(1000));
1808        // 100 + 1000 - (100 * 1000 / 1_000_000) = 1100 - 0 = 1100 (rounded down)
1809        assert_eq!(total_with_override, 1100);
1810    }
1811
1812    #[test]
1813    fn test_max_combined_fees_stays_valid() {
1814        // Test that even with max protocol + max LP fees, we stay under compute_swap_step limit
1815        let fees = UniswapV4Fees::new(1000, 1000, 1000);
1816        let total = fees.calculate_swap_fees_pips(true, Some(lp_fee::MAX_LP_FEE));
1817
1818        // Using UniswapV4 formula: 1000 + 1000000 - (1000 * 1000000 / 1_000_000)
1819        // = 1001000 - 1000 = 1000000
1820        assert_eq!(total, 1_000_000);
1821    }
1822
1823    #[test]
1824    fn test_get_limits_graceful_underflow() {
1825        // Verifies graceful handling of liquidity underflow in get_limits for V4
1826        let usv4_state = UniswapV4State::new(
1827            1000000,
1828            U256::from_str("79228162514264337593543950336").unwrap(), // 1:1 price
1829            UniswapV4Fees { zero_for_one: 0, one_for_zero: 0, lp_fee: 3000 },
1830            0,
1831            60,
1832            vec![
1833                // A tick with net_liquidity > current_liquidity
1834                // When zero_for_one=true, this gets negated and would cause underflow
1835                TickInfo {
1836                    index: -60,
1837                    net_liquidity: 2000000, // 2x current liquidity
1838                    sqrt_price: U256::from_str("79051508376726796163471739988").unwrap(),
1839                },
1840            ],
1841        )
1842        .unwrap();
1843
1844        let usdc = usdc();
1845        let weth = weth();
1846
1847        let (limit_in, limit_out) = usv4_state
1848            .get_limits(usdc.address.clone(), weth.address.clone())
1849            .unwrap();
1850
1851        // Should return some conservative limits
1852        assert!(limit_in > BigUint::zero());
1853        assert!(limit_out > BigUint::zero());
1854    }
1855
1856    // Tests based on Uniswap V4's ProtocolFeeLibrary.t.sol
1857    // See: https://github.com/Uniswap/v4-core/blob/main/test/libraries/ProtocolFeeLibrary.t.sol
1858
1859    /// Maximum protocol fee in pips (1000 = 0.1%)
1860    const MAX_PROTOCOL_FEE: u32 = 1000;
1861
1862    #[rstest]
1863    #[case::max_protocol_and_max_lp(MAX_PROTOCOL_FEE, lp_fee::MAX_LP_FEE, lp_fee::MAX_LP_FEE)]
1864    #[case::max_protocol_with_3000_lp(MAX_PROTOCOL_FEE, 3000, 3997)]
1865    #[case::max_protocol_with_zero_lp(MAX_PROTOCOL_FEE, 0, MAX_PROTOCOL_FEE)]
1866    #[case::zero_protocol_zero_lp(0, 0, 0)]
1867    #[case::zero_protocol_with_1000_lp(0, 1000, 1000)]
1868    fn test_calculate_swap_fees_uniswap_test_cases(
1869        #[case] protocol_fee: u32,
1870        #[case] lp_fee: u32,
1871        #[case] expected: u32,
1872    ) {
1873        let fees = UniswapV4Fees::new(protocol_fee, protocol_fee, lp_fee);
1874        let result = fees.calculate_swap_fees_pips(true, None);
1875        assert_eq!(result, expected);
1876    }
1877
1878    #[test]
1879    fn test_calculate_swap_fees_with_dynamic_fee() {
1880        // Test that dynamic fees default to 0 when no override is provided
1881        let fees = UniswapV4Fees::new(100, 90, lp_fee::DYNAMIC_FEE_FLAG);
1882
1883        // Without override, dynamic fee should be treated as 0
1884        let total_zero_for_one = fees.calculate_swap_fees_pips(true, None);
1885        // 100 + 0 - (100 * 0 / 1_000_000) = 100
1886        assert_eq!(total_zero_for_one, 100);
1887
1888        // With override, should use the override value
1889        let total_with_override = fees.calculate_swap_fees_pips(true, Some(500));
1890        // 100 + 500 - (100 * 500 / 1_000_000) = 600 - 0 = 600
1891        assert_eq!(total_with_override, 600);
1892    }
1893
1894    #[test]
1895    fn test_calculate_swap_fees_direction_matters() {
1896        // Test that zero_for_one direction affects which protocol fee is used
1897        let fees = UniswapV4Fees::new(100, 200, 500);
1898
1899        let zero_for_one_fee = fees.calculate_swap_fees_pips(true, None);
1900        // 100 + 500 - (100 * 500 / 1_000_000) = 600 - 0 = 600
1901        assert_eq!(zero_for_one_fee, 600);
1902
1903        let one_for_zero_fee = fees.calculate_swap_fees_pips(false, None);
1904        // 200 + 500 - (200 * 500 / 1_000_000) = 700 - 0 = 700
1905        assert_eq!(one_for_zero_fee, 700);
1906    }
1907
1908    #[rstest]
1909    #[case::high_lp_fee(1000, 500_000, 500_500)] // 1000 + 500k - 500 = 500.5k
1910    #[case::mid_fees(500, 500_000, 500_250)] // 500 + 500k - 250 = 500.25k
1911    #[case::low_fees(100, 100_000, 100_090)] // 100 + 100k - 10 = 100.09k
1912    fn test_calculate_swap_fees_formula_precision(
1913        #[case] protocol_fee: u32,
1914        #[case] lp_fee: u32,
1915        #[case] expected: u32,
1916    ) {
1917        // Test cases where the subtraction term (protocol * lp / 1M) significantly affects the
1918        // result
1919        let fees = UniswapV4Fees::new(protocol_fee, protocol_fee, lp_fee);
1920        let result = fees.calculate_swap_fees_pips(true, None);
1921        assert_eq!(result, expected, "Failed for protocol={}, lp={}", protocol_fee, lp_fee);
1922    }
1923
1924    #[test]
1925    fn test_calculate_swap_fees_override_takes_precedence() {
1926        // Test that lp_fee_override completely replaces stored lp_fee
1927        let fees = UniswapV4Fees::new(100, 100, 3000);
1928
1929        // With override, stored lp_fee should be ignored
1930        let result = fees.calculate_swap_fees_pips(true, Some(5000));
1931        // 100 + 5000 - (100 * 5000 / 1_000_000) = 5100 - 0 = 5100
1932        assert_eq!(result, 5100);
1933
1934        // Without override, should use stored lp_fee
1935        let result_no_override = fees.calculate_swap_fees_pips(true, None);
1936        // 100 + 3000 - (100 * 3000 / 1_000_000) = 3100 - 0 = 3100
1937        assert_eq!(result_no_override, 3100);
1938    }
1939
1940    #[test]
1941    fn test_calculate_swap_fees_zero_protocol_fee() {
1942        // When protocol fee is 0, formula simplifies to just lpFee
1943        let fees = UniswapV4Fees::new(0, 0, 3000);
1944        let result = fees.calculate_swap_fees_pips(true, None);
1945        // 0 + 3000 - (0 * 3000 / 1_000_000) = 3000
1946        assert_eq!(result, 3000);
1947    }
1948
1949    #[test]
1950    fn test_calculate_swap_fees_zero_lp_fee() {
1951        // When lp fee is 0, formula simplifies to just protocolFee
1952        let fees = UniswapV4Fees::new(500, 500, 0);
1953        let result = fees.calculate_swap_fees_pips(true, None);
1954        // 500 + 0 - (500 * 0 / 1_000_000) = 500
1955        assert_eq!(result, 500);
1956    }
1957
1958    // Helper to create a basic test pool for swap_to_price tests
1959    fn create_basic_v4_test_pool() -> UniswapV4State {
1960        let liquidity = 100_000_000_000_000_000_000u128; // 100e18
1961        let sqrt_price = get_sqrt_price_q96(U256::from(20_000_000u64), U256::from(10_000_000u64))
1962            .expect("Failed to calculate sqrt price");
1963        let tick = get_tick_at_sqrt_ratio(sqrt_price).expect("Failed to calculate tick");
1964
1965        let ticks = vec![TickInfo::new(0, 0).unwrap(), TickInfo::new(46080, 0).unwrap()];
1966
1967        UniswapV4State::new(
1968            liquidity,
1969            sqrt_price,
1970            UniswapV4Fees { zero_for_one: 0, one_for_zero: 0, lp_fee: 3000 }, // 0.3% fee
1971            tick,
1972            60, // tick spacing
1973            ticks,
1974        )
1975        .expect("Failed to create pool")
1976    }
1977
1978    fn create_tick_boundary_v4_test_pool() -> UniswapV4State {
1979        let sqrt_price = get_sqrt_ratio_at_tick(0).expect("Failed to calculate sqrt price");
1980        let ticks = vec![TickInfo::new(-120, 0).unwrap(), TickInfo::new(120, 0).unwrap()];
1981
1982        UniswapV4State::new(
1983            100_000_000_000_000_000_000u128,
1984            sqrt_price,
1985            UniswapV4Fees { zero_for_one: 0, one_for_zero: 0, lp_fee: 3000 },
1986            0,
1987            60,
1988            ticks,
1989        )
1990        .expect("Failed to create pool")
1991    }
1992
1993    #[test]
1994    fn test_partial_step_updates_tick_when_price_moves_without_crossing_initialized_tick() {
1995        let pool = create_tick_boundary_v4_test_pool();
1996        let amount = -I256::from_raw(U256::from(100_000_000_000_000_000u64));
1997
1998        let result = pool
1999            .swap(true, amount, None, None)
2000            .expect("swap should stay within the current liquidity range");
2001        let expected_tick =
2002            get_tick_at_sqrt_ratio(result.sqrt_price).expect("new sqrt price should map to a tick");
2003
2004        assert_ne!(result.sqrt_price, pool.sqrt_price);
2005        assert_ne!(result.sqrt_price, get_sqrt_ratio_at_tick(-120).unwrap());
2006        assert_ne!(expected_tick, pool.tick);
2007        assert_eq!(result.tick, expected_tick);
2008    }
2009
2010    #[test]
2011    fn test_swap_keeps_boundary_tick_when_price_does_not_move() {
2012        let mut pool = create_tick_boundary_v4_test_pool();
2013        pool.tick = -1;
2014        let amount = -I256::from_raw(U256::from(1u64));
2015
2016        let result = pool
2017            .swap(true, amount, None, None)
2018            .expect("swap should consume the input as fee without moving price");
2019
2020        assert_eq!(result.sqrt_price, pool.sqrt_price);
2021        assert_eq!(get_tick_at_sqrt_ratio(result.sqrt_price).unwrap(), 0);
2022        assert_eq!(result.tick, pool.tick);
2023    }
2024
2025    #[test]
2026    fn test_swap_to_price_price_too_high() {
2027        let pool = create_basic_v4_test_pool();
2028
2029        let token_x = token_x();
2030        let token_y = token_y();
2031
2032        // Price far above pool price - should return zero
2033        let target_price = Price::new(BigUint::from(10_000_000u64), BigUint::from(1_000_000u64));
2034
2035        let result = pool.query_pool_swap(&QueryPoolSwapParams::new(
2036            token_x,
2037            token_y,
2038            SwapConstraint::PoolTargetPrice {
2039                target: target_price,
2040                tolerance: 0f64,
2041                min_amount_in: None,
2042                max_amount_in: None,
2043            },
2044        ));
2045        assert!(result.is_err(), "Should return error when target price is unreachable");
2046    }
2047
2048    #[test]
2049    fn test_swap_to_price_no_liquidity() {
2050        // Test that swap_to_price returns zero for pool with no liquidity
2051        let pool = UniswapV4State::new(
2052            0, // No liquidity
2053            U256::from_str("79228162514264337593543950336").unwrap(),
2054            UniswapV4Fees { zero_for_one: 0, one_for_zero: 0, lp_fee: 3000 },
2055            0,
2056            60,
2057            vec![],
2058        )
2059        .unwrap();
2060
2061        let token_x = token_x();
2062        let token_y = token_y();
2063
2064        let target_price = Price::new(BigUint::from(2_000_000u64), BigUint::from(1_000_000u64));
2065
2066        let pool_swap = pool.query_pool_swap(&QueryPoolSwapParams::new(
2067            token_x,
2068            token_y,
2069            SwapConstraint::PoolTargetPrice {
2070                target: target_price,
2071                tolerance: 0f64,
2072                min_amount_in: None,
2073                max_amount_in: None,
2074            },
2075        ));
2076
2077        assert!(pool_swap.is_err());
2078    }
2079
2080    #[test]
2081    fn test_swap_to_price_with_protocol_fees() {
2082        let liquidity = 100_000_000_000_000_000_000u128;
2083        let sqrt_price = get_sqrt_price_q96(U256::from(20_000_000u64), U256::from(10_000_000u64))
2084            .expect("Failed to calculate sqrt price");
2085        let tick = get_tick_at_sqrt_ratio(sqrt_price).expect("Failed to calculate tick");
2086
2087        let ticks = vec![TickInfo::new(0, 0).unwrap(), TickInfo::new(46080, 0).unwrap()];
2088
2089        // Create pool with different protocol fees for each direction
2090        let pool = UniswapV4State::new(
2091            liquidity,
2092            sqrt_price,
2093            UniswapV4Fees {
2094                zero_for_one: 1000, // 0.1% protocol fee for zero_for_one
2095                one_for_zero: 200,  // 0.02% protocol fee for one_for_zero
2096                lp_fee: 3000,       // 0.3% LP fee
2097            },
2098            tick,
2099            60,
2100            ticks,
2101        )
2102        .expect("Failed to create pool");
2103
2104        let token_x = token_x();
2105        let token_y = token_y();
2106
2107        // Pool at 2.0 Y/X = 0.5 X/Y, swap_to_price moves price DOWN to target
2108
2109        // Test zero_for_one direction (X -> Y, uses zero_for_one fee)
2110        let target_price = Price::new(BigUint::from(2_000_000u64), BigUint::from(1_010_000u64));
2111        let pool_swap_forward = pool
2112            .query_pool_swap(&QueryPoolSwapParams::new(
2113                token_x.clone(),
2114                token_y.clone(),
2115                SwapConstraint::PoolTargetPrice {
2116                    target: target_price,
2117                    tolerance: 0f64,
2118                    min_amount_in: None,
2119                    max_amount_in: None,
2120                },
2121            ))
2122            .expect("swap_to_price failed");
2123
2124        // Test one_for_zero direction (Y -> X, uses one_for_zero fee)
2125        let target_price_reverse =
2126            Price::new(BigUint::from(1_010_000u64), BigUint::from(2_040_000u64));
2127        let pool_swap_backward = pool
2128            .query_pool_swap(&QueryPoolSwapParams::new(
2129                token_y,
2130                token_x,
2131                SwapConstraint::PoolTargetPrice {
2132                    target: target_price_reverse,
2133                    tolerance: 0f64,
2134                    min_amount_in: None,
2135                    max_amount_in: None,
2136                },
2137            ))
2138            .expect("swap_to_price failed");
2139
2140        assert!(
2141            pool_swap_backward.amount_out().clone() > BigUint::ZERO,
2142            "One for zero swap should return non-zero output"
2143        );
2144
2145        // Higher fees require more volume to reach the same price target
2146        // trade_zfo has 0.1% protocol fee, trade_ofz has 0.02% protocol fee
2147        assert!(
2148            pool_swap_forward.amount_out() < pool_swap_backward.amount_in(),
2149            "Backward fees should be lower therefore backward swap should be bigger"
2150        );
2151        assert!(
2152            pool_swap_forward.amount_in() < pool_swap_backward.amount_out(),
2153            "Backward fees should be lower therefore backward swap should be bigger"
2154        );
2155    }
2156
2157    #[test]
2158    fn test_swap_to_price_different_targets() {
2159        // Test with various target prices using working format
2160        let pool = create_basic_v4_test_pool();
2161
2162        let token_x = token_x();
2163        let token_y = token_y();
2164
2165        // Pool at 2.0 Y/X (20M/10M)
2166        // Test 1: Target close to spot (1.98 Y/X)
2167        let target_price = Price::new(BigUint::from(2_000_000u64), BigUint::from(1_010_000u64));
2168        let pool_swap_close = pool
2169            .query_pool_swap(&QueryPoolSwapParams::new(
2170                token_x.clone(),
2171                token_y.clone(),
2172                SwapConstraint::PoolTargetPrice {
2173                    target: target_price,
2174                    tolerance: 0f64,
2175                    min_amount_in: None,
2176                    max_amount_in: None,
2177                },
2178            ))
2179            .expect("swap_to_price failed");
2180        assert!(
2181            *pool_swap_close.amount_out() > BigUint::ZERO,
2182            "Expected non-zero for 1.98 Y/X target"
2183        );
2184
2185        // Test 2: Target further from spot (1.90 Y/X)
2186        let target_price = Price::new(BigUint::from(1_900_000u64), BigUint::from(1_000_000u64));
2187        let pool_swap_below = pool
2188            .query_pool_swap(&QueryPoolSwapParams::new(
2189                token_x.clone(),
2190                token_y.clone(),
2191                SwapConstraint::PoolTargetPrice {
2192                    target: target_price,
2193                    tolerance: 0f64,
2194                    min_amount_in: None,
2195                    max_amount_in: None,
2196                },
2197            ))
2198            .expect("swap_to_price failed");
2199        assert!(
2200            pool_swap_below.amount_out().clone() > BigUint::ZERO,
2201            "Expected non-zero for 1.90 Y/X target"
2202        );
2203
2204        // Test 3: Target far from spot (1.5 Y/X)
2205        let target_price = Price::new(BigUint::from(1_500_000u64), BigUint::from(1_000_000u64));
2206        let pool_swap_far = pool
2207            .query_pool_swap(&QueryPoolSwapParams::new(
2208                token_x,
2209                token_y,
2210                SwapConstraint::PoolTargetPrice {
2211                    target: target_price,
2212                    tolerance: 0f64,
2213                    min_amount_in: None,
2214                    max_amount_in: None,
2215                },
2216            ))
2217            .expect("swap_to_price failed");
2218        assert!(
2219            pool_swap_far.amount_out().clone() > BigUint::ZERO,
2220            "Expected non-zero for 1.5 Y/X target"
2221        );
2222
2223        // Verify that further targets require more volume
2224        assert!(
2225            pool_swap_close.amount_out().clone() < pool_swap_below.amount_out().clone(),
2226            "Closer target (1.98 Y/X) should require less volume than medium target (1.90 Y/X). \
2227             Got close: {}, medium: {}",
2228            pool_swap_close.amount_out().clone(),
2229            pool_swap_below.amount_out().clone()
2230        );
2231        assert!(
2232            pool_swap_below.amount_out().clone() < pool_swap_far.amount_out().clone(),
2233            "Medium target (1.90 Y/X) should require less volume than far target (1.5 Y/X). \
2234             Got medium: {}, far: {}",
2235            pool_swap_below.amount_out().clone(),
2236            pool_swap_far.amount_out().clone()
2237        );
2238    }
2239
2240    #[test]
2241    fn test_swap_to_price_around_spot_price() {
2242        let liquidity = 10_000_000_000_000_000u128;
2243        let sqrt_price =
2244            get_sqrt_price_q96(U256::from(2_000_000_000u64), U256::from(1_000_000_000u64))
2245                .expect("Failed to calculate sqrt price");
2246        let tick = get_tick_at_sqrt_ratio(sqrt_price).expect("Failed to calculate tick");
2247
2248        let ticks = vec![TickInfo::new(0, 0).unwrap(), TickInfo::new(46080, 0).unwrap()];
2249
2250        // Use FeeAmount::Low equivalent (500 pips = 0.05%)
2251        let pool = UniswapV4State::new(
2252            liquidity,
2253            sqrt_price,
2254            UniswapV4Fees {
2255                zero_for_one: 0,
2256                one_for_zero: 0,
2257                lp_fee: 500, // 0.05% to match V3 FeeAmount::Low
2258            },
2259            tick,
2260            60,
2261            ticks,
2262        )
2263        .expect("Failed to create pool");
2264
2265        let token_x = token_x();
2266        let token_y = token_y();
2267
2268        // Test 1: Price just above spot price, too little to cover fees
2269        let target_price = Price::new(BigUint::from(1_999_750u64), BigUint::from(1_000_250u64));
2270
2271        let result = pool.query_pool_swap(&QueryPoolSwapParams::new(
2272            token_x.clone(),
2273            token_y.clone(),
2274            SwapConstraint::PoolTargetPrice {
2275                target: target_price,
2276                tolerance: 0f64,
2277                min_amount_in: None,
2278                max_amount_in: None,
2279            },
2280        ));
2281        assert!(result.is_err(), "Should return error when target price is unreachable");
2282
2283        // Test 2: Price far enough from spot prices to enable trading despite fees (0.1% lower)
2284        let target_price = Price::new(BigUint::from(1_999_000u64), BigUint::from(1_001_000u64));
2285
2286        let pool_swap = pool
2287            .query_pool_swap(&QueryPoolSwapParams::new(
2288                token_x,
2289                token_y,
2290                SwapConstraint::PoolTargetPrice {
2291                    target: target_price,
2292                    tolerance: 0f64,
2293                    min_amount_in: None,
2294                    max_amount_in: None,
2295                },
2296            ))
2297            .expect("swap_to_price failed");
2298
2299        // Should match V3 output exactly with same fees
2300        let expected_amount_out =
2301            BigUint::from_str("7062236922008").expect("Failed to parse expected value");
2302        assert_eq!(
2303            pool_swap.amount_out().clone(),
2304            expected_amount_out,
2305            "V4 should match V3 output with same fees (0.05%)"
2306        );
2307    }
2308
2309    #[test]
2310    fn test_swap_to_price_matches_get_amount_out() {
2311        let pool = create_basic_v4_test_pool();
2312
2313        let token_x = token_x();
2314        let token_y = token_y();
2315
2316        // Get the trade from swap_to_price
2317        let target_price = Price::new(BigUint::from(2_000_000u64), BigUint::from(1_010_000u64));
2318        let pool_swap = pool
2319            .query_pool_swap(&QueryPoolSwapParams::new(
2320                token_x.clone(),
2321                token_y.clone(),
2322                SwapConstraint::PoolTargetPrice {
2323                    target: target_price,
2324                    tolerance: 0f64,
2325                    min_amount_in: None,
2326                    max_amount_in: None,
2327                },
2328            ))
2329            .expect("swap_to_price failed");
2330        assert!(*pool_swap.amount_in() > BigUint::ZERO, "Amount in should be positive");
2331
2332        // Use the amount_in from swap_to_price with get_amount_out
2333        let result = pool
2334            .get_amount_out(pool_swap.amount_in().clone(), &token_x, &token_y)
2335            .expect("get_amount_out failed");
2336
2337        // The amount_out from get_amount_out should be close to swap_to_price's amount_out
2338        // Allow for small rounding differences
2339        assert!(result.amount > BigUint::ZERO);
2340        assert!(result.amount >= *pool_swap.amount_out());
2341    }
2342
2343    #[test]
2344    fn test_swap_to_price_basic() {
2345        let liquidity = 100_000_000_000_000_000_000u128;
2346        let sqrt_price = get_sqrt_price_q96(U256::from(20_000_000u64), U256::from(10_000_000u64))
2347            .expect("Failed to calculate sqrt price");
2348        let tick = get_tick_at_sqrt_ratio(sqrt_price).expect("Failed to calculate tick");
2349
2350        let ticks = vec![TickInfo::new(0, 0).unwrap(), TickInfo::new(46080, 0).unwrap()];
2351
2352        let pool = UniswapV4State::new(
2353            liquidity,
2354            sqrt_price,
2355            UniswapV4Fees {
2356                zero_for_one: 0,
2357                one_for_zero: 0,
2358                lp_fee: 3000, // 0.3% LP fee
2359            },
2360            tick,
2361            60,
2362            ticks,
2363        )
2364        .expect("Failed to create pool");
2365
2366        let token_x = token_x();
2367        let token_y = token_y();
2368
2369        // Target price: 2_000_000/1_010_000 ≈ 1.98 Y/X
2370        let target_price = Price::new(BigUint::from(2_000_000u64), BigUint::from(1_010_000u64));
2371
2372        let pool_swap = pool
2373            .query_pool_swap(&QueryPoolSwapParams::new(
2374                token_x,
2375                token_y,
2376                SwapConstraint::PoolTargetPrice {
2377                    target: target_price,
2378                    tolerance: 0f64,
2379                    min_amount_in: None,
2380                    max_amount_in: None,
2381                },
2382            ))
2383            .expect("swap_to_price failed");
2384
2385        // Should match V3's output exactly with same fees (0.3%)
2386        let expected_amount_in = BigUint::from_str("246739021727519745").unwrap();
2387        let expected_amount_out = BigUint::from_str("490291909043340795").unwrap();
2388
2389        assert_eq!(
2390            *pool_swap.amount_in(),
2391            expected_amount_in,
2392            "amount_in should match expected value"
2393        );
2394        assert_eq!(
2395            *pool_swap.amount_out(),
2396            expected_amount_out,
2397            "amount_out should match expected value"
2398        );
2399    }
2400
2401    #[test]
2402    fn test_swap_price_limit_out_of_range_returns_error() {
2403        let pool = create_basic_v4_test_pool();
2404        let amount = -I256::from_raw(U256::from(1000u64)); // V4 uses negative for exact input
2405
2406        // zero_for_one: price_limit equal to sqrt_price is invalid (must be strictly less)
2407        let result = pool.swap(true, amount, Some(pool.sqrt_price), None);
2408        assert!(matches!(result, Err(SimulationError::InvalidInput(_, None))));
2409
2410        // zero_for_one: price_limit at MIN_SQRT_RATIO is invalid (must be strictly greater)
2411        let result = pool.swap(true, amount, Some(MIN_SQRT_RATIO), None);
2412        assert!(matches!(result, Err(SimulationError::InvalidInput(_, None))));
2413
2414        // one_for_zero: price_limit equal to sqrt_price is invalid (must be strictly greater)
2415        let result = pool.swap(false, amount, Some(pool.sqrt_price), None);
2416        assert!(matches!(result, Err(SimulationError::InvalidInput(_, None))));
2417
2418        // one_for_zero: price_limit at MAX_SQRT_RATIO is invalid (must be strictly less)
2419        let result = pool.swap(false, amount, Some(MAX_SQRT_RATIO), None);
2420        assert!(matches!(result, Err(SimulationError::InvalidInput(_, None))));
2421    }
2422
2423    #[test]
2424    fn test_swap_at_extreme_price_returns_error() {
2425        // Simulates the depth calculation scenario: pool sqrt_price is at MIN_SQRT_RATIO + 1,
2426        // so the default price limit for zero_for_one equals sqrt_price and fails validation.
2427        let sqrt_price = MIN_SQRT_RATIO + U256::from(1u64);
2428        let tick = get_tick_at_sqrt_ratio(sqrt_price).expect("Failed to calculate tick");
2429        // tick_spacing 60; ticks must be aligned
2430        let aligned_tick = (MIN_TICK / 60) * 60 + 60; // first multiple of 60 above MIN_TICK
2431        let ticks = vec![
2432            TickInfo::new(aligned_tick, 0).unwrap(),
2433            TickInfo::new(aligned_tick + 60, 0).unwrap(),
2434        ];
2435        let pool = UniswapV4State::new(
2436            100_000_000_000_000_000_000u128,
2437            sqrt_price,
2438            UniswapV4Fees { zero_for_one: 0, one_for_zero: 0, lp_fee: 3000 },
2439            tick,
2440            60,
2441            ticks,
2442        )
2443        .unwrap();
2444
2445        let amount = -I256::from_raw(U256::from(1000u64));
2446        // Default price limit for zero_for_one is MIN_SQRT_RATIO + 1 == sqrt_price, so invalid
2447        let result = pool.swap(true, amount, None, None);
2448        assert!(matches!(result, Err(SimulationError::InvalidInput(_, None))));
2449    }
2450}