waterpump-evm-pool-sdk 0.1.0

use alloy::primitives::{
    aliases::{I24, U160},
    U256,
};
use anyhow::{Context, Result};
use tracing::debug;
use uniswap_sdk_core::{
    entities::BaseCurrency,
    prelude::{Currency, CurrencyAmount, Price},
};

/// Trait for pool keys that have token_a and token_b fields
pub trait PoolKey {
    fn token_a(&self) -> &Currency;
    fn token_b(&self) -> &Currency;
}

// Implement PoolKey for common pool key types
impl PoolKey for crate::types::v3_pool_key::V3PoolKey {
    fn token_a(&self) -> &Currency { &self.token_a }

    fn token_b(&self) -> &Currency { &self.token_b }
}

impl PoolKey for waterpump_evm_core::pool_key::SlipstreamPoolKey {
    fn token_a(&self) -> &Currency { &self.token_a }

    fn token_b(&self) -> &Currency { &self.token_b }
}

impl PoolKey for waterpump_evm_core::pool_key::V4PoolKey {
    fn token_a(&self) -> &Currency { &self.token_a }

    fn token_b(&self) -> &Currency { &self.token_b }
}

/// Merge two pool prices based on how the pools are connected
///
/// This function determines the connection path between two pools and merges
/// their prices accordingly. The pools must share a common token (either
/// pool_a.token_b == pool_b.token_a, etc.)
///
/// # Arguments
/// * `pool_a_key` - The first pool key with token_a and token_b
/// * `pool_b_key` - The second pool key with token_a and token_b
/// * `price_a` - The price from pool_a (token_a/token_b)
/// * `price_b` - The price from pool_b (token_a/token_b)
///
/// # Returns
/// The merged price representing the path through both pools
pub fn merge_two_pool_prices(
    pool_a_key: &impl PoolKey,
    pool_b_key: &impl PoolKey,
    price_a: &Price<Currency, Currency>,
    price_b: &Price<Currency, Currency>,
) -> Result<Price<Currency, Currency>> {
    // Determine the path and merge prices
    // Pool_a: token_a_a/token_b_a, Pool_b: token_a_b/token_b_b
    // We need to find the connection and multiply prices appropriately
    let merged_price = if pool_a_key.token_b().address() == pool_b_key.token_a().address() {
        // Path: token_a_a -> token_b_a (== token_a_b) -> token_b_b
        price_a.multiply(price_b).context("Failed to multiply prices")?
    } else if pool_a_key.token_b().address() == pool_b_key.token_b().address() {
        // Path: token_a_a -> token_b_a (== token_b_b) -> token_a_b
        let price_b_inv = price_b.invert();
        price_a.multiply(&price_b_inv).context("Failed to multiply prices")?
    } else if pool_a_key.token_a().address() == pool_b_key.token_a().address() {
        // Path: token_b_a -> token_a_a (== token_a_b) -> token_b_b
        let price_a_inv = price_a.invert();
        price_a_inv.multiply(price_b).context("Failed to multiply prices")?
    } else if pool_a_key.token_a().address() == pool_b_key.token_b().address() {
        // Path: token_b_a -> token_a_a (== token_b_b) -> token_a_b
        let price_a_inv = price_a.invert();
        let price_b_inv = price_b.invert();
        price_a_inv.multiply(&price_b_inv).context("Failed to multiply prices")?
    } else {
        anyhow::bail!(
            "Pools are not connected: pool_a ({:?}/{:?}) and pool_b ({:?}/{:?})",
            pool_a_key.token_a().address(),
            pool_a_key.token_b().address(),
            pool_b_key.token_a().address(),
            pool_b_key.token_b().address()
        );
    };

    let merged_price_str =
        merged_price.to_significant(8, None).unwrap_or_else(|_| "N/A".to_string());
    debug!(
        merged_price = %merged_price_str,
        "Prices merged successfully"
    );

    Ok(merged_price)
}

