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