waterpump-evm-pool-sdk 0.1.0

use alloy::primitives::{
    aliases::{I24, U160},
    Address, Bytes, U256,
};
use alloy_sol_types::SolCall;
use anyhow::{Context, Result};
use uniswap_sdk_core::{
    entities::{BaseCurrency, FractionBase},
    prelude::ToBig,
    utils::FromBig,
};
use uniswap_v3_sdk::prelude::{encode_multicall, encode_permit, encode_refund_eth};
use waterpump_evm_core::pool_key::SlipstreamPoolKey;

use crate::{
    common::{mint_amounts_utils::calculate_mint_amounts, slippage_utils::ratios_after_slippage},
    pool_infusers::common::{
        calculate_position_amounts, calculate_slippage_context, AddLiquidityAmounts,
    },
    pool_swappers::common::MethodParameters,
    traits::pool_infuser::{
        AddBatchLiquidityParams, AddLiquidityOptions, AddLiquidityParams,
        AddLiquiditySpecificOptions, CollectOptions, CollectParams, RemoveBatchLiquidityParams,
        RemoveLiquidityOptions, RemoveLiquidityParams,
    },
};

// Define Slipstream INonfungiblePositionManager interface calls
alloy_sol_types::sol! {
    interface INonfungiblePositionManager {
        struct MintParams {
            address token0;
            address token1;
            int24 tickSpacing;
            int24 tickLower;
            int24 tickUpper;
            uint256 amount0Desired;
            uint256 amount1Desired;
            uint256 amount0Min;
            uint256 amount1Min;
            address recipient;
            uint256 deadline;
            uint160 sqrtPriceX96;
        }

        /// @notice Creates a new position wrapped in a NFT
        /// @dev Call this when the pool does exist and is initialized. Note that if the pool is created but not initialized
        /// a method does not exist, i.e. the pool is assumed to be initialized.
        /// @param params The params necessary to mint a position, encoded as `MintParams` in calldata
        /// @return tokenId The ID of the token that represents the minted position
        /// @return liquidity The amount of liquidity for this position
        /// @return amount0 The amount of token0
        /// @return amount1 The amount of token1
        function mint(MintParams calldata params)
            external
            payable
            returns (uint256 tokenId, uint128 liquidity, uint256 amount0, uint256 amount1);

        struct IncreaseLiquidityParams {
            uint256 tokenId;
            uint256 amount0Desired;
            uint256 amount1Desired;
            uint256 amount0Min;
            uint256 amount1Min;
            uint256 deadline;
        }

        /// @notice Increases the amount of liquidity in a position, with tokens paid by the `msg.sender`
        /// @param params tokenId The ID of the token for which liquidity is being increased,
        /// amount0Desired The desired amount of token0 to be spent,
        /// amount1Desired The desired amount of token1 to be spent,
        /// amount0Min The minimum amount of token0 to spend, which serves as a slippage check,
        /// amount1Min The minimum amount of token1 to spend, which serves as a slippage check,
        /// deadline The time by which the transaction must be included to effect the change
        /// @return liquidity The new liquidity amount as a result of the increase
        /// @return amount0 The amount of token0 to acheive resulting liquidity
        /// @return amount1 The amount of token1 to acheive resulting liquidity
        function increaseLiquidity(IncreaseLiquidityParams calldata params)
            external
            payable
            returns (uint128 liquidity, uint256 amount0, uint256 amount1);

        struct DecreaseLiquidityParams {
            uint256 tokenId;
            uint128 liquidity;
            uint256 amount0Min;
            uint256 amount1Min;
            uint256 deadline;
        }

        /// @notice Decreases the amount of liquidity in a position and accounts it to the position
        /// @param params tokenId The ID of the token for which liquidity is being decreased,
        /// amount The amount by which liquidity will be decreased,
        /// amount0Min The minimum amount of token0 that should be accounted for the burned liquidity,
        /// amount1Min The minimum amount of token1 that should be accounted for the burned liquidity,
        /// deadline The time by which the transaction must be included to effect the change
        /// @return amount0 The amount of token0 accounted to the position's tokens owed
        /// @return amount1 The amount of token1 accounted to the position's tokens owed
        function decreaseLiquidity(DecreaseLiquidityParams calldata params)
            external
            payable
            returns (uint256 amount0, uint256 amount1);

        struct CollectParams {
            uint256 tokenId;
            address recipient;
            uint128 amount0Max;
            uint128 amount1Max;
        }

        /// @notice Collects up to a maximum amount of fees owed to a specific position to the recipient
        /// @notice Used to update staked positions before deposit and withdraw
        /// @param params tokenId The ID of the NFT for which tokens are being collected,
        /// recipient The account that should receive the tokens,
        /// amount0Max The maximum amount of token0 to collect,
        /// amount1Max The maximum amount of token1 to collect
        /// @return amount0 The amount of fees collected in token0
        /// @return amount1 The amount of fees collected in token1
        function collect(CollectParams calldata params) external payable returns (uint256 amount0, uint256 amount1);

        /// @notice Burns a token ID, which deletes it from the NFT contract. The token must have 0 liquidity and all tokens
        /// must be collected first.
        /// @param tokenId The ID of the token that is being burned
        function burn(uint256 tokenId) external payable;
    }
}

