waterpump-evm-pool-sdk 0.1.0

use alloy::{
    primitives::{aliases::U128, Address, U256},
    rpc::types::TransactionReceipt,
};
use alloy_sol_types::sol;
use anyhow::{Context, Result};
use tracing::debug;
use uniswap_sdk_core::prelude::{Currency, CurrencyAmount, ToBig};

use crate::{
    pool_swappers::common::{find_event, find_events},
    types::swap_results::{
        AddLiquidityResult, AddLiquidityResultItem, CollectResult, RemoveLiquidityResult,
        RemoveLiquidityResultItem,
    },
};

// Define event structs for NonfungiblePositionManager events
// Note: PancakeSwap V3 uses the same event structure as Uniswap V3
sol! {
    interface INonfungiblePositionManagerEvents {
        /// IncreaseLiquidity(uint256 indexed tokenId, uint128 liquidity, uint256 amount0, uint256 amount1)
        event IncreaseLiquidity(
            uint256 indexed tokenId,
            uint128 liquidity,
            uint256 amount0,
            uint256 amount1
        );

        /// DecreaseLiquidity(uint256 indexed tokenId, uint128 liquidity, uint256 amount0, uint256 amount1)
        event DecreaseLiquidity(
            uint256 indexed tokenId,
            uint128 liquidity,
            uint256 amount0,
            uint256 amount1
        );

        /// Collect(uint256 indexed tokenId, address recipient, uint256 amount0, uint256 amount1)
        event Collect(
            uint256 indexed tokenId,
            address recipient,
            uint256 amount0,
            uint256 amount1
        );

        /// Transfer(address indexed from, address indexed to, uint256 indexed tokenId)
        event Transfer(
            address indexed from,
            address indexed to,
            uint256 indexed tokenId
        );
    }
}

// Define event structs for PancakeSwap V3 Pool events
// Note: PancakeSwap V3 uses the same event structure as Uniswap V3
sol! {
    interface IPancakeV3PoolEvents {
        /// @notice Emitted when liquidity is minted for a given position
        /// @param sender The address that minted the liquidity
        /// @param owner The owner of the position and recipient of any minted liquidity
        /// @param tickLower The lower tick of the position
        /// @param tickUpper The upper tick of the position
        /// @param amount The amount of liquidity minted to the position range
        /// @param amount0 How much token0 was required for the minted liquidity
        /// @param amount1 How much token1 was required for the minted liquidity
        event Mint(
            address sender,
            address indexed owner,
            int24 indexed tickLower,
            int24 indexed tickUpper,
            uint128 amount,
            uint256 amount0,
            uint256 amount1
        );

        /// @notice Emitted when a position's liquidity is removed
        /// @dev Does not withdraw any fees earned by the liquidity position, which must be withdrawn via #collect
        /// @param owner The owner of the position for which liquidity is removed
        /// @param tickLower The lower tick of the position
        /// @param tickUpper The upper tick of the position
        /// @param amount The amount of liquidity to remove
        /// @param amount0 The amount of token0 withdrawn
        /// @param amount1 The amount of token1 withdrawn
        event Burn(
            address indexed owner,
            int24 indexed tickLower,
            int24 indexed tickUpper,
            uint128 amount,
            uint256 amount0,
            uint256 amount1
        );

        /// @notice Emitted when fees are collected by the owner of a position
        /// @dev Collect events may be emitted with zero amount0 and amount1 when the caller chooses not to collect fees
        /// @param owner The owner of the position for which fees are collected
        /// @param recipient The address that received the collected fees
        /// @param tickLower The lower tick of the position
        /// @param tickUpper The upper tick of the position
        /// @param amount0 The amount of token0 fees collected
        /// @param amount1 The amount of token1 fees collected
        event Collect(
            address indexed owner,
            address recipient,
            int24 indexed tickLower,
            int24 indexed tickUpper,
            uint128 amount0,
            uint128 amount1
        );
    }
}