/// Calculate the total value of a position in a target currency
///
/// This function takes token amounts, fees, and a price, and calculates the
/// total value in the target currency by converting the non-target currency
/// amounts using the price.
///
/// # Arguments
/// * `amount0` - Amount of token0 in the position
/// * `amount1` - Amount of token1 in the position
/// * `fee0` - Uncollected fees in token0 (optional)
/// * `fee1` - Uncollected fees in token1 (optional)
/// * `price` - Price from token0 to token1 (Price<Currency0, Currency1>)
/// * `target_currency0` - If true, calculate value in currency0; if false,
///   calculate in currency1
///
/// # Returns
/// Total value of the position in the target currency
pub fn calculate_position_value(
    amount0: CurrencyAmount<Currency>,
    amount1: CurrencyAmount<Currency>,
    fee0: Option<CurrencyAmount<Currency>>,
    fee1: Option<CurrencyAmount<Currency>>,
    price: Price<Currency, Currency>,
    target_currency0: bool,
) -> Result<CurrencyAmount<Currency>> {
    // Sum amounts and fees for each currency (fees default to zero if None)
    let total_amount0 = if let Some(fee) = fee0 {
        amount0.add(&fee).map_err(|e| anyhow::anyhow!("Failed to add amount0 and fee0: {:?}", e))?
    } else {
        amount0
    };
    let total_amount1 = if let Some(fee) = fee1 {
        amount1.add(&fee).map_err(|e| anyhow::anyhow!("Failed to add amount1 and fee1: {:?}", e))?
    } else {
        amount1
    };

    let total_value =
        calculate_value_in_target_currency(total_amount0, total_amount1, price, target_currency0)
            .map_err(|e| anyhow::anyhow!("Failed to calculate position value: {:?}", e))?;

    Ok(total_value)
}

/// Calculate the value of a position in a target currency
///
/// This function takes token amounts, and a price, and calculates the value
/// in the target currency by converting the non-target currency amounts using
/// the price.
///
/// # Arguments
/// * `amount0` - Amount of token0 in the position
/// * `amount1` - Amount of token1 in the position
/// * `price` - Price from token0 to token1 (Price<Currency0, Currency1>)
/// * `target_currency0` - If true, calculate value in currency0; if false,
///   calculate in currency1
///
/// # Returns
/// Total value of the position in the target currency
pub fn calculate_value_in_target_currency(
    amount0: CurrencyAmount<Currency>,
    amount1: CurrencyAmount<Currency>,
    price: Price<Currency, Currency>,
    target_currency0: bool,
) -> Result<CurrencyAmount<Currency>> {
    if target_currency0 {
        // Total value in currency0 = amount0 + (amount1) converted to currency0
        // price is Price<Currency0, Currency1>, so we need inverted price to convert
        // currency1 -> currency0
        let inverted_price = price.invert();
        let amount1_in_currency0 = inverted_price
            .quote(&amount1)
            .map_err(|e| anyhow::anyhow!("Failed to quote amount1 to currency0: {:?}", e))?;
        amount0
            .add(&amount1_in_currency0)
            .map_err(|e| anyhow::anyhow!("Failed to add amounts in currency0: {:?}", e))
    } else {
        // Total value in currency1 = amount1 + (amount0) converted to currency1
        // price is Price<Currency0, Currency1>, so we can use it directly to convert
        // currency0 -> currency1
        let amount0_in_currency1 = price
            .quote(&amount0)
            .map_err(|e| anyhow::anyhow!("Failed to quote amount0 to currency1: {:?}", e))?;
        amount1
            .add(&amount0_in_currency1)
            .map_err(|e| anyhow::anyhow!("Failed to add amounts in currency1: {:?}", e))
    }
}

