waterpump-evm-pool-sdk 0.1.0

use alloy::{
    primitives::{uint, utils::keccak256, B256, U256},
    sol_types::SolValue,
};
use uniswap_v3_sdk::prelude::TickIndex;

/// Constants for pool state slots
pub const POOLS_SLOT: U256 = uint!(6_U256);
pub const FEE_GROWTH_GLOBAL0_OFFSET: U256 = uint!(1_U256);
pub const LIQUIDITY_OFFSET: U256 = uint!(3_U256);
pub const TICKS_OFFSET: U256 = uint!(4_U256);
pub const TICK_BITMAP_OFFSET: U256 = uint!(5_U256);
pub const POSITIONS_OFFSET: U256 = uint!(6_U256);

/// Get the storage slot for a pool's state structure
///
/// In Uniswap V4 contract, pools is a mapping(bytes32 => State) stored at slot
/// 6 (POOLS_SLOT).
///
/// The State struct storage layout:
/// ```
/// struct State {
///     Slot0 slot0;                                    // state_slot + 0
///     uint256 feeGrowthGlobal0X128;                   // state_slot + 1
///     uint256 feeGrowthGlobal1X128;                   // state_slot + 2
///     uint128 liquidity;                              // state_slot + 3 (LIQUIDITY_OFFSET)
///     mapping(int24 => TickInfo) ticks;               // mapping slot: state_slot + 4 (TICKS_OFFSET)
///     mapping(int16 => uint256) tickBitmap;           // mapping slot: state_slot + 5 (TICK_BITMAP_OFFSET)
///     mapping(bytes32 => Position.State) positions;   // mapping slot: state_slot + 6 (POSITIONS_OFFSET)
/// }
/// ```
///
/// For mapping-type fields, calculate the storage location for a key:
/// slot = keccak256(key, mapping_slot)
///
/// For regular fields, storage location is: state_slot + offset
///
/// ## Arguments
/// * `pool_id` - Unique identifier for the pool (B256)
///
/// ## Returns
/// The starting storage slot position of the State struct (state_slot)
pub fn get_pool_state_slot(pool_id: B256) -> U256 {
    // The starting position of State struct for pools[pool_id] = keccak256(pool_id,
    // POOLS_SLOT) where POOLS_SLOT = 6 is the storage location of the pools
    // mapping
    U256::from_be_bytes(keccak256((pool_id, POOLS_SLOT).abi_encode()).0)
}

/// Get the storage slot for a pool's liquidity value
///
/// liquidity is a uint128 field in the State struct, stored at state_slot + 3.
/// Note: Due to V4's special storage layout, keccak256 is used to calculate the
/// final location.
///
/// ## Arguments
/// * `pool_id` - Unique identifier for the pool
///
/// ## Returns
/// The storage slot location of the liquidity field
pub fn get_liquidity_slot(pool_id: B256) -> U256 {
    let state_slot = get_pool_state_slot(pool_id);
    // liquidity is at state_slot + LIQUIDITY_OFFSET (i.e., +3)
    // Use keccak256 to calculate the final storage location
    U256::from_be_bytes(keccak256((state_slot + LIQUIDITY_OFFSET).abi_encode()).0)
}

/// Get the storage slot for a specific tick in the tick bitmap
///
/// tickBitmap is a mapping(int16 wordPos => uint256) used for fast lookup of
/// words containing active ticks. Each word can represent 256 ticks (each bit
/// represents one tick).
///
/// ## Arguments
/// * `pool_id` - Unique identifier for the pool
/// * `tick` - Tick value, which will be converted to the corresponding word
///   position (int16)
///
/// ## Returns
/// The storage slot location of this word position in the tickBitmap mapping
pub fn get_tick_bitmap_slot<I: TickIndex>(pool_id: B256, tick: I) -> U256 {
    let state_slot = get_pool_state_slot(pool_id);
    // The "slot marker" for tickBitmap mapping is at state_slot +
    // TICK_BITMAP_OFFSET (i.e., +5)
    let tick_bitmap_mapping = state_slot + TICK_BITMAP_OFFSET;
    // Calculate storage location of tickBitmap[wordPos] = keccak256(wordPos,
    // tick_bitmap_mapping) wordPos = tick.to_i24().as_i16() (integer division
    // of tick by 256)
    U256::from_be_bytes(keccak256((tick.to_i24().as_i16(), tick_bitmap_mapping).abi_encode()).0)
}

/// Get the storage slot for a specific tick's information
///
/// ticks is a mapping(int24 => TickInfo) that stores detailed information for
/// each tick, including total liquidity, net liquidity, fee growth outside,
/// etc.
///
/// ## Arguments
/// * `pool_id` - Unique identifier for the pool
/// * `tick` - Tick value (int24)
///
/// ## Returns
/// The storage slot location of the TickInfo data for this tick in the ticks
/// mapping
pub fn get_tick_info_slot<I: TickIndex>(pool_id: B256, tick: I) -> U256 {
    let state_slot = get_pool_state_slot(pool_id);
    // The "slot marker" for ticks mapping is at state_slot + TICKS_OFFSET (i.e.,
    // +4)
    let ticks_mapping_slot = state_slot + TICKS_OFFSET;
    // Calculate storage location of ticks[tick] = keccak256(tick,
    // ticks_mapping_slot)
    U256::from_be_bytes(keccak256((tick.to_i24(), ticks_mapping_slot).abi_encode()).0)
}

/// Get the storage slot for a specific position's information
///
/// positions is a mapping(bytes32 positionKey => Position.State) that stores
/// state information for each liquidity position.
///
/// ## Arguments
/// * `pool_id` - Unique identifier for the pool
/// * `position_key` - Unique identifier for the position (typically
///   keccak256(owner, tickLower, tickUpper))
///
/// ## Returns
/// The storage slot location of the Position.State data for this position in
/// the positions mapping
pub fn get_position_info_slot(pool_id: B256, position_key: B256) -> U256 {
    let state_slot = get_pool_state_slot(pool_id);
    // The "slot marker" for positions mapping is at state_slot + POSITIONS_OFFSET
    // (i.e., +6)
    let position_mapping_slot = state_slot + POSITIONS_OFFSET;
    // Calculate storage location of positions[positionKey] = keccak256(positionKey,
    // position_mapping_slot)
    U256::from_be_bytes(keccak256((position_key, position_mapping_slot).abi_encode()).0)
}