Merge pull request #16 from shuhuiluo/feature/pool

Add `address` method to pool entity and extend pool utilities

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "uniswap-v3-sdk"
-version = "0.11.0"
+version = "0.12.0"
 edition = "2021"
 authors = ["Shuhui Luo <twitter.com/aureliano_law>"]
 description = "Uniswap V3 SDK for Rust"
@@ -27,14 +27,15 @@ once_cell = "1.19"
 regex = { version = "1.10", optional = true }
 ruint = "1.11"
 thiserror = "1.0"
-uniswap-sdk-core = "0.10.0"
+uniswap-sdk-core = "0.11.0"
 uniswap_v3_math = "0.4.1"
 
 [features]
 extensions = ["aperture-lens", "ethers", "regex"]
 
 [dev-dependencies]
 criterion = "0.5.1"
+ethers = "2.0"
 tokio = { version = "1.35", features = ["full"] }
 
 [[bench]]

src/entities/pool.rs

@@ -1,5 +1,5 @@
 use crate::prelude::*;
-use alloy_primitives::{Address, B256, I256, U256};
+use alloy_primitives::{Address, ChainId, B256, I256, U256};
 use anyhow::Result;
 use num_bigint::BigUint;
 use once_cell::sync::Lazy;
@@ -95,7 +95,22 @@ impl Pool {
         })
     }
 