/// Decode add liquidity result from transaction receipt
///
/// Decodes the IncreaseLiquidity event to extract tokenId, liquidity, amount0,
/// and amount1. For mint operations, also decodes the Mint event from the pool
/// to get amount and tick information, and checks the Transfer event to get the
/// tokenId. Verifies that events are from the expected
/// NonfungiblePositionManager and pool addresses.
pub fn decode_add_liquidity_result(
    receipt: &TransactionReceipt,
    position_manager_address: Address,
    pool_address: Address,
    currency0: Currency,
    currency1: Currency,
) -> Result<AddLiquidityResult> {
    let tx_hash = receipt.transaction_hash;

    // Decode IncreaseLiquidity event and verify it's from the position manager
    let increase_event = find_event(receipt, |log, _| {
        // Verify the event is from the expected position manager address
        if log.address() != position_manager_address {
            return Err(anyhow::anyhow!(
                "IncreaseLiquidity event from unexpected address: expected {:?}, got {:?}",
                position_manager_address,
                log.address()
            ));
        }

        log.log_decode::<INonfungiblePositionManagerEvents::IncreaseLiquidity>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not an IncreaseLiquidity event: {:?}", e))
    })
    .context("Failed to find IncreaseLiquidity event in transaction receipt")?;

    let token_id = increase_event.tokenId;
    let liquidity = U128::from(increase_event.liquidity);
    let amount0 = CurrencyAmount::from_raw_amount(currency0, increase_event.amount0.to_big_int())?;
    let amount1 = CurrencyAmount::from_raw_amount(currency1, increase_event.amount1.to_big_int())?;

    // Extract tick information from Mint event (available for both mint and
    // increase operations) The pool emits a Mint event whenever liquidity is
    // added, including when increasing existing positions
    let mint_event = find_event(receipt, |log, _| {
        // Verify the event is from the expected pool address
        if log.address() != pool_address {
            return Err(anyhow::anyhow!(
                "Mint event from unexpected address: expected {:?}, got {:?}",
                pool_address,
                log.address()
            ));
        }

        log.log_decode::<IPancakeV3PoolEvents::Mint>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not a Mint event: {:?}", e))
    })
    .context("Failed to find Mint event in transaction receipt")?;

    debug!(
        tick_lower = ?mint_event.tickLower,
        tick_upper = ?mint_event.tickUpper,
        amount = ?mint_event.amount,
        amount0 = ?mint_event.amount0,
        amount1 = ?mint_event.amount1,
        "Decoded Mint event from pool"
    );

    let (final_tick_lower, final_tick_upper) = (mint_event.tickLower, mint_event.tickUpper);

    Ok(AddLiquidityResult {
        token_id,
        liquidity,
        amount0,
        amount1,
        tick_upper: final_tick_upper,
        tick_lower: final_tick_lower,
        tx_hash,
    })
}

