waterpump-evm-pool-sdk 0.1.0

use std::time::{SystemTime, UNIX_EPOCH};

use alloy::primitives::{aliases::U24, Address, Bytes, U256};
use alloy_sol_types::SolCall;
use anyhow::{Context, Result};
use tracing::{debug, info};
use uniswap_sdk_core::{
    entities::{BaseCurrency, BaseCurrencyCore, FractionBase},
    prelude::{BigInt, TradeType},
    utils::FromBig,
};
use uniswap_v3_sdk::prelude::{encode_multicall, encode_permit, encode_refund_eth, FeeAmount};

use crate::{
    contract_libs::iswap_router::ISwapRouter,
    pool_swappers::common::{determine_two_hop_path_tokens, MethodParameters},
    types::{
        swap_params::{SwapOptions, SwapParams},
        v3_pool_key::V3PoolKey,
    },
};

/// Encode a two-hop path for V3 quoter/router
/// Path format: token0 (20 bytes) + fee0 (3 bytes, packed as uint24) + token1
/// (20 bytes) + fee1 (3 bytes) + token2 (20 bytes) For ExactOutput, the path
/// should be in reverse order
fn encode_two_hop_path_v3(
    token0: Address,
    fee0: FeeAmount,
    token1: Address,
    fee1: FeeAmount,
    token2: Address,
) -> Bytes {
    let mut path = Vec::with_capacity(20 + 3 + 20 + 3 + 20);

    // token0 (20 bytes)
    path.extend_from_slice(token0.as_slice());
    info!(
        token0 = ?token0,
        fee0 = ?to_u24(fee0),
        token1 = ?token1,
        fee1 = ?to_u24(fee1),
        token2 = ?token2,
        "Encoded two-hop path for PancakeSwap V3"
    );

    // fee0 (3 bytes as uint24)
    // FeeAmount implements Into<U24> which can be converted to u32
    let fee0_u24: U24 = to_u24(fee0);

    path.extend_from_slice(&fee0_u24.to_be_bytes::<3>());

    // token1 (20 bytes)
    path.extend_from_slice(token1.as_slice());

    // fee1 (3 bytes as uint24)
    let fee1_u24: U24 = to_u24(fee1);

    path.extend_from_slice(&fee1_u24.to_be_bytes::<3>());

    // token2 (20 bytes)
    path.extend_from_slice(token2.as_slice());

    path.into()
}

