tycho-simulation 0.456.0

use std::{any::Any, collections::HashMap, sync::Arc};

use alloy::primitives::{Sign, I256, U256};
use num_bigint::BigUint;
use num_traits::Zero;
use serde::{Deserialize, Serialize};
use tracing::trace;
use tycho_common::{
    dto::ProtocolStateDelta,
    models::token::Token,
    simulation::{
        errors::{SimulationError, TransitionError},
        protocol_sim::{
            Balances, GetAmountOutResult, PoolSwap, Price, ProtocolSim, QueryPoolSwapParams,
            SwapConstraint,
        },
    },
    Bytes,
};

use crate::evm::{
    protocol::{
        clmm::clmm_swap_to_price,
        safe_math::{safe_add_u256, safe_sub_u256},
        u256_num::u256_to_biguint,
        utils::{
            slipstreams::{dynamic_fee_module::ZERO_FEE_INDICATOR, raw_target_price},
            uniswap::{
                liquidity_math,
                pool_tick::PoolTick,
                sqrt_price_math::{get_amount0_delta, get_amount1_delta, sqrt_price_q96_to_f64},
                swap_math::{self, SwapStepResult},
                swap_step_cache::{StepCacheKey, StepCacheRun, SwapStepCache},
                tick_list::{TickInfo, TickList, TickListErrorKind},
                tick_math::{
                    get_sqrt_ratio_at_tick_cached, MAX_SQRT_RATIO, MAX_TICK, MIN_SQRT_RATIO,
                    MIN_TICK,
                },
                StepComputation, SwapResults, SwapState,
            },
        },
    },
    query_pool_swap::price_to_f64_with_decimals,
};

// The names of the constants reflect the exact method from the tenderly log.
const GAS_PER_TICK: u64 = 25_000;
// nextInitializedTickWithinOneWord +  computeSwapStep + calculateFees
const GAS_PER_LOOP: u64 = 10_000;

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct VelodromeSlipstreamsState {
    liquidity: u128,
    sqrt_price: U256,
    default_fee: u32,
    custom_fee: u32,
    tick_spacing: i32,
    tick: PoolTick,
    /// The tick list, shared by this state, its clones and the states its swaps return, so a quote
    /// does not copy it. Change it only through `Arc::make_mut`, which copies the list while
    /// another state shares it.
    ticks: Arc<TickList>,
    /// Swap steps already taken from this state, shared by its clones. The steps hold only for
    /// this state's price, tick, liquidity, fee and ticks. `delta_transition` and
    /// `clone_post_swap` replace the cache with an empty one, and any other code that changes
    /// those fields must do the same.
    #[serde(skip)]
    step_cache: SwapStepCache,
}