/// Decode remove liquidity result from transaction receipt
///
/// Decodes the Burn event from the pool to extract tick and amount information,
/// and the DecreaseLiquidity event from the position manager to extract
/// token_id, amount0 and amount1. Verifies that events are from the expected
/// NonfungiblePositionManager and pool addresses.
pub fn decode_remove_liquidity_result(
    receipt: &TransactionReceipt,
    position_manager_address: Address,
    pool_address: Address,
    currency0: Currency,
    currency1: Currency,
) -> Result<RemoveLiquidityResult> {
    let tx_hash = receipt.transaction_hash;

    // Decode Burn event from the pool to get tick and amount information
    let burn_event = find_event(receipt, |log, _| {
        // Verify the event is from the expected pool address
        if log.address() != pool_address {
            return Err(anyhow::anyhow!(
                "Burn event from unexpected address: expected {:?}, got {:?}",
                pool_address,
                log.address()
            ));
        }

        log.log_decode::<IPancakeV3PoolEvents::Burn>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not a Burn event: {:?}", e))
    })
    .context("Failed to find Burn event in transaction receipt")?;

    debug!(
        tick_lower = ?burn_event.tickLower,
        tick_upper = ?burn_event.tickUpper,
        amount = ?burn_event.amount,
        amount0 = ?burn_event.amount0,
        amount1 = ?burn_event.amount1,
        "Decoded Burn event from pool"
    );

    // Decode DecreaseLiquidity event from the position manager to get token_id
    let decrease_event = find_event(receipt, |log, _| {
        // Verify the event is from the expected position manager address
        if log.address() != position_manager_address {
            return Err(anyhow::anyhow!(
                "DecreaseLiquidity event from unexpected address: expected {:?}, got {:?}",
                position_manager_address,
                log.address()
            ));
        }

        log.log_decode::<INonfungiblePositionManagerEvents::DecreaseLiquidity>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not a DecreaseLiquidity event: {:?}", e))
    })
    .context("Failed to find DecreaseLiquidity event in transaction receipt")?;

    // Decode Collect event from the pool to get collected fees and recipient
    let collect_event = find_event(receipt, |log, _| {
        // Verify the event is from the expected pool address
        if log.address() != pool_address {
            return Err(anyhow::anyhow!(
                "Collect event from pool: unexpected address: expected {:?}, got {:?}",
                pool_address,
                log.address()
            ));
        }

        log.log_decode::<IPancakeV3PoolEvents::Collect>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not a Collect event from pool: {:?}", e))
    })
    .context("Failed to find Collect event in transaction receipt")?;

    // Use amounts and liquidity from Burn event (from pool) as they represent the
    // actual amounts withdrawn Use token_id from DecreaseLiquidity event (from
    // position manager) Use collected amounts and recipient from Collect event
    // (from pool)
    Ok(RemoveLiquidityResult {
        token_id: decrease_event.tokenId,
        liquidity: U128::from(burn_event.amount),
        amount0_from_liquidity: CurrencyAmount::from_raw_amount(
            currency0.clone(),
            burn_event.amount0.to_big_int(),
        )?,
        amount1_from_liquidity: CurrencyAmount::from_raw_amount(
            currency1.clone(),
            burn_event.amount1.to_big_int(),
        )?,
        amount0: CurrencyAmount::from_raw_amount(
            currency0,
            U256::from(collect_event.amount0).to_big_int(),
        )?,
        amount1: CurrencyAmount::from_raw_amount(
            currency1,
            U256::from(collect_event.amount1).to_big_int(),
        )?,
        recipient: Some(collect_event.recipient),
        tick_upper: burn_event.tickUpper,
        tick_lower: burn_event.tickLower,
        tx_hash,
    })
}

/// Decode collect result from transaction receipt
///
/// Decodes the Collect event from the pool to extract recipient, tick, and
/// amount information, and the Collect event from the position manager to
/// extract token_id. Verifies that events are from the expected
/// NonfungiblePositionManager and pool addresses.
pub fn decode_collect_result(
    receipt: &TransactionReceipt,
    position_manager_address: Address,
    pool_address: Address,
    currency0: Currency,
    currency1: Currency,
) -> Result<CollectResult> {
    let tx_hash = receipt.transaction_hash;

    // Decode Collect event from the pool to get recipient, tick, and amount
    // information
    let pool_collect_event = find_event(receipt, |log, _| {
        // Verify the event is from the expected pool address
        if log.address() != pool_address {
            return Err(anyhow::anyhow!(
                "Collect event from pool: unexpected address: expected {:?}, got {:?}",
                pool_address,
                log.address()
            ));
        }

        log.log_decode::<IPancakeV3PoolEvents::Collect>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not a Collect event from pool: {:?}", e))
    })
    .context("Failed to find Collect event from pool in transaction receipt")?;

    debug!(
        recipient = ?pool_collect_event.recipient,
        tick_lower = ?pool_collect_event.tickLower,
        tick_upper = ?pool_collect_event.tickUpper,
        amount0 = ?pool_collect_event.amount0,
        amount1 = ?pool_collect_event.amount1,
        "Decoded Collect event from pool"
    );

    // Decode Collect event from the position manager to get token_id
    let position_manager_collect_event = find_event(receipt, |log, _| {
        // Verify the event is from the expected position manager address
        if log.address() != position_manager_address {
            return Err(anyhow::anyhow!(
                "Collect event from position manager: unexpected address: expected {:?}, got {:?}",
                position_manager_address,
                log.address()
            ));
        }

        log.log_decode::<INonfungiblePositionManagerEvents::Collect>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not a Collect event from position manager: {:?}", e))
    })
    .context("Failed to find Collect event from position manager in transaction receipt")?;

    // Use amounts from pool Collect event (as they represent the actual amounts
    // collected) Use token_id from position manager Collect event
    // Use recipient and ticks from pool Collect event
    Ok(CollectResult {
        token_id: position_manager_collect_event.tokenId,
        recipient: pool_collect_event.recipient,
        amount0: CurrencyAmount::from_raw_amount(
            currency0,
            U256::from(pool_collect_event.amount0).to_big_int(),
        )?,
        amount1: CurrencyAmount::from_raw_amount(
            currency1,
            U256::from(pool_collect_event.amount1).to_big_int(),
        )?,
        tick_upper: pool_collect_event.tickUpper,
        tick_lower: pool_collect_event.tickLower,
        reward_amounts: None,
        tx_hash,
    })
}