#[inline]
#[tracing::instrument(skip(pool_key, swap_params, swap_options), fields(token_a = ?pool_key.token_a.address(), token_b = ?pool_key.token_b.address(), fee = ?pool_key.fee, is_to_b = ?swap_params.is_to_b, trade_type = ?swap_params.trade_type))]
pub fn build_swap_call_parameters(
    pool_key: &V3PoolKey,
    swap_params: SwapParams,
    swap_options: SwapOptions,
) -> Result<MethodParameters> {
    let SwapOptions {
        amount_in_maximum,
        amount_out_minimum,
        recipient,
        input_token_permit,
        sqrt_price_limit_x96,
        fee: _fee,
    } = swap_options;

    if pool_key.token_a == pool_key.token_b {
        return Err(anyhow::anyhow!("Token A and Token B cannot be the same"));
    }
    let (token_in, token_out) = if swap_params.is_to_b {
        (pool_key.token_a.clone(), pool_key.token_b.clone())
    } else {
        (pool_key.token_b.clone(), pool_key.token_a.clone())
    };

    debug!(
        token_in = ?token_in.address(),
        token_out = ?token_out.address(),
        is_to_b = ?swap_params.is_to_b,
        "Determined swap direction for PancakeSwap V3"
    );

    let trade_type = swap_params.trade_type;

    // Get deadline (current time + 20 minutes)
    let deadline = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .context("Failed to get current time")?
        .as_secs()
        + 1200; // 20 minutes

    let mut calldatas: Vec<Bytes> = Vec::with_capacity(1 + 3);

    // encode permit if necessary
    if let Some(input_token_permit) = input_token_permit {
        assert!(!token_in.is_native(), "NON_TOKEN_PERMIT");
        calldatas.push(encode_permit(&token_in, input_token_permit));
    }

    let amount_in_max_clone = amount_in_maximum.clone();
    let amount_out_min_clone = amount_out_minimum.clone();

    let swap_calldata = match trade_type {
        TradeType::ExactInput => {
            let amount_out_minimum = if let Some(amount_out_minimum) = amount_out_min_clone {
                amount_out_minimum.quotient()
            } else {
                BigInt::ZERO
            };
            let amount_in = swap_params.amount.quotient();
            debug!(
                amount_in = ?amount_in,
                sqrt_price_limit_x96 = ?sqrt_price_limit_x96,
                amount_out_minimum = ?amount_out_minimum,
                deadline = ?deadline,
                "Building PancakeSwap V3 ExactInput swap"
            );

            ISwapRouter::exactInputSingleCall {
                params: ISwapRouter::ExactInputSingleParams {
                    tokenIn: token_in.address(),
                    tokenOut: token_out.address(),
                    fee: pool_key.fee.into(),
                    recipient,
                    deadline: U256::from(deadline),
                    amountIn: U256::from_big_int(swap_params.amount.quotient()),
                    amountOutMinimum: U256::from_big_int(amount_out_minimum),
                    sqrtPriceLimitX96: sqrt_price_limit_x96.unwrap_or_default(),
                },
            }
            .abi_encode()
            .into()
        }
        TradeType::ExactOutput => {
            let amount_in_maximum = if let Some(amount_in_maximum) = amount_in_max_clone {
                amount_in_maximum.quotient()
            } else {
                BigInt::MAX
            };

            ISwapRouter::exactOutputSingleCall {
                params: ISwapRouter::ExactOutputSingleParams {
                    tokenIn: token_in.address(),
                    tokenOut: token_out.address(),
                    fee: pool_key.fee.into(),
                    recipient,
                    deadline: U256::from(deadline),
                    amountOut: U256::from_big_int(swap_params.amount.quotient()),
                    amountInMaximum: U256::from_big_int(amount_in_maximum),
                    sqrtPriceLimitX96: sqrt_price_limit_x96.unwrap_or_default(),
                },
            }
            .abi_encode()
            .into()
        }
    };
    calldatas.push(swap_calldata);

    // flag for whether a refund needs to happen
    let must_refund = token_in.is_native() && trade_type == TradeType::ExactOutput;
    // refund
    if must_refund {
        calldatas.push(encode_refund_eth());
    }

    let value = if token_in.is_native() {
        match trade_type {
            TradeType::ExactOutput => {
                amount_in_maximum.context("Failed to get amount in maximum")?.quotient()
            }
            TradeType::ExactInput => swap_params.amount.quotient(),
        }
    } else {
        BigInt::ZERO
    };

    let calldata_count = calldatas.len();
    let final_calldata = encode_multicall(calldatas);
    debug!(
        calldata_count = calldata_count,
        final_calldata_len = final_calldata.len(),
        value = ?U256::from_big_int(value),
        trade_type = ?trade_type,
        "PancakeSwap V3 swap call parameters finalized"
    );
    Ok(MethodParameters { calldata: final_calldata, value: U256::from_big_int(value) })
}