/// Build calldata for mint or increase liquidity operation
fn build_add_liquidity_calldata(
    pool_key: &SlipstreamPoolKey,
    tick_lower: I24,
    tick_upper: I24,
    amounts: &AddLiquidityAmounts,
    specific_opts: AddLiquiditySpecificOptions,
    deadline: U256,
    sqrt_ratio_x96_current: U160,
) -> Bytes {
    match specific_opts {
        AddLiquiditySpecificOptions::Mint(opts) => INonfungiblePositionManager::mintCall {
            params: INonfungiblePositionManager::MintParams {
                token0: pool_key.token_a.address(),
                token1: pool_key.token_b.address(),
                tickSpacing: pool_key.tick_spacing,
                tickLower: tick_lower,
                tickUpper: tick_upper,
                amount0Desired: amounts.amount0_desired,
                amount1Desired: amounts.amount1_desired,
                amount0Min: amounts.amount0_min,
                amount1Min: amounts.amount1_min,
                recipient: opts.recipient,
                deadline,
                sqrtPriceX96: sqrt_ratio_x96_current,
            },
        }
        .abi_encode()
        .into(),
        AddLiquiditySpecificOptions::Increase(opts) => {
            INonfungiblePositionManager::increaseLiquidityCall {
                params: INonfungiblePositionManager::IncreaseLiquidityParams {
                    tokenId: opts.token_id,
                    amount0Desired: amounts.amount0_desired,
                    amount1Desired: amounts.amount1_desired,
                    amount0Min: amounts.amount0_min,
                    amount1Min: amounts.amount1_min,
                    deadline,
                },
            }
            .abi_encode()
            .into()
        }
    }
}

/// Build call parameters for adding liquidity (mint or increase)
///
/// Uses Position::mint_amounts_with_slippage approach with
/// ratios_after_slippage for accurate minimum amounts calculation.
pub fn build_add_liquidity_call_parameters(
    pool_key: &SlipstreamPoolKey,
    params: AddLiquidityParams,
    options: AddLiquidityOptions,
) -> Result<MethodParameters> {
    let mut calldatas: Vec<Bytes> = Vec::with_capacity(5);

    let deadline = options.deadline;
    let amount0_desired = U256::from_big_int(params.amount0.quotient());
    let amount1_desired = U256::from_big_int(params.amount1.quotient());

    // Calculate slippage context
    let slippage_ctx =
        calculate_slippage_context(&params.token0_price, &options.slippage_tolerance)?;

    // Calculate position amounts with slippage
    let amounts = calculate_position_amounts(
        params.tick_lower,
        params.tick_upper,
        amount0_desired,
        amount1_desired,
        &slippage_ctx,
    )?;

    // Create pool if needed (for mint)
    if let AddLiquiditySpecificOptions::Mint(opts) = params.specific_opts {
        if opts.create_pool {
            // Slipstream may have a different pool creation mechanism
            // For now, we'll skip this as it's typically not needed
            // TODO: Implement Slipstream pool creation if needed
        }
    }

    // Mint or increase liquidity
    calldatas.push(build_add_liquidity_calldata(
        pool_key,
        params.tick_lower,
        params.tick_upper,
        &amounts,
        params.specific_opts,
        deadline,
        slippage_ctx.sqrt_ratio_x96_current,
    ));

    // Handle permits if needed
    if let Some(token0_permit) = options.token0_permit {
        calldatas.push(encode_permit(&pool_key.token_a, token0_permit));
    }

    if let Some(token1_permit) = options.token1_permit {
        calldatas.push(encode_permit(&pool_key.token_b, token1_permit));
    }

    // Wrap ETH if needed (for native token)
    if options.use_native.is_some() {
        calldatas.push(encode_refund_eth());
    }

    Ok(MethodParameters { calldata: encode_multicall(calldatas), value: U256::ZERO })
}