impl VelodromeSlipstreamsState {
    /// Creates a new instance of `AerodromeSlipstreamsState`.
    ///
    /// # Arguments
    /// - `liquidity`: The initial liquidity of the pool.
    /// - `sqrt_price`: The square root of the current price.
    /// - `default_fee`: The default fee for the pool.
    /// - `custom_fee`: The custom fee for the pool.
    /// - `tick_spacing`: The tick spacing for the pool.
    /// - `tick`: The current tick of the pool.
    /// - `ticks`: A vector of `TickInfo` representing the tick information for the pool.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        liquidity: u128,
        sqrt_price: U256,
        default_fee: u32,
        custom_fee: u32,
        tick_spacing: i32,
        tick: i32,
        ticks: Vec<TickInfo>,
    ) -> Result<Self, SimulationError> {
        let tick_list = TickList::from(tick_spacing as u16, ticks)?;
        Ok(VelodromeSlipstreamsState {
            liquidity,
            sqrt_price,
            default_fee,
            custom_fee,
            tick_spacing,
            tick: tick.into(),
            ticks: Arc::new(tick_list),
            step_cache: SwapStepCache::default(),
        })
    }

    /// Returns the swap fee in pips. A custom fee of zero means that none is set, and
    /// [`ZERO_FEE_INDICATOR`] is the value the fee module stores for an explicit zero fee.
    fn get_fee(&self) -> u32 {
        match self.custom_fee {
            0 => self.default_fee,
            ZERO_FEE_INDICATOR => 0,
            custom_fee => custom_fee,
        }
    }

    /// Swaps token_in until `spot_price(token_in, token_out)` reaches the middle of the band
    /// `[target, target * (1 + tolerance)]`. Falls back to the numerical search on errors, and
    /// when the native swap stops outside the band.
    fn swap_to_target_price(
        &self,
        params: &QueryPoolSwapParams,
        target: &Price,
        tolerance: f64,
    ) -> Result<PoolSwap, SimulationError> {
        let token_in = params.token_in();
        let token_out = params.token_out();
        let target_f64 = price_to_f64_with_decimals(target, token_in.decimals, token_out.decimals)?;
        if target_f64 == self.spot_price(token_in, token_out)? {
            return Ok(PoolSwap::new(BigUint::ZERO, BigUint::ZERO, Box::new(self.clone()), None));
        }
        // `spot_price` has no fee markup, so the limit and the swap both use fee 0.
        let Some(limit) = raw_target_price(target, tolerance, 0) else {
            return crate::evm::query_pool_swap::query_pool_swap(self, params);
        };
        let Ok((amount_in, amount_out, result)) = clmm_swap_to_price(
            self.sqrt_price,
            &token_in.address,
            &token_out.address,
            &limit,
            0,
            Sign::Positive,
            |zero_for_one, amount_specified, sqrt_price_limit| {
                self.swap(zero_for_one, amount_specified, Some(sqrt_price_limit))
            },
        ) else {
            return crate::evm::query_pool_swap::query_pool_swap(self, params);
        };

        let new_state = if amount_in.is_zero() {
            self.clone()
        } else {
            self.clone_post_swap(result.sqrt_price, result.tick, result.liquidity)
        };
        // The native swap can spend its whole input allowance before it reaches the limit.
        let final_spot = new_state.spot_price(token_in, token_out)?;
        if final_spot < target_f64 || final_spot > target_f64 * (1.0 + tolerance) {
            return crate::evm::query_pool_swap::query_pool_swap(self, params);
        }
        Ok(PoolSwap::new(amount_in, amount_out, Box::new(new_state), None))
    }

    /// A clone of this state with the final price, tick and liquidity of a swap, and an empty
    /// step cache.
    fn clone_post_swap(&self, sqrt_price: U256, tick: PoolTick, liquidity: u128) -> Self {
        let mut state = self.clone();
        state.sqrt_price = sqrt_price;
        state.tick = tick;
        state.liquidity = liquidity;
        state.step_cache = SwapStepCache::default();
        state
    }

    fn swap(
        &self,
        zero_for_one: bool,
        amount_specified: I256,
        sqrt_price_limit: Option<U256>,
    ) -> Result<SwapResults, SimulationError> {
        if self.liquidity == 0 {
            return Err(SimulationError::RecoverableError("No liquidity".to_string()));
        }
        let price_limit = if let Some(limit) = sqrt_price_limit {
            limit
        } else if zero_for_one {
            safe_add_u256(MIN_SQRT_RATIO, U256::from(1u64))?
        } else {
            safe_sub_u256(MAX_SQRT_RATIO, U256::from(1u64))?
        };

        let price_limit_valid = if zero_for_one {
            price_limit > MIN_SQRT_RATIO && price_limit < self.sqrt_price
        } else {
            price_limit < MAX_SQRT_RATIO && price_limit > self.sqrt_price
        };
        if !price_limit_valid {
            return Err(SimulationError::InvalidInput("Price limit out of range".into(), None));
        }

        let exact_input = amount_specified > I256::from_raw(U256::from(0u64));

        let mut state = SwapState {
            amount_remaining: amount_specified,
            amount_calculated: I256::from_raw(U256::from(0u64)),
            sqrt_price: self.sqrt_price,
            tick: self.tick.value().into(),
            liquidity: self.liquidity,
        };
        let mut gas_used = U256::from(130_000);

        let fee = self.get_fee();
        // Only exact-input swaps with no price limit use the step cache. The cache finds a step by
        // its input, which an exact-output swap does not know. A cached step can also lie past a
        // price limit.
        let cache_key = if exact_input && sqrt_price_limit.is_none() {
            StepCacheKey::new(self.sqrt_price, state.tick.value(), self.liquidity, fee, gas_used)
        } else {
            None
        };
        let mut cache_run = StepCacheRun::new(
            &self.step_cache,
            zero_for_one,
            cache_key,
            amount_specified.unsigned_abs(),
        );
        if let Some(progress) = cache_run.apply_cached_progress(&mut state) {
            gas_used = progress.gas;
        }
        while state.amount_remaining != I256::from_raw(U256::from(0u64)) &&
            state.sqrt_price != price_limit
        {
            let (mut next_tick, next_tick_info) = match self
                .ticks
                .next_initialized_tick_within_one_word(state.tick.value(), zero_for_one)
            {
                Ok((tick, info)) => (tick, info),
                Err(tick_err) => match tick_err.kind {
                    TickListErrorKind::TicksExeeded => {
                        let new_state = self.clone_post_swap(
                            state.sqrt_price,
                            state.tick.clone(),
                            state.liquidity,
                        );
                        return Err(SimulationError::InvalidInput(
                            "Ticks exceeded".into(),
                            Some(GetAmountOutResult::new(
                                u256_to_biguint(state.amount_calculated.abs().into_raw()),
                                u256_to_biguint(gas_used),
                                Box::new(new_state),
                            )),
                        ));
                    }
                    _ => return Err(SimulationError::FatalError("Unknown error".to_string())),
                },
            };

            next_tick = next_tick.clamp(MIN_TICK, MAX_TICK);

            let sqrt_price_start = state.sqrt_price;
            let sqrt_price_next = match next_tick_info {
                Some(info) => info.sqrt_price,
                None => get_sqrt_ratio_at_tick_cached(next_tick)?,
            };
            let SwapStepResult { sqrt_price, amount_in_with_fee, amount_out } =
                swap_math::compute_swap_step(
                    state.sqrt_price,
                    VelodromeSlipstreamsState::get_sqrt_ratio_target(
                        sqrt_price_next,
                        price_limit,
                        zero_for_one,
                    ),
                    state.liquidity,
                    state.amount_remaining,
                    fee,
                )?;
            state.sqrt_price = sqrt_price;

            let step = StepComputation {
                sqrt_price_start,
                tick_next: next_tick,
                sqrt_price_next,
                amount_in_with_fee,
                amount_out,
            };
            if exact_input {
                state.amount_remaining -=
                    I256::checked_from_sign_and_abs(Sign::Positive, step.amount_in_with_fee)
                        .unwrap();
                state.amount_calculated -=
                    I256::checked_from_sign_and_abs(Sign::Positive, step.amount_out).unwrap();
            } else {
                state.amount_remaining +=
                    I256::checked_from_sign_and_abs(Sign::Positive, step.amount_out).unwrap();
                state.amount_calculated +=
                    I256::checked_from_sign_and_abs(Sign::Positive, step.amount_in_with_fee)
                        .unwrap();
            }
            if state.sqrt_price == step.sqrt_price_next {
                if let Some(info) = next_tick_info {
                    let liquidity_raw = info.net_liquidity;
                    let liquidity_net = if zero_for_one { -liquidity_raw } else { liquidity_raw };
                    state.liquidity =
                        liquidity_math::add_liquidity_delta(state.liquidity, liquidity_net)?;
                    gas_used = safe_add_u256(gas_used, U256::from(GAS_PER_TICK))?;
                }
                state.tick =
                    PoolTick::from(if zero_for_one { step.tick_next - 1 } else { step.tick_next });
            } else if state.sqrt_price != step.sqrt_price_start {
                state.tick = PoolTick::at_sqrt_price(state.sqrt_price)?;
            }
            gas_used = safe_add_u256(gas_used, U256::from(GAS_PER_LOOP))?;
            cache_run.cache_progress(&step, &state, gas_used);
        }
        Ok(SwapResults {
            amount_calculated: state.amount_calculated,
            amount_specified,
            amount_remaining: state.amount_remaining,
            sqrt_price: state.sqrt_price,
            liquidity: state.liquidity,
            tick: state.tick,
            gas_used,
        })
    }

    fn get_sqrt_ratio_target(
        sqrt_price_next: U256,
        sqrt_price_limit: U256,
        zero_for_one: bool,
    ) -> U256 {
        let cond1 = if zero_for_one {
            sqrt_price_next < sqrt_price_limit
        } else {
            sqrt_price_next > sqrt_price_limit
        };

        if cond1 {
            sqrt_price_limit
        } else {
            sqrt_price_next
        }
    }
}

