Merge pull request #110 from propeller-heads/zz/uniswap_v4/add-state-…

…struct

feat(uniswap_v4): Implement `UniswapV4State` struct

src/evm/protocol/mod.rs

@@ -3,5 +3,6 @@ pub mod safe_math;
 pub mod u256_num;
 pub mod uniswap_v2;
 pub mod uniswap_v3;
+pub mod uniswap_v4;
 pub(crate) mod utils;
 pub mod vm;

src/evm/protocol/uniswap_v3/mod.rs

@@ -5,6 +5,4 @@ mod solidity_math;
 mod sqrt_price_math;
 pub mod state;
 mod swap_math;
-pub mod tick_list;
-mod tick_math;
 pub mod tycho_decoder;

src/evm/protocol/uniswap_v3/state.rs

@@ -6,18 +6,20 @@ use tracing::trace;
 use tycho_core::{dto::ProtocolStateDelta, Bytes};
 
 use super::{
-    enums::FeeAmount,
-    liquidity_math,
-    sqrt_price_math::sqrt_price_q96_to_f64,
-    swap_math,
-    tick_list::{TickInfo, TickList},
-    tick_math,
+    enums::FeeAmount, liquidity_math, sqrt_price_math::sqrt_price_q96_to_f64, swap_math,
     tycho_decoder::i24_be_bytes_to_i32,
 };
 use crate::{
     evm::protocol::{
         safe_math::{safe_add_u256, safe_sub_u256},
         u256_num::u256_to_biguint,
+        utils::uniswap::{
+            tick_list::{TickInfo, TickList, TickListErrorKind},
+            tick_math::{
+                get_sqrt_ratio_at_tick, get_tick_at_sqrt_ratio, MAX_SQRT_RATIO, MAX_TICK,
+                MIN_SQRT_RATIO, MIN_TICK,
+            },
+        },
     },
     models::Token,
     protocol::{
@@ -107,16 +109,16 @@ impl UniswapV3State {
         let price_limit = if let Some(limit) = sqrt_price_limit {
             limit
         } else if zero_for_one {
-            safe_add_u256(tick_math::MIN_SQRT_RATIO, U256::from(1u64))?
+            safe_add_u256(MIN_SQRT_RATIO, U256::from(1u64))?
         } else {
-            safe_sub_u256(tick_math::MAX_SQRT_RATIO, U256::from(1u64))?
+            safe_sub_u256(MAX_SQRT_RATIO, U256::from(1u64))?
         };
 
         if zero_for_one {
-            assert!(price_limit > tick_math::MIN_SQRT_RATIO);
+            assert!(price_limit > MIN_SQRT_RATIO);
             assert!(price_limit < self.sqrt_price);
         } else {
-            assert!(price_limit < tick_math::MAX_SQRT_RATIO);
+            assert!(price_limit < MAX_SQRT_RATIO);
             assert!(price_limit > self.sqrt_price);
         }
 
@@ -140,7 +142,7 @@ impl UniswapV3State {
             {
                 Ok((tick, init)) => (tick, init),
                 Err(tick_err) => match tick_err.kind {
-                    super::tick_list::TickListErrorKind::TicksExeeded => {
+                    TickListErrorKind::TicksExeeded => {
                         let mut new_state = self.clone();
                         new_state.liquidity = state.liquidity;
                         new_state.tick = state.tick;
@@ -158,9 +160,9 @@ impl UniswapV3State {
                 },
             };
 
-            next_tick = next_tick.clamp(tick_math::MIN_TICK, tick_math::MAX_TICK);
+            next_tick = next_tick.clamp(MIN_TICK, MAX_TICK);
 
-            let sqrt_price_next = tick_math::get_sqrt_ratio_at_tick(next_tick)?;
+            let sqrt_price_next = get_sqrt_ratio_at_tick(next_tick)?;
             let (sqrt_price, amount_in, amount_out, fee_amount) = swap_math::compute_swap_step(
                 state.sqrt_price,
                 UniswapV3State::get_sqrt_ratio_target(sqrt_price_next, price_limit, zero_for_one),
@@ -209,7 +211,7 @@ impl UniswapV3State {
                 }
                 state.tick = if zero_for_one { step.tick_next - 1 } else { step.tick_next };
             } else if state.sqrt_price != step.sqrt_price_start {
-                state.tick = tick_math::get_tick_at_sqrt_ratio(state.sqrt_price)?;
+                state.tick = get_tick_at_sqrt_ratio(state.sqrt_price)?;
             }
             gas_used = safe_add_u256(gas_used, U256::from(2000))?;
         }

src/evm/protocol/uniswap_v3/tycho_decoder.rs

@@ -4,8 +4,9 @@ use alloy_primitives::U256;
 use tycho_client::feed::{synchronizer::ComponentWithState, Header};
 use tycho_core::Bytes;
 
-use super::{enums::FeeAmount, state::UniswapV3State, tick_list::TickInfo};
+use super::{enums::FeeAmount, state::UniswapV3State};
 use crate::{
+    evm::protocol::utils::uniswap::tick_list::TickInfo,
     models::Token,
     protocol::{errors::InvalidSnapshotError, models::TryFromWithBlock},
 };

src/evm/protocol/uniswap_v4/mod.rs

@@ -0,0 +1 @@
+pub mod state;

src/evm/protocol/uniswap_v4/state.rs

@@ -0,0 +1,48 @@
+use alloy_primitives::U256;
+
+use crate::evm::protocol::utils::uniswap::tick_list::{TickInfo, TickList};
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct UniswapV4State {
+    liquidity: u128,
+    sqrt_price: U256,
+    lp_fee: i32,
+    protocol_fees: UniswapV4ProtocolFees,
+    tick: i32,
+    ticks: TickList,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct UniswapV4ProtocolFees {
+    zero2one: i32,
+    one2zero: i32,
+}
+
+impl UniswapV4ProtocolFees {
+    pub fn new(zero2one: i32, one2zero: i32) -> Self {
+        Self { zero2one, one2zero }
+    }
+}
+
+impl UniswapV4State {
+    /// Creates a new `UniswapV4State` with specified values.
+    pub fn new(
+        liquidity: u128,
+        sqrt_price: U256,
+        lp_fee: i32,
+        protocol_fees: UniswapV4ProtocolFees,
+        tick: i32,
+        tick_spacing: i32,
+        ticks: Vec<TickInfo>,
+    ) -> Self {
+        let tick_list = TickList::from(
+            tick_spacing
+                .try_into()
+                // even though it's given as int24, tick_spacing must be positive, see here:
+                // https://github.com/Uniswap/v4-core/blob/a22414e4d7c0d0b0765827fe0a6c20dfd7f96291/src/libraries/TickMath.sol#L25-L28
+                .expect("tick_spacing should always be positive"),
+            ticks,
+        );
+        UniswapV4State { liquidity, sqrt_price, lp_fee, protocol_fees, tick, ticks: tick_list }
+    }
+}

src/evm/protocol/utils.rs → src/evm/protocol/utils/mod.rs

@@ -1,3 +1,5 @@
+pub(crate) mod uniswap;
+
 use alloy_primitives::Address;
 use tycho_core::Bytes;
 

src/evm/protocol/utils/uniswap/mod.rs

@@ -0,0 +1,2 @@
+pub(crate) mod tick_list;
+pub(crate) mod tick_math;

src/evm/protocol/uniswap_v3/tick_list.rs → src/evm/protocol/utils/uniswap/tick_list.rs

@@ -6,9 +6,9 @@ use super::tick_math;
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub struct TickInfo {
-    pub(super) index: i32,
-    pub(super) net_liquidity: i128,
-    pub(super) sqrt_price: U256,
+    pub(crate) index: i32,
+    pub(crate) net_liquidity: i128,
+    pub(crate) sqrt_price: U256,
 }
 
 impl TickInfo {
@@ -27,26 +27,26 @@ impl PartialOrd for TickInfo {
 }
 
 #[derive(Debug)]
-pub(super) struct TickListError {
-    pub(super) kind: TickListErrorKind,
+pub(crate) struct TickListError {
+    pub(crate) kind: TickListErrorKind,
 }
 
 #[derive(Debug, PartialEq)]
-pub(super) enum TickListErrorKind {
+pub(crate) enum TickListErrorKind {
     NotFound,
     BelowSmallest,
     AtOrAboveLargest,
     TicksExeeded,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub(super) struct TickList {
+pub(crate) struct TickList {
     tick_spacing: u16,
     ticks: Vec<TickInfo>,
 }
 
 impl TickList {
-    pub(super) fn from(spacing: u16, ticks: Vec<TickInfo>) -> Self {
+    pub(crate) fn from(spacing: u16, ticks: Vec<TickInfo>) -> Self {
         let tick_list = TickList { tick_spacing: spacing, ticks };
         let valid = tick_list.valid_ticks();
         if valid.is_ok() {
@@ -113,7 +113,7 @@ impl TickList {
         }
     }
 
-    pub(super) fn set_tick_liquidity(&mut self, tick: i32, liquidity: i128) {
+    pub(crate) fn set_tick_liquidity(&mut self, tick: i32, liquidity: i128) {
         match self
             .ticks
             .binary_search_by(|t| t.index.cmp(&tick))
@@ -152,7 +152,7 @@ impl TickList {
         tick >= maximum
     }
 
-    pub(super) fn get_tick(&self, index: i32) -> Result<&TickInfo, TickListError> {
+    pub(crate) fn get_tick(&self, index: i32) -> Result<&TickInfo, TickListError> {
         match self
             .ticks
             .binary_search_by(|el| el.index.cmp(&index))
@@ -196,7 +196,7 @@ impl TickList {
         }
     }
 
-    pub(super) fn next_initialized_tick_within_one_word(
+    pub(crate) fn next_initialized_tick_within_one_word(
         &self,
         tick: i32,
         lte: bool,

src/evm/protocol/uniswap_v3/tick_math.rs → src/evm/protocol/utils/uniswap/tick_math.rs

@@ -7,17 +7,17 @@ use crate::{
     protocol::errors::SimulationError,
 };
 
-pub(super) const MIN_TICK: i32 = -887272;
-pub(super) const MAX_TICK: i32 = 887272;
+pub(crate) const MIN_TICK: i32 = -887272;
+pub(crate) const MAX_TICK: i32 = 887272;
 
 // MIN_SQRT_RATIO: 4295128739
-pub(super) const MIN_SQRT_RATIO: U256 = U256::from_limbs([4295128739u64, 0, 0, 0]);
+pub(crate) const MIN_SQRT_RATIO: U256 = U256::from_limbs([4295128739u64, 0, 0, 0]);
 
 // MAX_SQRT_RATIO: 1461446703485210103287273052203988822378723970342
-pub(super) const MAX_SQRT_RATIO: U256 =
+pub(crate) const MAX_SQRT_RATIO: U256 =
     U256::from_limbs([6743328256752651558u64, 17280870778742802505u64, 4294805859u64, 0]);
 
-pub(super) fn get_sqrt_ratio_at_tick(tick: i32) -> Result<U256, SimulationError> {
+pub(crate) fn get_sqrt_ratio_at_tick(tick: i32) -> Result<U256, SimulationError> {
     assert!(tick.abs() <= MAX_TICK);
     let abs_tick = U256::from(tick.unsigned_abs());
     let mut ratio = if abs_tick.bit(0) {
@@ -154,7 +154,7 @@ fn most_significant_bit(x: U256) -> usize {
     x.bit_len() - 1
 }
 
-pub(super) fn get_tick_at_sqrt_ratio(sqrt_price: U256) -> Result<i32, SimulationError> {
+pub(crate) fn get_tick_at_sqrt_ratio(sqrt_price: U256) -> Result<i32, SimulationError> {
     assert!(sqrt_price >= MIN_SQRT_RATIO && sqrt_price < MAX_SQRT_RATIO);
     let ratio_x128 = sqrt_price << 32;
     let msb = most_significant_bit(ratio_x128);