-    pub fn chain_id(&self) -> u32 {
+    /// Returns the pool address
+    pub fn address(
+        &self,
+        init_code_hash_manual_override: Option<B256>,
+        factory_address_override: Option<Address>,
+    ) -> Address {
+        Self::get_address(
+            &self.token0,
+            &self.token1,
+            self.fee,
+            init_code_hash_manual_override,
+            factory_address_override,
+        )
+    }
+
+    pub fn chain_id(&self) -> ChainId {
         self.token0.chain_id()
     }
 

src/extensions/mod.rs

@@ -1,9 +1,11 @@
 //! Extensions to the core library.
 
 mod ephemeral_tick_data_provider;
+mod pool;
 mod price_tick_conversions;
 
 pub use ephemeral_tick_data_provider::EphemeralTickDataProvider;
+pub use pool::*;
 pub use price_tick_conversions::*;
 
 use crate::prelude::*;

src/extensions/pool.rs

@@ -0,0 +1,226 @@
+use crate::prelude::*;
+use alloy_primitives::{Address, ChainId, B256};
+use anyhow::Result;
+use aperture_lens::prelude::{
+    get_populated_ticks_in_range,
+    i_uniswap_v3_pool::{IUniswapV3Pool, Slot0Return},
+    ierc20_metadata::IERC20Metadata,
+};
+use ethers::prelude::*;
+use num_integer::Integer;
+use std::sync::Arc;
+use uniswap_sdk_core::{prelude::Token, token};
+use uniswap_v3_math::utils::u256_to_ruint;
+
+pub fn get_pool_contract<M: Middleware>(
+    factory: Address,
+    token_a: Address,
+    token_b: Address,
+    fee: FeeAmount,
+    client: Arc<M>,
+) -> IUniswapV3Pool<M> {
+    IUniswapV3Pool::new(
+        compute_pool_address(factory, token_a, token_b, fee, None).into_array(),
+        client,
+    )
+}
+
+pub async fn get_pool<M: Middleware>(
+    chain_id: ChainId,
+    factory: Address,
+    token_a: Address,
+    token_b: Address,
+    fee: FeeAmount,
+    client: Arc<M>,
+    block_id: Option<BlockId>,
+) -> Result<Pool, MulticallError<M>> {
+    let pool_contract = get_pool_contract(factory, token_a, token_b, fee, client.clone());
+    let token_a_contract = IERC20Metadata::new(token_a.into_array(), client.clone());
+    let token_b_contract = IERC20Metadata::new(token_b.into_array(), client.clone());
+    let mut multicall = Multicall::new_with_chain_id(client, None, Some(chain_id)).unwrap();
+    multicall.block = block_id;
+    multicall
+        .add_call(pool_contract.slot_0(), false)
+        .add_call(pool_contract.liquidity(), false)
+        .add_call(token_a_contract.decimals(), false)
+        .add_call(token_a_contract.name(), false)
+        .add_call(token_a_contract.symbol(), false)
+        .add_call(token_b_contract.decimals(), false)
+        .add_call(token_b_contract.name(), false)
+        .add_call(token_b_contract.symbol(), false);
+    let (
+        slot_0,
+        liquidity,
+        token_a_decimals,
+        token_a_name,
+        token_a_symbol,
+        token_b_decimals,
+        token_b_name,
+        token_b_symbol,
+    ): (Slot0Return, u128, u8, String, String, u8, String, String) = multicall.call().await?;
+    let sqrt_price_x96 = slot_0.sqrt_price_x96;
+    if sqrt_price_x96.is_zero() {
+        panic!("Pool has been created but not yet initialized");
+    }
+    Ok(Pool::new(
+        token!(
+            chain_id,
+            token_a,
+            token_a_decimals,
+            token_a_symbol,
+            token_a_name
+        ),
+        token!(
+            chain_id,
+            token_b,
+            token_b_decimals,
+            token_b_symbol,
+            token_b_name
+        ),
+        fee,
+        u256_to_ruint(sqrt_price_x96),
+        liquidity,
+        None,
+    )
+    .unwrap())
+}
+
+/// Normalizes the specified tick range.
+fn normalize_ticks(
+    tick_current: i32,
+    tick_spacing: i32,
+    tick_lower: i32,
+    tick_upper: i32,
+) -> (i32, i32, i32) {
+    assert!(tick_lower <= tick_upper, "tickLower > tickUpper");
+    // The current tick must be within the specified tick range.
+    let tick_current_aligned = tick_current.div_mod_floor(&tick_spacing).0 * tick_spacing;
+    let tick_lower = tick_lower.max(MIN_TICK).min(tick_current_aligned);
+    let tick_upper = tick_upper.min(MAX_TICK).max(tick_current_aligned);
+    (tick_current_aligned, tick_lower, tick_upper)
+}
+
+fn reconstruct_liquidity_array(
+    tick_array: Vec<(i32, i128)>,
+    tick_current_aligned: i32,
+    current_liquidity: u128,
+) -> Result<Vec<(i32, u128)>> {
+    // Locate the tick in the populated ticks array with the current liquidity.
+    let current_index = tick_array
+        .iter()
+        .position(|&(tick, _)| tick > tick_current_aligned)
+        .unwrap()
+        - 1;
+    // Accumulate the liquidity from the current tick to the end of the populated ticks array.
+    let mut cumulative_liquidity = current_liquidity;
+    let mut liquidity_array = vec![(0, 0); tick_array.len()];
+    for (i, &(tick, liquidity_net)) in tick_array.iter().enumerate().skip(current_index + 1) {
+        // added when tick is crossed from left to right
+        cumulative_liquidity = add_delta(cumulative_liquidity, liquidity_net)?;
+        liquidity_array[i] = (tick, cumulative_liquidity);
+    }
+    cumulative_liquidity = current_liquidity;
+    for (i, &(tick, liquidity_net)) in tick_array.iter().enumerate().take(current_index + 1).rev() {
+        liquidity_array[i] = (tick, cumulative_liquidity);
+        // subtracted when tick is crossed from right to left
+        cumulative_liquidity = add_delta(cumulative_liquidity, -liquidity_net)?;
+    }
+    Ok(liquidity_array)
+}
+
+pub async fn get_liquidity_array_for_pool<M: Middleware>(
+    pool: Pool,
+    tick_lower: i32,
+    tick_upper: i32,
+    client: Arc<M>,
+    block_id: Option<BlockId>,
+    init_code_hash_manual_override: Option<B256>,
+    factory_address_override: Option<Address>,
+) -> Result<Vec<(i32, u128)>, ContractError<M>> {
+    let (tick_current_aligned, tick_lower, tick_upper) = normalize_ticks(
+        pool.tick_current,
+        pool.tick_spacing(),
+        tick_lower,
+        tick_upper,
+    );
+    let ticks = get_populated_ticks_in_range(
+        pool.address(init_code_hash_manual_override, factory_address_override)
+            .into_array()
+            .into(),
+        tick_lower,
+        tick_upper,
+        client,
+        block_id,
+    )
+    .await?;
+    Ok(reconstruct_liquidity_array(
+        ticks
+            .into_iter()
+            .map(|tick| (tick.tick, tick.liquidity_net))
+            .collect(),
+        tick_current_aligned,
+        pool.liquidity,
+    )
+    .unwrap())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use alloy_primitives::address;
+
+    async fn pool() -> Pool {
+        get_pool(
+            1,
+            FACTORY_ADDRESS,
+            address!("2260FAC5E5542a773Aa44fBCfeDf7C193bc2C599"),
+            address!("C02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"),
+            FeeAmount::LOW,
+            Arc::new(MAINNET.provider()),
+            Some(BlockId::from(17000000)),
+        )
+        .await
+        .unwrap()
+    }
+
+    #[tokio::test]
+    async fn test_get_pool() {
+        let pool = pool().await;
+        assert_eq!(pool.token0.symbol.unwrap(), "WBTC");
+        assert_eq!(pool.token1.symbol.unwrap(), "WETH");
+        assert_eq!(pool.tick_current, 257344);
+        assert_eq!(pool.liquidity, 786352807736110014);
+    }
+
+    #[tokio::test]
+    async fn test_get_liquidity_array_for_pool() {
+        let pool = pool().await;
+        const DOUBLE_TICK: i32 = 6932;
+        let tick_current_aligned =
+            pool.tick_current.div_mod_floor(&pool.tick_spacing()).0 * pool.tick_spacing();
+        let liquidity = pool.liquidity;
+        let tick_lower = pool.tick_current - DOUBLE_TICK;
+        let tick_upper = pool.tick_current + DOUBLE_TICK;
+        let liquidity_array = get_liquidity_array_for_pool(
+            pool,
+            tick_lower,
+            tick_upper,
+            Arc::new(MAINNET.provider()),
+            Some(BlockId::from(17000000)),
+            None,
+            None,
+        )
+        .await
+        .unwrap();
+        assert!(!liquidity_array.is_empty());
+        assert_eq!(
+            liquidity_array[liquidity_array
+                .iter()
+                .position(|&(tick, _)| tick > tick_current_aligned)
+                .unwrap()
+                - 1]
+            .1,
+            liquidity
+        );
+    }
+}