#[typetag::serde]
impl ProtocolSim for VelodromeSlipstreamsState {
    fn fee(&self) -> f64 {
        self.get_fee() as f64 / 1_000_000.0
    }

    fn spot_price(&self, a: &Token, b: &Token) -> Result<f64, SimulationError> {
        if a < b {
            sqrt_price_q96_to_f64(self.sqrt_price, a.decimals, b.decimals)
        } else {
            sqrt_price_q96_to_f64(self.sqrt_price, b.decimals, a.decimals)
                .map(|price| 1.0f64 / price)
        }
    }

    fn get_amount_out(
        &self,
        amount_in: BigUint,
        token_a: &Token,
        token_b: &Token,
    ) -> Result<GetAmountOutResult, SimulationError> {
        let zero_for_one = token_a < token_b;
        let amount_specified = I256::checked_from_sign_and_abs(
            Sign::Positive,
            U256::from_be_slice(&amount_in.to_bytes_be()),
        )
        .ok_or_else(|| {
            SimulationError::InvalidInput("I256 overflow: amount_in".to_string(), None)
        })?;

        let result = self.swap(zero_for_one, amount_specified, None)?;

        trace!(?amount_in, ?token_a, ?token_b, ?zero_for_one, ?result, "SLIPSTREAMS SWAP");
        let new_state = self.clone_post_swap(result.sqrt_price, result.tick, result.liquidity);

        Ok(GetAmountOutResult::new(
            u256_to_biguint(
                result
                    .amount_calculated
                    .abs()
                    .into_raw(),
            ),
            u256_to_biguint(result.gas_used),
            Box::new(new_state),
        ))
    }