/// Calculate fee growth inside a position's tick range and compute tokens owed
///
/// This function implements the fee growth calculation logic from Uniswap
/// V3/V4. Reference: https://github.com/Uniswap/v4-core/blob/f630c8ca8c669509d958353200953762fd15761a/contracts/libraries/Pool.sol#L566
///
/// # Arguments
/// * `tick` - Current tick of the pool
/// * `tick_lower` - Lower tick of the position
/// * `tick_upper` - Upper tick of the position
/// * `fee_growth_global_0x128` - Global fee growth for token0 (Q128 format)
/// * `fee_growth_global_1x128` - Global fee growth for token1 (Q128 format)
/// * `fee_growth_outside_0x128_lower` - Fee growth outside for token0 at lower
///   tick
/// * `fee_growth_outside_1x128_lower` - Fee growth outside for token1 at lower
///   tick
/// * `fee_growth_outside_0x128_upper` - Fee growth outside for token0 at upper
///   tick
/// * `fee_growth_outside_1x128_upper` - Fee growth outside for token1 at upper
///   tick
/// * `fee_growth_inside_0_last_x128` - Last recorded fee growth inside for
///   token0
/// * `fee_growth_inside_1_last_x128` - Last recorded fee growth inside for
///   token1
/// * `liquidity` - Liquidity of the position
///
/// # Returns
/// A tuple of (tokens_owed_0, tokens_owed_1) representing the uncollected fees
#[allow(clippy::too_many_arguments)]
pub fn calculate_tokens_owed_unrealized(
    tick: I24,
    tick_lower: I24,
    tick_upper: I24,
    fee_growth_global_0x128: U256,
    fee_growth_global_1x128: U256,
    fee_growth_outside_0x128_lower: U256,
    fee_growth_outside_1x128_lower: U256,
    fee_growth_outside_0x128_upper: U256,
    fee_growth_outside_1x128_upper: U256,
    fee_growth_inside_0_last_x128: U256,
    fee_growth_inside_1_last_x128: U256,
    liquidity: u128,
) -> Result<(U256, U256)> {
    use uniswap_v3_sdk::prelude::get_tokens_owed;

    // Convert I24 to i32 for comparison
    let tick_bytes = tick.to_le_bytes::<3>();
    let tick_lower_bytes = tick_lower.to_le_bytes::<3>();
    let tick_upper_bytes = tick_upper.to_le_bytes::<3>();

    let sign_extend = |bytes: [u8; 3]| -> [u8; 4] {
        let sign_byte = if bytes[2] & 0x80 != 0 { 0xFF } else { 0x00 };
        [bytes[0], bytes[1], bytes[2], sign_byte]
    };

    let tick_i32 = i32::from_le_bytes(sign_extend(tick_bytes));
    let tick_lower_i32 = i32::from_le_bytes(sign_extend(tick_lower_bytes));
    let tick_upper_i32 = i32::from_le_bytes(sign_extend(tick_upper_bytes));
    // Calculate fee growth inside based on tick position
    // Reference: uniswap-v3-sdk/src/extensions/position.rs and Uniswap V3 core
    // contract This matches the exact logic from the official Uniswap V3 SDK
    // Uniswap V3 fee growth accumulators use modular (wrapping) 2^256 arithmetic.
    // Plain subtraction would panic on underflow; wrapping_sub is required.
    let (fee_growth_inside_0x128, fee_growth_inside_1x128) = if tick_i32 < tick_lower_i32 {
        (
            fee_growth_outside_0x128_lower.wrapping_sub(fee_growth_outside_0x128_upper),
            fee_growth_outside_1x128_lower.wrapping_sub(fee_growth_outside_1x128_upper),
        )
    } else if tick_i32 >= tick_upper_i32 {
        (
            fee_growth_outside_0x128_upper.wrapping_sub(fee_growth_outside_0x128_lower),
            fee_growth_outside_1x128_upper.wrapping_sub(fee_growth_outside_1x128_lower),
        )
    } else {
        (
            fee_growth_global_0x128
                .wrapping_sub(fee_growth_outside_0x128_lower)
                .wrapping_sub(fee_growth_outside_0x128_upper),
            fee_growth_global_1x128
                .wrapping_sub(fee_growth_outside_1x128_lower)
                .wrapping_sub(fee_growth_outside_1x128_upper),
        )
    };
    // Calculate tokens owed using the fee growth difference
    let (tokens_owed_0, tokens_owed_1) = get_tokens_owed(
        fee_growth_inside_0_last_x128,
        fee_growth_inside_1_last_x128,
        liquidity,
        fee_growth_inside_0x128,
        fee_growth_inside_1x128,
    );

    Ok((tokens_owed_0, tokens_owed_1))
}