/// Build swap call parameters for a two-hop swap
#[inline]
#[tracing::instrument(skip(pool1_key, pool2_key, swap_params, swap_options), fields(token_a = ?pool1_key.token_a.address(), token_b = ?pool1_key.token_b.address(), token_c = ?pool2_key.token_b.address(), is_to_b = ?swap_params.is_to_b, trade_type = ?swap_params.trade_type))]
pub fn build_two_hop_swap_call_parameters(
    pool1_key: &V3PoolKey,
    pool2_key: &V3PoolKey,
    swap_params: SwapParams,
    swap_options: SwapOptions,
) -> Result<MethodParameters> {
    // Forbid ExactOutput for two-hop swaps
    if swap_params.trade_type == TradeType::ExactOutput {
        return Err(anyhow::anyhow!(
            "ExactOutput is not supported for two-hop swaps. Use ExactInput instead."
        ));
    }

    let SwapOptions {
        amount_in_maximum: _amount_in_maximum, // Not used for ExactInput
        amount_out_minimum,
        recipient,
        input_token_permit,
        sqrt_price_limit_x96: _sqrt_price_limit_x96, // Not used for multi-hop
        fee: _fee,
    } = swap_options;

    // Determine tokens for the path
    let (token0, token1, token2) =
        determine_two_hop_path_tokens(pool1_key, pool2_key, swap_params.is_to_b)?;

    debug!(
        token0 = ?token0.address(),
        token1 = ?token1.address(),
        token2 = ?token2.address(),
        "Determined two-hop swap path for PancakeSwap V3"
    );

    let (fee0, fee1) = if swap_params.is_to_b {
        (pool1_key.fee, pool2_key.fee)
    } else {
        (pool2_key.fee, pool1_key.fee)
    };

    info!(
        fee0 = ?fee0,
        fee1 = ?fee1,
        "Determined two-hop swap fees for PancakeSwap V3"
    );

    // Only ExactInput is supported for two-hop swaps
    let path =
        encode_two_hop_path_v3(token0.address(), fee0, token1.address(), fee1, token2.address());

    // Convert amounts for ExactInput
    let amount_in = U256::from_big_int(swap_params.amount.quotient());
    let amount_out_minimum = amount_out_minimum
        .as_ref()
        .map(|amt| U256::from_big_int(amt.quotient()))
        .unwrap_or(U256::ZERO);

    // Calculate ETH value for native tokens
    let eth_value = if token0.is_native() { swap_params.amount.quotient() } else { BigInt::ZERO };

    info!(
        token0 = ?token0.address(),
        token1 = ?token1.address(),
        token2 = ?token2.address(),
        recipient = ?recipient,
        amount_in = ?amount_in,
        amount_out_minimum = ?amount_out_minimum,
        "Building PancakeSwap V3 two-hop swap with exactInput"
    );

    // Get deadline (current time + 20 minutes)
    let deadline = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .context("Failed to get current time")?
        .as_secs()
        + 1200; // 20 minutes

    let swap_calldata = ISwapRouter::exactInputCall {
        params: ISwapRouter::ExactInputParams {
            path,
            recipient,
            deadline: U256::from(deadline),
            amountIn: amount_in,
            amountOutMinimum: amount_out_minimum,
        },
    }
    .abi_encode();

    let mut calldatas: Vec<Bytes> = Vec::with_capacity(1 + 3);

    // encode permit if necessary
    if let Some(input_token_permit) = input_token_permit {
        assert!(!token0.is_native(), "NON_TOKEN_PERMIT");
        calldatas.push(encode_permit(&token0, input_token_permit));
    }

    calldatas.push(swap_calldata.into());

    let calldata_count = calldatas.len();
    let final_calldata = encode_multicall(calldatas);
    info!(
        calldata_count = calldata_count,
        final_calldata_len = final_calldata.len(),
        value = ?U256::from_big_int(eth_value),
        "PancakeSwap V3 two-hop swap call parameters finalized"
    );
    Ok(MethodParameters { calldata: final_calldata, value: U256::from_big_int(eth_value) })
}

fn to_u24(fee: FeeAmount) -> U24 {
    match fee {
        FeeAmount::LOWEST => U24::from(100),
        FeeAmount::LOW_200 => U24::from(200),
        FeeAmount::LOW_300 => U24::from(300),
        FeeAmount::LOW_400 => U24::from(400),
        FeeAmount::LOW => U24::from(500),
        FeeAmount::MEDIUM => U24::from(3000),
        FeeAmount::HIGH => U24::from(10000),
        FeeAmount::CUSTOM(val) => U24::from(val),
    }
}