    fn get_limits(
        &self,
        token_in: Bytes,
        token_out: Bytes,
    ) -> Result<(BigUint, BigUint), SimulationError> {
        // If the pool has no liquidity, return zeros for both limits
        if self.liquidity == 0 {
            return Ok((BigUint::zero(), BigUint::zero()));
        }

        let zero_for_one = token_in < token_out;
        let mut current_tick = self.tick.value();
        let mut current_sqrt_price = self.sqrt_price;
        let mut current_liquidity = self.liquidity;
        let mut total_amount_in = U256::from(0u64);
        let mut total_amount_out = U256::from(0u64);

        // Iterate through all ticks in the direction of the swap
        // Continues until there is no more liquidity in the pool or no more ticks to process
        while let Ok((tick, next_tick_info)) = self
            .ticks
            .next_initialized_tick_within_one_word(current_tick, zero_for_one)
        {
            // Clamp the tick value to ensure it's within valid range
            let next_tick = tick.clamp(MIN_TICK, MAX_TICK);

            let sqrt_price_next = match next_tick_info {
                Some(info) => info.sqrt_price,
                None => get_sqrt_ratio_at_tick_cached(next_tick)?,
            };

            // Calculate the amount of tokens swapped when moving from current_sqrt_price to
            // sqrt_price_next. Direction determines which token is being swapped in vs out
            let (amount_in, amount_out) = if zero_for_one {
                let amount0 = get_amount0_delta(
                    sqrt_price_next,
                    current_sqrt_price,
                    current_liquidity,
                    true,
                )?;
                let amount1 = get_amount1_delta(
                    sqrt_price_next,
                    current_sqrt_price,
                    current_liquidity,
                    false,
                )?;
                (amount0, amount1)
            } else {
                let amount0 = get_amount0_delta(
                    sqrt_price_next,
                    current_sqrt_price,
                    current_liquidity,
                    false,
                )?;
                let amount1 = get_amount1_delta(
                    sqrt_price_next,
                    current_sqrt_price,
                    current_liquidity,
                    true,
                )?;
                (amount1, amount0)
            };

            // Accumulate total amounts for this tick range
            total_amount_in = safe_add_u256(total_amount_in, amount_in)?;
            total_amount_out = safe_add_u256(total_amount_out, amount_out)?;

            // If this tick is "initialized" (meaning its someone's position boundary), update the
            // liquidity when crossing it
            // For zero_for_one, liquidity is removed when crossing a tick
            // For one_for_zero, liquidity is added when crossing a tick
            if let Some(info) = next_tick_info {
                let liquidity_raw = info.net_liquidity;
                let liquidity_delta = if zero_for_one { -liquidity_raw } else { liquidity_raw };
                current_liquidity =
                    liquidity_math::add_liquidity_delta(current_liquidity, liquidity_delta)?;
            }

            // Move to the next tick position
            current_tick = if zero_for_one { next_tick - 1 } else { next_tick };
            current_sqrt_price = sqrt_price_next;
        }

        Ok((u256_to_biguint(total_amount_in), u256_to_biguint(total_amount_out)))
    }

    fn delta_transition(
        &mut self,
        delta: ProtocolStateDelta,
        _tokens: &HashMap<Bytes, Token>,
        _balances: &Balances,
    ) -> Result<(), TransitionError> {
        // Swap steps taken from the old state do not hold for the new one.
        self.step_cache = SwapStepCache::default();
        // apply attribute changes
        if let Some(liquidity) = delta
            .updated_attributes
            .get("liquidity")
        {
            self.liquidity = u128::from(liquidity.clone());
        }
        if let Some(sqrt_price) = delta
            .updated_attributes
            .get("sqrt_price_x96")
        {
            self.sqrt_price = U256::from_be_slice(sqrt_price);
        }
        if let Some(default_fee) = delta
            .updated_attributes
            .get("default_fee")
        {
            self.default_fee = u32::from(default_fee.clone());
        }
        if let Some(custom_fee) = delta
            .updated_attributes
            .get("custom_fee")
        {
            self.custom_fee = u32::from(custom_fee.clone());
        }
        if let Some(tick) = delta.updated_attributes.get("tick") {
            self.tick = i32::from(tick.clone()).into();
        }

        // apply tick & observations changes
        for (key, value) in delta.updated_attributes.iter() {
            // tick liquidity keys are in the format "ticks/{tick_index}/net_liquidity"
            if key.starts_with("ticks/") {
                let parts: Vec<&str> = key.split('/').collect();
                Arc::make_mut(&mut self.ticks)
                    .set_tick_liquidity(
                        parts[1]
                            .parse::<i32>()
                            .map_err(|err| TransitionError::DecodeError(err.to_string()))?,
                        i128::from(value.clone()),
                    )
                    .map_err(|err| TransitionError::DecodeError(err.to_string()))?;
            }
        }
        // delete ticks - ignores deletes for attributes other than tick liquidity
        for key in delta.deleted_attributes.iter() {
            // tick liquidity keys are in the format "ticks/{tick_index}/net_liquidity"
            if key.starts_with("ticks/") {
                let parts: Vec<&str> = key.split('/').collect();
                Arc::make_mut(&mut self.ticks)
                    .set_tick_liquidity(
                        parts[1]
                            .parse::<i32>()
                            .map_err(|err| TransitionError::DecodeError(err.to_string()))?,
                        0,
                    )
                    .map_err(|err| TransitionError::DecodeError(err.to_string()))?;
            }
        }
        Ok(())
    }