/// Decode batch add liquidity result from transaction receipt
///
/// Decodes all IncreaseLiquidity events from the position manager and all Mint
/// events from the pool to extract position information for each position
/// created in a batch multicall transaction.
pub fn decode_add_batch_liquidity_result(
    receipt: &TransactionReceipt,
    position_manager_address: Address,
    pool_address: Address,
    currency0: Currency,
    currency1: Currency,
) -> Result<Vec<AddLiquidityResultItem>> {
    // Decode all IncreaseLiquidity events from the position manager
    let increase_events: Vec<_> = find_events(receipt, |log, _| {
        if log.address() != position_manager_address {
            return Err(anyhow::anyhow!(
                "IncreaseLiquidity event from unexpected address: expected {:?}, got {:?}",
                position_manager_address,
                log.address()
            ));
        }

        log.log_decode::<INonfungiblePositionManagerEvents::IncreaseLiquidity>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not an IncreaseLiquidity event: {:?}", e))
    })?;

    // Decode all Mint events from the pool
    let mint_events: Vec<_> = find_events(receipt, |log, _| {
        if log.address() != pool_address {
            return Err(anyhow::anyhow!(
                "Mint event from unexpected address: expected {:?}, got {:?}",
                pool_address,
                log.address()
            ));
        }

        log.log_decode::<IPancakeV3PoolEvents::Mint>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not a Mint event: {:?}", e))
    })?;

    debug!(
        increase_events_count = increase_events.len(),
        mint_events_count = mint_events.len(),
        "Decoded batch add liquidity events"
    );

    // Match IncreaseLiquidity events with Mint events by order
    // They should be emitted in the same order within the multicall
    if increase_events.len() != mint_events.len() {
        return Err(anyhow::anyhow!(
            "Mismatch between IncreaseLiquidity events ({}) and Mint events ({})",
            increase_events.len(),
            mint_events.len()
        ));
    }

    let mut results = Vec::with_capacity(increase_events.len());
    for (increase_event, mint_event) in increase_events.into_iter().zip(mint_events.into_iter()) {
        let token_id = increase_event.tokenId;
        let liquidity = U128::from(increase_event.liquidity);
        let amount0 = CurrencyAmount::from_raw_amount(
            currency0.clone(),
            increase_event.amount0.to_big_int(),
        )?;
        let amount1 = CurrencyAmount::from_raw_amount(
            currency1.clone(),
            increase_event.amount1.to_big_int(),
        )?;

        debug!(
            token_id = ?token_id,
            liquidity = ?liquidity,
            tick_lower = ?mint_event.tickLower,
            tick_upper = ?mint_event.tickUpper,
            "Decoded position in batch"
        );

        results.push(AddLiquidityResultItem {
            token_id,
            liquidity,
            amount0,
            amount1,
            tick_upper: mint_event.tickUpper,
            tick_lower: mint_event.tickLower,
        });
    }

    Ok(results)
}