/// Calculate token amounts from liquidity and tick range using Uniswap V3 SDK
/// functions
///
/// This is a pure function that calculates the token0 and token1 amounts in a
/// position based on the liquidity, tick range, and current pool state.
///
/// # Arguments
/// * `liquidity` - The liquidity value of the position
/// * `tick_lower` - The lower tick of the position's price range
/// * `tick_upper` - The upper tick of the position's price range
/// * `current_tick` - The current tick of the pool
/// * `sqrt_price_x96` - The current sqrt price of the pool (Q96 format)
/// * `token0` - The token0 currency
/// * `token1` - The token1 currency
///
/// # Returns
/// A tuple of (amount0, amount1) as CurrencyAmount representing the token
/// amounts in the position
pub fn calculate_position_token_amounts(
    liquidity: u128,
    tick_lower: I24,
    tick_upper: I24,
    current_tick: I24,
    sqrt_price_x96: U160,
) -> Result<(U256, U256)> {
    use waterpump_evm_amm_math::{sqrt_price_math, tick_math::get_sqrt_ratio_at_tick};

    // Convert I24 ticks to i32 for evm-amm-math
    let i24_to_i32 = |tick: I24| -> i32 {
        let bytes = tick.to_le_bytes::<3>();
        let sign_byte = if bytes[2] & 0x80 != 0 { 0xFF } else { 0x00 };
        i32::from_le_bytes([bytes[0], bytes[1], bytes[2], sign_byte])
    };

    let tick_lower_i32 = i24_to_i32(tick_lower);
    let tick_upper_i32 = i24_to_i32(tick_upper);
    let current_tick_i32 = i24_to_i32(current_tick);

    // Get sqrt prices for ticks using evm-amm-math (returns U256)
    let sqrt_price_lower = get_sqrt_ratio_at_tick(tick_lower_i32)?;
    let sqrt_price_upper = get_sqrt_ratio_at_tick(tick_upper_i32)?;

    // Convert sqrt_price_x96 (U160) to U256 for evm-amm-math
    let sqrt_price_current = U256::from(sqrt_price_x96);

    // Calculate amounts based on current tick position using get_amount_0_delta and
    // get_amount_1_delta These functions return U256, so we need to handle that
    let (amount0_raw, amount1_raw) = if current_tick_i32 < tick_lower_i32 {
        // Position is entirely in token0 (price is below the range)
        let amount0 = sqrt_price_math::get_amount_0_delta(
            sqrt_price_lower,
            sqrt_price_upper,
            liquidity,
            false, // round_up = false for exact calculation
        )?;
        (amount0, U256::from(0u128))
    } else if current_tick_i32 >= tick_upper_i32 {
        // Position is entirely in token1 (price is above the range)
        let amount1 = sqrt_price_math::get_amount_1_delta(
            sqrt_price_lower,
            sqrt_price_upper,
            liquidity,
            false, // round_up = false for exact calculation
        )?;
        (U256::from(0u128), amount1)
    } else {
        // Position spans the current price
        let amount0 = sqrt_price_math::get_amount_0_delta(
            sqrt_price_current,
            sqrt_price_upper,
            liquidity,
            false,
        )?;

        let amount1 = sqrt_price_math::get_amount_1_delta(
            sqrt_price_lower,
            sqrt_price_current,
            liquidity,
            false,
        )?;

        (amount0, amount1)
    };

    Ok((amount0_raw, amount1_raw))
}

#[cfg(test)]
mod tests {
    use alloy::primitives::{aliases::I24, U256};

    use super::*;

    /// Helper to create an I24 from an i32 value.
    fn i24_from_i32(val: i32) -> I24 {
        let bytes = val.to_le_bytes();
        I24::from_le_bytes::<3>([bytes[0], bytes[1], bytes[2]])
    }