    /// Answers [`SwapConstraint::PoolTargetPrice`] with one native swap and no `price_points`.
    /// [`SwapConstraint::TradeLimitPrice`] uses the numerical search.
    fn query_pool_swap(&self, params: &QueryPoolSwapParams) -> Result<PoolSwap, SimulationError> {
        match params.swap_constraint() {
            SwapConstraint::TradeLimitPrice { .. } => {
                crate::evm::query_pool_swap::query_pool_swap(self, params)
            }
            SwapConstraint::PoolTargetPrice {
                target,
                tolerance,
                min_amount_in: _,
                max_amount_in: _,
            } => self.swap_to_target_price(params, target, *tolerance),
        }
    }

    fn clone_box(&self) -> Box<dyn ProtocolSim> {
        Box::new(self.clone())
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    fn eq(&self, other: &dyn ProtocolSim) -> bool {
        if let Some(other_state) = other
            .as_any()
            .downcast_ref::<VelodromeSlipstreamsState>()
        {
            self.liquidity == other_state.liquidity &&
                self.sqrt_price == other_state.sqrt_price &&
                self.get_fee() == other_state.get_fee() &&
                self.tick == other_state.tick &&
                self.ticks == other_state.ticks
        } else {
            false
        }
    }
}

#[cfg(test)]
mod tests {
    use alloy::primitives::{Sign, I256, U256};
    use rstest::rstest;
    use tycho_common::{models::Chain, simulation::errors::SimulationError};

    use super::*;
    use crate::evm::{
        protocol::utils::uniswap::{
            tick_list::TickInfo,
            tick_math::{
                get_sqrt_ratio_at_tick, get_tick_at_sqrt_ratio, MAX_SQRT_RATIO, MIN_SQRT_RATIO,
                MIN_TICK,
            },
        },
        query_pool_swap::test_helpers::{target_price_params, to_price},
    };

    fn create_basic_test_pool() -> VelodromeSlipstreamsState {
        let sqrt_price = get_sqrt_ratio_at_tick(0).expect("Failed to calculate sqrt price");
        let ticks = vec![TickInfo::new(-120, 0).unwrap(), TickInfo::new(120, 0).unwrap()];
        VelodromeSlipstreamsState::new(
            100_000_000_000_000_000_000u128,
            sqrt_price,
            3000,
            0,
            1,
            0,
            ticks,
        )
        .expect("Failed to create pool")
    }

    #[rstest]
    #[case::no_custom_fee(0, 0.003)]
    #[case::zero_fee_indicator(420, 0.0)]
    #[case::custom_fee(500, 0.0005)]
    fn test_fee_resolves_custom_fee(#[case] custom_fee: u32, #[case] expected: f64) {
        let sqrt_price = get_sqrt_ratio_at_tick(0).unwrap();
        let ticks = vec![TickInfo::new(-120, 0).unwrap(), TickInfo::new(120, 0).unwrap()];
        let pool = VelodromeSlipstreamsState::new(
            10u128.pow(20),
            sqrt_price,
            3000,
            custom_fee,
            1,
            0,
            ticks,
        )
        .unwrap();

        assert_eq!(pool.fee(), expected);
    }