/// Decode batch remove liquidity result from transaction receipt
///
/// Decodes all DecreaseLiquidity events from the position manager, all Burn
/// events from the pool, and all Collect events from the pool to extract
/// position information for each position removed in a batch multicall
/// transaction.
pub fn decode_remove_batch_liquidity_result(
    receipt: &TransactionReceipt,
    position_manager_address: Address,
    pool_address: Address,
    currency0: Currency,
    currency1: Currency,
) -> Result<Vec<RemoveLiquidityResultItem>> {
    // Decode all DecreaseLiquidity events from the position manager
    let decrease_events: Vec<_> = find_events(receipt, |log, _| {
        if log.address() != position_manager_address {
            return Err(anyhow::anyhow!(
                "DecreaseLiquidity event from unexpected address: expected {:?}, got {:?}",
                position_manager_address,
                log.address()
            ));
        }

        log.log_decode::<INonfungiblePositionManagerEvents::DecreaseLiquidity>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not a DecreaseLiquidity event: {:?}", e))
    })?;

    // Decode all Burn events from the pool
    let burn_events: Vec<_> = find_events(receipt, |log, _| {
        if log.address() != pool_address {
            return Err(anyhow::anyhow!(
                "Burn event from unexpected address: expected {:?}, got {:?}",
                pool_address,
                log.address()
            ));
        }

        log.log_decode::<IPancakeV3PoolEvents::Burn>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not a Burn event: {:?}", e))
    })?;

    // Decode all Collect events from the pool
    let collect_events: Vec<_> = find_events(receipt, |log, _| {
        if log.address() != pool_address {
            return Err(anyhow::anyhow!(
                "Collect event from unexpected address: expected {:?}, got {:?}",
                pool_address,
                log.address()
            ));
        }

        log.log_decode::<IPancakeV3PoolEvents::Collect>()
            .map(|decoded| decoded.inner)
            .map_err(|e| anyhow::anyhow!("Not a Collect event: {:?}", e))
    })?;

    debug!(
        decrease_events_count = decrease_events.len(),
        burn_events_count = burn_events.len(),
        collect_events_count = collect_events.len(),
        "Decoded batch remove liquidity events"
    );

    // Validate event counts match
    if decrease_events.len() != burn_events.len() || decrease_events.len() != collect_events.len() {
        return Err(anyhow::anyhow!(
            "Mismatch between DecreaseLiquidity events ({}), Burn events ({}), and Collect events \
             ({})",
            decrease_events.len(),
            burn_events.len(),
            collect_events.len()
        ));
    }

    let mut results = Vec::with_capacity(decrease_events.len());
    for ((decrease_event, burn_event), collect_event) in
        decrease_events.into_iter().zip(burn_events.into_iter()).zip(collect_events.into_iter())
    {
        let token_id = decrease_event.tokenId;
        let liquidity = U128::from(burn_event.amount);

        let amount0_from_liquidity =
            CurrencyAmount::from_raw_amount(currency0.clone(), burn_event.amount0.to_big_int())?;
        let amount1_from_liquidity =
            CurrencyAmount::from_raw_amount(currency1.clone(), burn_event.amount1.to_big_int())?;
        let amount0 = CurrencyAmount::from_raw_amount(
            currency0.clone(),
            U256::from(collect_event.amount0).to_big_int(),
        )?;
        let amount1 = CurrencyAmount::from_raw_amount(
            currency1.clone(),
            U256::from(collect_event.amount1).to_big_int(),
        )?;

        debug!(
            token_id = ?token_id,
            liquidity = ?liquidity,
            tick_lower = ?burn_event.tickLower,
            tick_upper = ?burn_event.tickUpper,
            "Decoded position removal in batch"
        );

        results.push(RemoveLiquidityResultItem {
            token_id,
            liquidity,
            amount0,
            amount1,
            amount0_from_liquidity,
            amount1_from_liquidity,
            recipient: Some(collect_event.recipient),
            tick_upper: burn_event.tickUpper,
            tick_lower: burn_event.tickLower,
        });
    }

    Ok(results)
}