waterpump-evm-pool-sdk 0.1.0

use alloy::{
    network::Ethereum,
    primitives::Address,
    providers::{DynProvider, Provider},
};
use anyhow::Result;
use tracing::{debug, instrument, trace};
use uniswap_lens::bindings::{
    ierc20::IERC20::IERC20Instance, iuniswapv3pool::IUniswapV3Pool::IUniswapV3PoolInstance,
};
use uniswap_sdk_core::{
    entities::BaseCurrency,
    prelude::{Currency, CurrencyAmount, Price, ToBig},
};
use uniswap_v3_sdk::prelude::FeeAmount;

use super::utils::sqrt_price_x96_to_price;
use crate::{
    impl_pool_base, impl_v3_pool_state, impl_v3_pool_viewer, types::v3_pool_key::V3PoolKey,
};

type PoolPriceAndBalances =
    (Price<Currency, Currency>, CurrencyAmount<Currency>, CurrencyAmount<Currency>);

#[derive(Clone, Debug)]
pub struct V3WithHolderPoolViewer {
    pub pool_address: Address,
    pub pool_key: V3PoolKey,
    pub is_token_a_base: bool,
    pub holder_address: Address,
    pub provider: DynProvider<Ethereum>,
}

impl V3WithHolderPoolViewer {
    #[instrument(skip(token_a, token_b), fields(token_a_addr = ?token_a.address(), token_b_addr = ?token_b.address(), fee = ?fee, pool_address = ?pool_address, holder_address = ?holder_address))]
    pub fn new(
        token_a: Currency,
        token_b: Currency,
        fee: FeeAmount,
        is_token_a_base: bool,
        pool_address: Address,
        holder_address: Address,
        provider: DynProvider<Ethereum>,
    ) -> Result<Self> {
        debug!(pool_address = ?pool_address, "Using provided pool address");
        Ok(Self {
            pool_address,
            pool_key: V3PoolKey { token_a, token_b, fee },
            is_token_a_base,
            holder_address,
            provider,
        })
    }

    #[instrument(skip(self), fields(holder_address = ?self.holder_address, token_a = ?self.pool_key.token_a.address(), token_b = ?self.pool_key.token_b.address()))]
    pub async fn get_holder_balances(
        &self,
    ) -> Result<(CurrencyAmount<Currency>, CurrencyAmount<Currency>)> {
        let token_a = IERC20Instance::new(self.pool_key.token_a.address(), self.provider.clone());
        let token_b = IERC20Instance::new(self.pool_key.token_b.address(), self.provider.clone());
        let mut multicall = self.provider.multicall().dynamic();
        multicall = multicall.add_dynamic(token_a.balanceOf(self.holder_address));
        multicall = multicall.add_dynamic(token_b.balanceOf(self.holder_address));
        trace!("Executing multicall for balance queries");
        let results = multicall.aggregate().await?;
        let balance0 = results[0].to_big_int();
        let balance1 = results[1].to_big_int();
        debug!(balance0 = ?balance0, balance1 = ?balance1, "Raw balances retrieved");
        let amount_a = CurrencyAmount::from_raw_amount(self.pool_key.token_a.clone(), balance0)?;
        let amount_b = CurrencyAmount::from_raw_amount(self.pool_key.token_b.clone(), balance1)?;
        debug!(token_a_balance = %amount_a.to_exact(), token_b_balance = %amount_b.to_exact(), "Balances retrieved successfully");
        Ok((amount_a, amount_b))
    }

    #[instrument(skip(self), fields(pool_address = ?self.pool_address, holder_address = ?self.holder_address, is_token_a_base = is_token_a_base))]
    pub async fn get_pool_price_and_balances(
        &self,
        is_token_a_base: bool,
    ) -> Result<PoolPriceAndBalances> {
        // Use non-dynamic multicall (.add()) to mix different call types - returns a
        // tuple
        let pool_contract = IUniswapV3PoolInstance::new(self.pool_address, self.provider.clone());
        let token_a = IERC20Instance::new(self.pool_key.token_a.address(), self.provider.clone());
        let token_b = IERC20Instance::new(self.pool_key.token_b.address(), self.provider.clone());
        let multicall = self
            .provider
            .multicall()
            .add(pool_contract.slot0())
            .add(token_a.balanceOf(self.holder_address))
            .add(token_b.balanceOf(self.holder_address));

        trace!("Executing multicall for slot0 and balance queries");
        let (slot0_result, balance0_result, balance1_result) = multicall.aggregate().await?;
        let balance0 = balance0_result.to_big_int();
        let balance1 = balance1_result.to_big_int();
        debug!(sqrt_price_x96 = ?slot0_result.sqrtPriceX96, tick = ?slot0_result.tick, "Slot0 result retrieved");
        debug!(balance0 = ?balance0, balance1 = ?balance1, "Raw balances retrieved");

        let price = sqrt_price_x96_to_price(
            self.pool_key.token_a.clone(),
            self.pool_key.token_b.clone(),
            slot0_result.sqrtPriceX96,
        );
        let price = if is_token_a_base {
            debug!("Using token_a as base currency");
            price
        } else {
            debug!("Using token_b as base currency (inverting price)");
            price.invert()
        };

        let amount_a = CurrencyAmount::from_raw_amount(self.pool_key.token_a.clone(), balance0)?;
        let amount_b = CurrencyAmount::from_raw_amount(self.pool_key.token_b.clone(), balance1)?;
        debug!(
            price = %price.to_significant(8, None).unwrap_or_else(|_| "N/A".to_string()),
            token_a_balance = %amount_a.to_exact(),
            token_b_balance = %amount_b.to_exact(),
            "Price and balances retrieved successfully"
        );
        Ok((price, amount_a, amount_b))
    }

    /// Get the pool address
    pub fn pool_address(&self) -> Address { self.pool_address }

    /// Get the pool key
    pub fn pool_key(&self) -> &V3PoolKey { &self.pool_key }
}

impl_pool_base!(V3WithHolderPoolViewer);
impl_v3_pool_viewer!(V3WithHolderPoolViewer);
impl_v3_pool_state!(V3WithHolderPoolViewer);