    #[test]
    fn test_fee_growth_wrapping_tick_below_lower() {
        // Scenario: tick < tick_lower, and lower < upper in raw value
        // so plain subtraction would underflow.
        // fee_growth_outside_lower = 10, fee_growth_outside_upper = 20
        // wrapping result: 10 - 20 = U256::MAX - 9
        let result = calculate_tokens_owed_unrealized(
            i24_from_i32(-100),  // tick (below lower)
            i24_from_i32(-50),   // tick_lower
            i24_from_i32(50),    // tick_upper
            U256::from(1000u64), // fee_growth_global_0
            U256::from(1000u64), // fee_growth_global_1
            U256::from(10u64),   // fee_growth_outside_0_lower
            U256::from(10u64),   // fee_growth_outside_1_lower
            U256::from(20u64),   // fee_growth_outside_0_upper
            U256::from(20u64),   // fee_growth_outside_1_upper
            U256::ZERO,          // fee_growth_inside_0_last
            U256::ZERO,          // fee_growth_inside_1_last
            0u128,               // liquidity (zero so tokens_owed = 0)
        );
        // Should not panic - the wrapping subtraction handles the underflow
        assert!(result.is_ok());
    }

    #[test]
    fn test_fee_growth_wrapping_tick_above_upper() {
        // Scenario: tick >= tick_upper, upper < lower in raw value
        // plain subtraction would underflow.
        let result = calculate_tokens_owed_unrealized(
            i24_from_i32(100),   // tick (above upper)
            i24_from_i32(-50),   // tick_lower
            i24_from_i32(50),    // tick_upper
            U256::from(1000u64), // fee_growth_global_0
            U256::from(1000u64), // fee_growth_global_1
            U256::from(30u64),   // fee_growth_outside_0_lower
            U256::from(30u64),   // fee_growth_outside_1_lower
            U256::from(5u64),    // fee_growth_outside_0_upper (less than lower)
            U256::from(5u64),    // fee_growth_outside_1_upper
            U256::ZERO,          // fee_growth_inside_0_last
            U256::ZERO,          // fee_growth_inside_1_last
            0u128,
        );
        assert!(result.is_ok());
    }

    #[test]
    fn test_fee_growth_wrapping_tick_in_range() {
        // Scenario: tick in range, global < lower + upper
        // plain subtraction would underflow.
        let result = calculate_tokens_owed_unrealized(
            i24_from_i32(0),    // tick (in range)
            i24_from_i32(-50),  // tick_lower
            i24_from_i32(50),   // tick_upper
            U256::from(10u64),  // fee_growth_global_0 (small)
            U256::from(10u64),  // fee_growth_global_1
            U256::from(100u64), // fee_growth_outside_0_lower (large)
            U256::from(100u64), // fee_growth_outside_1_lower
            U256::from(200u64), // fee_growth_outside_0_upper (large)
            U256::from(200u64), // fee_growth_outside_1_upper
            U256::ZERO,         // fee_growth_inside_0_last
            U256::ZERO,         // fee_growth_inside_1_last
            0u128,
        );
        assert!(result.is_ok());
    }

    #[test]
    fn test_fee_growth_normal_no_wrap() {
        // Scenario: tick in range, normal case where no wrapping needed.
        // global=1000, lower=100, upper=200 => inside = 1000-100-200 = 700
        // last = 0, liquidity = 1
        // tokens_owed = (700 - 0) * 1 / 2^128
        // With these small values, tokens_owed = 0 (integer division)
        let result = calculate_tokens_owed_unrealized(
            i24_from_i32(0),
            i24_from_i32(-50),
            i24_from_i32(50),
            U256::from(1000u64),
            U256::from(2000u64),
            U256::from(100u64),
            U256::from(200u64),
            U256::from(200u64),
            U256::from(300u64),
            U256::ZERO,
            U256::ZERO,
            1u128,
        );
        assert!(result.is_ok());
        let (owed0, owed1) = result.unwrap();
        // fee_growth_inside_0 = 1000 - 100 - 200 = 700
        // fee_growth_inside_1 = 2000 - 200 - 300 = 1500
        // tokens_owed = (700 * 1) / 2^128 = 0 (too small)
        assert_eq!(owed0, U256::ZERO);
        assert_eq!(owed1, U256::ZERO);
    }
}