    #[test]
    fn test_partial_step_updates_tick_when_price_moves_without_crossing_initialized_tick() {
        let pool = create_basic_test_pool();
        let amount =
            I256::checked_from_sign_and_abs(Sign::Positive, U256::from(100_000_000_000_000_000u64))
                .unwrap();

        let result = pool
            .swap(true, amount, None)
            .expect("swap should stay within the current liquidity range");
        let expected_tick =
            get_tick_at_sqrt_ratio(result.sqrt_price).expect("new sqrt price should map to a tick");

        assert_ne!(result.sqrt_price, pool.sqrt_price);
        assert_ne!(result.sqrt_price, get_sqrt_ratio_at_tick(-120).unwrap());
        assert_ne!(expected_tick, pool.tick.value());
        assert_eq!(result.tick.value(), expected_tick);
    }

    #[test]
    fn test_swap_keeps_boundary_tick_when_price_does_not_move() {
        let mut pool = create_basic_test_pool();
        pool.tick = PoolTick::from(-1);
        let amount = I256::checked_from_sign_and_abs(Sign::Positive, U256::from(1u64)).unwrap();

        let result = pool
            .swap(true, amount, None)
            .expect("swap should consume the input as fee without moving price");

        assert_eq!(result.sqrt_price, pool.sqrt_price);
        assert_eq!(get_tick_at_sqrt_ratio(result.sqrt_price).unwrap(), 0);
        assert_eq!(result.tick, pool.tick);
    }

    #[test]
    fn test_swap_price_limit_out_of_range_returns_error() {
        let pool = create_basic_test_pool();
        let amount = I256::checked_from_sign_and_abs(Sign::Positive, U256::from(1000u64)).unwrap();

        let result = pool.swap(true, amount, Some(pool.sqrt_price));
        assert!(matches!(result, Err(SimulationError::InvalidInput(_, None))));

        let result = pool.swap(true, amount, Some(MIN_SQRT_RATIO));
        assert!(matches!(result, Err(SimulationError::InvalidInput(_, None))));

        let result = pool.swap(false, amount, Some(pool.sqrt_price));
        assert!(matches!(result, Err(SimulationError::InvalidInput(_, None))));

        let result = pool.swap(false, amount, Some(MAX_SQRT_RATIO));
        assert!(matches!(result, Err(SimulationError::InvalidInput(_, None))));
    }

    #[test]
    fn test_swap_at_extreme_price_returns_error() {
        let sqrt_price = MIN_SQRT_RATIO + U256::from(1u64);
        let tick = get_tick_at_sqrt_ratio(sqrt_price).expect("Failed to calculate tick");
        let ticks =
            vec![TickInfo::new(MIN_TICK, 0).unwrap(), TickInfo::new(MIN_TICK + 1, 0).unwrap()];
        let pool = VelodromeSlipstreamsState::new(
            100_000_000_000_000_000_000u128,
            sqrt_price,
            3000,
            0,
            1,
            tick,
            ticks,
        )
        .expect("Failed to create pool");

        let amount = I256::checked_from_sign_and_abs(Sign::Positive, U256::from(1000u64)).unwrap();
        let result = pool.swap(true, amount, None);
        assert!(matches!(result, Err(SimulationError::InvalidInput(_, None))));
    }

    fn token_pair() -> (Token, Token) {
        let token_x = Token::new(&Bytes::from([0x01; 20]), "X", 18, 0, &[], Chain::Ethereum, 100);
        let token_y = Token::new(&Bytes::from([0x02; 20]), "Y", 18, 0, &[], Chain::Ethereum, 100);
        (token_x, token_y)
    }

    #[rstest]
    fn test_query_pool_swap_target_price_lands_in_band(#[values(true, false)] sell_x: bool) {
        let pool = create_basic_test_pool();
        let (token_x, token_y) = token_pair();
        let (token_in, token_out) =
            if sell_x { (&token_x, &token_y) } else { (&token_y, &token_x) };
        let target = 0.995;
        let tolerance = 1e-4;
        let params = target_price_params(
            token_in,
            token_out,
            to_price(target, token_in, token_out),
            tolerance,
        );

        let swap = pool.query_pool_swap(&params).unwrap();

        let new_spot = swap
            .new_state()
            .spot_price(token_in, token_out)
            .unwrap();
        assert!(new_spot >= target && new_spot <= target * (1.0 + tolerance), "spot {new_spot}");
        assert!(swap.price_points().is_none(), "the native path returns no price points");
        let quote = pool
            .get_amount_out(swap.amount_in().clone(), token_in, token_out)
            .unwrap();
        assert_eq!(&quote.amount, swap.amount_out());
    }

    /// A target less than half the tolerance below spot puts the swap limit above spot.
    #[test]
    fn test_query_pool_swap_target_price_falls_back_to_search() {
        let pool = create_basic_test_pool();
        let (token_x, token_y) = token_pair();
        let params =
            target_price_params(&token_x, &token_y, to_price(1.0 - 1e-5, &token_x, &token_y), 1e-4);

        let swap = pool.query_pool_swap(&params).unwrap();

        assert!(swap.price_points().is_some(), "the numerical search returns price points");
    }