/// Build call parameters for adding batch liquidity (multiple positions in one
/// transaction)
///
/// Uses multicall to batch multiple mint operations together for gas
/// efficiency. Uses a shared token0_price for all positions.
pub fn build_add_batch_liquidity_call_parameters(
    pool_key: &SlipstreamPoolKey,
    params: AddBatchLiquidityParams,
    options: AddLiquidityOptions,
) -> Result<MethodParameters> {
    let mut calldatas: Vec<Bytes> = Vec::with_capacity(params.items.len() + 4);

    let deadline = options.deadline;

    // Calculate shared slippage context once
    let slippage_ctx =
        calculate_slippage_context(&params.token0_price, &options.slippage_tolerance)?;

    // Permits if necessary (once for all positions)
    if let Some(permit) = options.token0_permit {
        calldatas.push(encode_permit(&pool_key.token_a, permit));
    }
    if let Some(permit) = options.token1_permit {
        calldatas.push(encode_permit(&pool_key.token_b, permit));
    }

    for item in &params.items {
        let amount0_desired = U256::from_big_int(item.amount0.quotient());
        let amount1_desired = U256::from_big_int(item.amount1.quotient());

        // Calculate position amounts with slippage
        let amounts = calculate_position_amounts(
            item.tick_lower,
            item.tick_upper,
            amount0_desired,
            amount1_desired,
            &slippage_ctx,
        )?;

        // Mint or increase liquidity
        calldatas.push(build_add_liquidity_calldata(
            pool_key,
            item.tick_lower,
            item.tick_upper,
            &amounts,
            item.specific_opts,
            deadline,
            slippage_ctx.sqrt_ratio_x96_current,
        ));
    }

    // Wrap ETH if needed (for native token)
    if options.use_native.is_some() {
        calldatas.push(encode_refund_eth());
    }

    Ok(MethodParameters { calldata: encode_multicall(calldatas), value: U256::ZERO })
}

/// Parameters for building remove liquidity calldata for a single position
struct RemovePositionParams {
    token_id: U256,
    liquidity: U256,
    tick_lower: I24,
    tick_upper: I24,
    recipient: Address,
    deadline: U256,
    burn_token: bool,
}

/// Build calldatas for removing liquidity from a single position (internal
/// helper)
///
/// Generates decrease liquidity, collect, and optionally burn calldatas.
fn build_remove_position_calldatas(
    pool_key: &SlipstreamPoolKey,
    params: &RemovePositionParams,
    sqrt_ratio_x96_lower_slippage: U160,
    sqrt_ratio_x96_upper_slippage: U160,
    collect_options: &CollectOptions,
) -> Result<Vec<Bytes>> {
    let mut calldatas: Vec<Bytes> = Vec::with_capacity(4);

    // Calculate minimum amounts with slippage following
    // Position::burn_amounts_with_slippage:
    // - amount0 is smaller when price is higher (use upper slippage price)
    // - amount1 is smaller when price is lower (use lower slippage price)
    let mint_amounts_upper = calculate_mint_amounts(
        params.tick_lower,
        params.tick_upper,
        params.liquidity.to::<u128>(),
        sqrt_ratio_x96_upper_slippage,
    )
    .context("Failed to calculate mint amounts at upper slippage price")?;

    let mint_amounts_lower = calculate_mint_amounts(
        params.tick_lower,
        params.tick_upper,
        params.liquidity.to::<u128>(),
        sqrt_ratio_x96_lower_slippage,
    )
    .context("Failed to calculate mint amounts at lower slippage price")?;

    let amount0_min = mint_amounts_upper.amount0;
    let amount1_min = mint_amounts_lower.amount1;

    // Decrease liquidity
    calldatas.push(
        INonfungiblePositionManager::decreaseLiquidityCall {
            params: INonfungiblePositionManager::DecreaseLiquidityParams {
                tokenId: params.token_id,
                liquidity: params.liquidity.to::<u128>(),
                amount0Min: U256::from(amount0_min),
                amount1Min: U256::from(amount1_min),
                deadline: params.deadline,
            },
        }
        .abi_encode()
        .into(),
    );

    // Collect
    let collect_params = CollectParams { token_id: params.token_id, recipient: params.recipient };
    let collect_calldatas = build_collect_calldatas(pool_key, &collect_params, collect_options)?;
    calldatas.extend(collect_calldatas);

    // Burn token if needed
    if params.burn_token {
        calldatas.push(
            INonfungiblePositionManager::burnCall { tokenId: params.token_id }.abi_encode().into(),
        );
    }

    Ok(calldatas)
}