    /// With maximum liquidity, the native swap spends its whole input allowance far above a very
    /// low target.
    #[test]
    fn test_query_pool_swap_target_price_out_of_reach_falls_back_to_search() {
        let sqrt_price = get_sqrt_ratio_at_tick(0).unwrap();
        let ticks = vec![TickInfo::new(MIN_TICK, 0).unwrap(), TickInfo::new(-MIN_TICK, 0).unwrap()];
        let pool =
            VelodromeSlipstreamsState::new(u128::MAX, sqrt_price, 3000, 0, 1, 0, ticks).unwrap();
        let (token_x, token_y) = token_pair();
        let target = Price::new(BigUint::from(1u8), BigUint::from(1u8) << 100usize);
        let target_f64 = 2f64.powi(-100);
        let tolerance = 1e-4;
        let params = target_price_params(&token_x, &token_y, target, tolerance);

        let swap = pool.query_pool_swap(&params).unwrap();

        let new_spot = swap
            .new_state()
            .spot_price(&token_x, &token_y)
            .unwrap();
        assert!(
            new_spot >= target_f64 && new_spot <= target_f64 * (1.0 + tolerance),
            "spot {new_spot} is outside the band of target {target_f64}"
        );
        assert!(swap.price_points().is_some(), "the numerical search returns price points");
    }

    #[test]
    fn test_query_pool_swap_target_price_at_spot() {
        let pool = create_basic_test_pool();
        let (token_x, token_y) = token_pair();
        let params = target_price_params(
            &token_x,
            &token_y,
            Price::new(BigUint::from(1u64), BigUint::from(1u64)),
            1e-4,
        );

        let swap = pool.query_pool_swap(&params).unwrap();

        assert!(swap.amount_in().is_zero());
        assert!(swap.amount_out().is_zero());
        assert!(swap.new_state().eq(&pool));
    }

    #[test]
    fn test_query_pool_swap_target_price_above_spot() {
        let pool = create_basic_test_pool();
        let (token_x, token_y) = token_pair();
        let params =
            target_price_params(&token_x, &token_y, to_price(1.01, &token_x, &token_y), 1e-4);

        let result = pool.query_pool_swap(&params);

        assert!(matches!(result, Err(SimulationError::InvalidInput(..))), "got {result:?}");
    }
}

#[cfg(test)]
mod tick_list_sharing_tests {
    use std::collections::{HashMap, HashSet};

    use tycho_common::{dto::ProtocolStateDelta, hex_bytes::Bytes};

    use super::*;
    use crate::evm::protocol::utils::uniswap::tick_math::get_sqrt_ratio_at_tick;

    #[test]
    fn test_delta_transition_leaves_clones_ticks_unchanged() {
        let original = VelodromeSlipstreamsState::new(
            100_000_000_000_000_000_000u128,
            get_sqrt_ratio_at_tick(0).unwrap(),
            3000,
            0,
            1,
            0,
            vec![TickInfo::new(-120, 10000).unwrap(), TickInfo::new(120, -10000).unwrap()],
        )
        .unwrap();
        let mut updated = original.clone();
        let delta = ProtocolStateDelta {
            component_id: "State1".to_owned(),
            updated_attributes: HashMap::from([(
                "ticks/-120/net_liquidity".to_string(),
                Bytes::from(20000_i128.to_be_bytes().to_vec()),
            )]),
            deleted_attributes: HashSet::new(),
        };

        updated
            .delta_transition(delta, &HashMap::new(), &Balances::default())
            .unwrap();

        assert_eq!(
            updated
                .ticks
                .get_tick(-120)
                .unwrap()
                .net_liquidity,
            20000
        );
        assert_eq!(
            original
                .ticks
                .get_tick(-120)
                .unwrap()
                .net_liquidity,
            10000
        );
    }
}

#[cfg(test)]
mod step_cache_tests {
    use std::{
        collections::{HashMap, HashSet},
        str::FromStr,
    };

    use tycho_common::{dto::ProtocolStateDelta, hex_bytes::Bytes};

    use super::*;
    use crate::evm::protocol::utils::uniswap::swap_step_cache::test_fixtures::{
        amount_orders, test_amounts, test_ticks, wbtc_weth, LIQUIDITY, SQRT_PRICE, TICK,
    };

    fn test_pool() -> VelodromeSlipstreamsState {
        VelodromeSlipstreamsState::new(
            LIQUIDITY,
            U256::from_str(SQRT_PRICE).unwrap(),
            500,
            0,
            10,
            TICK,
            test_ticks(),
        )
        .unwrap()
    }