/// Build call parameters for removing liquidity
pub fn build_remove_liquidity_call_parameters(
    pool_key: &SlipstreamPoolKey,
    params: RemoveLiquidityParams,
    options: RemoveLiquidityOptions,
) -> Result<MethodParameters> {
    let mut calldatas: Vec<Bytes> = Vec::with_capacity(6);

    // NFT permit if needed
    if let Some(permit) = options.permit {
        // Define IERC721Permit interface for Slipstream
        alloy_sol_types::sol! {
            interface IERC721Permit {
                function permit(
                    address spender,
                    uint256 tokenId,
                    uint256 deadline,
                    uint8 v,
                    bytes32 r,
                    bytes32 s
                ) external;
            }
        }
        calldatas.push(
            IERC721Permit::permitCall {
                spender: permit.spender,
                tokenId: params.token_id,
                deadline: permit.deadline,
                v: permit.signature.v() as u8 + 27,
                r: permit.signature.r().into(),
                s: permit.signature.s().into(),
            }
            .abi_encode()
            .into(),
        );
    }

    // Get slippage-adjusted sqrt ratios
    let (sqrt_ratio_x96_lower_slippage, sqrt_ratio_x96_upper_slippage) =
        ratios_after_slippage(&params.token0_price, &options.slippage_tolerance)
            .context("Failed to calculate ratios after slippage")?;

    // Build remove position calldatas
    let remove_params = RemovePositionParams {
        token_id: params.token_id,
        liquidity: params.liquidity,
        tick_lower: params.tick_lower,
        tick_upper: params.tick_upper,
        recipient: params.recipient,
        deadline: params.deadline,
        burn_token: options.burn_token,
    };
    let position_calldatas = build_remove_position_calldatas(
        pool_key,
        &remove_params,
        sqrt_ratio_x96_lower_slippage,
        sqrt_ratio_x96_upper_slippage,
        &options.collect_options,
    )?;
    calldatas.extend(position_calldatas);

    Ok(MethodParameters { calldata: encode_multicall(calldatas), value: U256::ZERO })
}

/// Build call parameters for batch removing liquidity (multiple positions in
/// one transaction)
///
/// Uses multicall to batch multiple decrease liquidity operations together for
/// gas efficiency. Uses a shared token0_price for all positions.
pub fn build_remove_batch_liquidity_call_parameters(
    pool_key: &SlipstreamPoolKey,
    params: RemoveBatchLiquidityParams,
    options: RemoveLiquidityOptions,
) -> Result<MethodParameters> {
    let mut calldatas: Vec<Bytes> = Vec::with_capacity(params.items.len() * 3 + 1);

    // Get slippage-adjusted sqrt ratios (shared for all positions)
    let (sqrt_ratio_x96_lower_slippage, sqrt_ratio_x96_upper_slippage) =
        ratios_after_slippage(&params.token0_price, &options.slippage_tolerance)
            .context("Failed to calculate ratios after slippage")?;

    for item in &params.items {
        let remove_params = RemovePositionParams {
            token_id: item.token_id,
            liquidity: item.liquidity,
            tick_lower: item.tick_lower,
            tick_upper: item.tick_upper,
            recipient: params.recipient,
            deadline: params.deadline,
            burn_token: item.burn_token,
        };
        let position_calldatas = build_remove_position_calldatas(
            pool_key,
            &remove_params,
            sqrt_ratio_x96_lower_slippage,
            sqrt_ratio_x96_upper_slippage,
            &options.collect_options,
        )?;
        calldatas.extend(position_calldatas);
    }

    Ok(MethodParameters { calldata: encode_multicall(calldatas), value: U256::ZERO })
}

/// Build calldatas for collecting fees (internal helper)
fn build_collect_calldatas(
    _pool_key: &SlipstreamPoolKey,
    params: &CollectParams,
    options: &CollectOptions,
) -> Result<Vec<Bytes>> {
    let mut calldatas: Vec<Bytes> = Vec::with_capacity(1);

    // Convert CurrencyAmount to u128 for amount0Max and amount1Max
    // Use max values if not specified (u128::MAX)
    let amount0_max = options
        .expected_currency_owed0
        .as_ref()
        .map(|c| c.quotient().to_big_uint())
        .unwrap_or_else(|| U256::from(u128::MAX).to_big_uint());
    let amount1_max = options
        .expected_currency_owed1
        .as_ref()
        .map(|c| c.quotient().to_big_uint())
        .unwrap_or_else(|| U256::from(u128::MAX).to_big_uint());
    // Convert to u128 (Slipstream uses u128 for collect amounts)
    let amount0_max_u128 = U256::from_big_uint(amount0_max).to::<u128>();
    let amount1_max_u128 = U256::from_big_uint(amount1_max).to::<u128>();

    // Collect
    calldatas.push(
        INonfungiblePositionManager::collectCall {
            params: INonfungiblePositionManager::CollectParams {
                tokenId: params.token_id,
                recipient: params.recipient,
                amount0Max: amount0_max_u128,
                amount1Max: amount1_max_u128,
            },
        }
        .abi_encode()
        .into(),
    );

    Ok(calldatas)
}

/// Build call parameters for collecting fees
pub fn build_collect_call_parameters(
    pool_key: &SlipstreamPoolKey,
    params: &CollectParams,
    options: &CollectOptions,
) -> Result<MethodParameters> {
    let calldatas = build_collect_calldatas(pool_key, params, options)?;
    Ok(MethodParameters { calldata: encode_multicall(calldatas), value: U256::ZERO })
}