    /// A quote's amount, gas and new state, or its error, as text.
    fn format_quote_result(
        pool: &VelodromeSlipstreamsState,
        amount: &BigUint,
        sell: &Token,
        buy: &Token,
    ) -> String {
        match pool.get_amount_out(amount.clone(), sell, buy) {
            Ok(result) => {
                let state = result
                    .new_state
                    .as_any()
                    .downcast_ref::<VelodromeSlipstreamsState>()
                    .unwrap();
                format!(
                    "{} {} {} {} {}",
                    result.amount,
                    result.gas,
                    state.sqrt_price,
                    state.tick.value(),
                    state.liquidity
                )
            }
            Err(error) => format!("error {error}"),
        }
    }

    #[test]
    fn test_step_cache_quotes_match_a_fresh_state_in_any_order() {
        let (wbtc, weth) = wbtc_weth();
        for (sell, buy, smallest) in [(&wbtc, &weth, 1_000u64), (&weth, &wbtc, 1_000_000_000)] {
            for order in amount_orders(&test_amounts(smallest)) {
                let shared = test_pool();
                for amount in &order {
                    assert_eq!(
                        format_quote_result(&shared, amount, sell, buy),
                        format_quote_result(&test_pool(), amount, sell, buy),
                        "selling {amount} {}",
                        sell.symbol
                    );
                }
                let steps = shared
                    .step_cache
                    .min_inputs(sell < buy)
                    .len();
                assert!(steps > 2, "only {steps} steps cached");
            }
        }
    }

    #[test]
    fn test_step_cache_quotes_match_a_fresh_state_at_each_cached_step() {
        let (wbtc, weth) = wbtc_weth();
        // Amounts that run past the last tick, so the cache covers the whole tick list. Selling
        // WETH needs a far larger amount to leave its first tick.
        for (sell, buy, doublings) in [(&wbtc, &weth, 30u32), (&weth, &wbtc, 70)] {
            let cached = test_pool();
            let _ = cached.get_amount_out(BigUint::from(10u32) << doublings, sell, buy);
            let boundaries = cached.step_cache.min_inputs(sell < buy);
            assert!(boundaries.len() > 2, "only {} steps cached", boundaries.len());
            for boundary in boundaries.into_iter().skip(1) {
                for amount in [boundary - U256::from(1u64), boundary, boundary + U256::from(1u64)] {
                    let amount = u256_to_biguint(amount);
                    assert_eq!(
                        format_quote_result(&cached, &amount, sell, buy),
                        format_quote_result(&test_pool(), &amount, sell, buy),
                        "selling {amount} {}",
                        sell.symbol
                    );
                }
            }
        }
    }

    #[test]
    fn test_step_cache_of_a_swapped_state() {
        let (wbtc, weth) = wbtc_weth();
        let pool = test_pool();
        let swapped = pool
            .get_amount_out(BigUint::from(300_000_000u64), &wbtc, &weth)
            .unwrap()
            .new_state;
        let swapped = swapped
            .as_any()
            .downcast_ref::<VelodromeSlipstreamsState>()
            .unwrap();
        // A serde round trip gives the same state with an empty step cache.
        let fresh: VelodromeSlipstreamsState =
            serde_json::from_value(serde_json::to_value(swapped).unwrap()).unwrap();

        for amount in test_amounts(1_000).iter().rev() {
            assert_eq!(
                format_quote_result(swapped, amount, &wbtc, &weth),
                format_quote_result(&fresh, amount, &wbtc, &weth),
                "selling {amount} WBTC"
            );
        }
    }

    #[test]
    fn test_step_cache_restarts_after_delta_transition() {
        let (wbtc, weth) = wbtc_weth();
        let large = BigUint::from(3_000_000_000u64);
        // Only a crossed tick's liquidity changes, so the price, tick and liquidity the cache
        // starts from stay the same.
        let delta = || ProtocolStateDelta {
            component_id: "State1".to_owned(),
            updated_attributes: HashMap::from([(
                "ticks/255820/net_liquidity".to_string(),
                Bytes::from(
                    1_000_000_000_000_000i128
                        .to_be_bytes()
                        .to_vec(),
                ),
            )]),
            deleted_attributes: HashSet::new(),
        };
        let mut cached = test_pool();
        cached
            .get_amount_out(large.clone(), &wbtc, &weth)
            .unwrap();
        let mut fresh = test_pool();

        cached
            .delta_transition(delta(), &HashMap::new(), &Balances::default())
            .unwrap();
        fresh
            .delta_transition(delta(), &HashMap::new(), &Balances::default())
            .unwrap();

        assert_eq!(
            format_quote_result(&cached, &large, &wbtc, &weth),
            format_quote_result(&fresh, &large, &wbtc, &weth)
        );
    }
}