Merge pull request #111 from propeller-heads/zz/uniswap_v4/implement-…

…swap

feat(uniswap_v4): Implement `swap()` for `UniswapV4State`

src/evm/protocol/uniswap_v3/mod.rs

@@ -1,8 +1,4 @@
 //! Uniswap V3 Decentralized Exchange
 pub mod enums;
-mod liquidity_math;
-mod solidity_math;
-mod sqrt_price_math;
 pub mod state;
-mod swap_math;
 pub mod tycho_decoder;

src/evm/protocol/uniswap_v3/state.rs

@@ -5,20 +5,21 @@ use num_bigint::BigUint;
 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,
-    tycho_decoder::i24_be_bytes_to_i32,
-};
+use super::{enums::FeeAmount, 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::{
+            liquidity_math,
+            sqrt_price_math::sqrt_price_q96_to_f64,
+            swap_math,
             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,
             },
+            StepComputation, SwapResults, SwapState,
         },
     },
     models::Token,
@@ -38,35 +39,6 @@ pub struct UniswapV3State {
     ticks: TickList,
 }
 
-#[derive(Debug)]
-struct SwapState {
-    amount_remaining: I256,
-    amount_calculated: I256,
-    sqrt_price: U256,
-    tick: i32,
-    liquidity: u128,
-}
-
-#[derive(Debug)]
-struct StepComputation {
-    sqrt_price_start: U256,
-    tick_next: i32,
-    initialized: bool,
-    sqrt_price_next: U256,
-    amount_in: U256,
-    amount_out: U256,
-    fee_amount: U256,
-}
-
-#[derive(Debug)]
-struct SwapResults {
-    amount_calculated: I256,
-    sqrt_price: U256,
-    liquidity: u128,
-    tick: i32,
-    gas_used: U256,
-}
-
 impl UniswapV3State {
     /// Creates a new instance of `UniswapV3State`.
     ///

src/evm/protocol/uniswap_v4/state.rs

@@ -1,26 +1,51 @@
-use alloy_primitives::U256;
+#![allow(dead_code)] //TODO: remove when used
 
-use crate::evm::protocol::utils::uniswap::tick_list::{TickInfo, TickList};
+use alloy_primitives::{Sign, I256, U256};
+
+use crate::{
+    evm::protocol::{
+        safe_math::{safe_add_u256, safe_sub_u256},
+        u256_num::u256_to_biguint,
+        utils::uniswap::{
+            liquidity_math, swap_math,
+            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,
+            },
+            StepComputation, SwapResults, SwapState,
+        },
+    },
+    protocol::{errors::SimulationError, models::GetAmountOutResult, state::ProtocolSim},
+};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct UniswapV4State {
     liquidity: u128,
     sqrt_price: U256,
-    lp_fee: i32,
-    protocol_fees: UniswapV4ProtocolFees,
+    fees: UniswapV4Fees,
     tick: i32,
     ticks: TickList,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub struct UniswapV4ProtocolFees {
-    zero2one: i32,
-    one2zero: i32,
+pub struct UniswapV4Fees {
+    // Protocol fees in the zero for one direction
+    zero_for_one: u32,
+    // Protocol fees in the one for zero direction
+    one_for_zero: u32,
+    // Liquidity providers fees
+    lp_fee: u32,
 }
 
-impl UniswapV4ProtocolFees {
-    pub fn new(zero2one: i32, one2zero: i32) -> Self {
-        Self { zero2one, one2zero }
+impl UniswapV4Fees {
+    pub fn new(zero_for_one: u32, one_for_zero: u32, lp_fee: u32) -> Self {
+        Self { zero_for_one, one_for_zero, lp_fee }
+    }
+
+    fn calculate_swap_fees_pips(&self, zero_for_one: bool) -> u32 {
+        let protocol_fees = if zero_for_one { self.zero_for_one } else { self.one_for_zero };
+        protocol_fees + self.lp_fee
     }
 }
 
@@ -29,8 +54,7 @@ impl UniswapV4State {
     pub fn new(
         liquidity: u128,
         sqrt_price: U256,
-        lp_fee: i32,
-        protocol_fees: UniswapV4ProtocolFees,
+        fees: UniswapV4Fees,
         tick: i32,
         tick_spacing: i32,
         ticks: Vec<TickInfo>,
@@ -43,6 +67,200 @@ impl UniswapV4State {
                 .expect("tick_spacing should always be positive"),
             ticks,
         );
-        UniswapV4State { liquidity, sqrt_price, lp_fee, protocol_fees, tick, ticks: tick_list }
+        UniswapV4State { liquidity, sqrt_price, fees, tick, ticks: tick_list }
+    }
+
+    fn swap(
+        &self,
+        zero_for_one: bool,
+        amount_specified: I256,
+        sqrt_price_limit: Option<U256>,
+    ) -> Result<SwapResults, SimulationError> {
+        if self.liquidity == 0 {
+            return Err(SimulationError::RecoverableError("No liquidity".to_string()));
+        }
+        let price_limit = if let Some(limit) = sqrt_price_limit {
+            limit
+        } else if zero_for_one {
+            safe_add_u256(MIN_SQRT_RATIO, U256::from(1u64))?
+        } else {
+            safe_sub_u256(MAX_SQRT_RATIO, U256::from(1u64))?
+        };
+
+        if zero_for_one {
+            assert!(price_limit > MIN_SQRT_RATIO);
+            assert!(price_limit < self.sqrt_price);
+        } else {
+            assert!(price_limit < MAX_SQRT_RATIO);
+            assert!(price_limit > self.sqrt_price);
+        }
+
+        let exact_input = amount_specified > I256::from_raw(U256::from(0u64));
+
+        let mut state = SwapState {
+            amount_remaining: amount_specified,
+            amount_calculated: I256::from_raw(U256::from(0u64)),
+            sqrt_price: self.sqrt_price,
+            tick: self.tick,
+            liquidity: self.liquidity,
+        };
+        let mut gas_used = U256::from(130_000);
+
+        while state.amount_remaining != I256::from_raw(U256::from(0u64)) &&
+            state.sqrt_price != price_limit
+        {
+            let (mut next_tick, initialized) = match self
+                .ticks
+                .next_initialized_tick_within_one_word(state.tick, zero_for_one)
+            {
+                Ok((tick, init)) => (tick, init),
+                Err(tick_err) => match tick_err.kind {
+                    TickListErrorKind::TicksExeeded => {
+                        let mut new_state = self.clone();
+                        new_state.liquidity = state.liquidity;
+                        new_state.tick = state.tick;
+                        new_state.sqrt_price = state.sqrt_price;
+                        return Err(SimulationError::InvalidInput(
+                            "Ticks exceeded".into(),
+                            Some(GetAmountOutResult::new(
+                                u256_to_biguint(state.amount_calculated.abs().into_raw()),
+                                u256_to_biguint(gas_used),
+                                Box::new(new_state),
+                            )),
+                        ));
+                    }
+                    _ => return Err(SimulationError::FatalError("Unknown error".to_string())),
+                },
+            };
+
+            next_tick = next_tick.clamp(MIN_TICK, MAX_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,
+                UniswapV4State::get_sqrt_ratio_target(sqrt_price_next, price_limit, zero_for_one),
+                state.liquidity,
+                state.amount_remaining,
+                self.fees
+                    .calculate_swap_fees_pips(zero_for_one),
+            )?;
+            state.sqrt_price = sqrt_price;
+
+            let step = StepComputation {
+                sqrt_price_start: state.sqrt_price,
+                tick_next: next_tick,
+                initialized,
+                sqrt_price_next,
+                amount_in,
+                amount_out,
+                fee_amount,
+            };
+            if exact_input {
+                state.amount_remaining -= I256::checked_from_sign_and_abs(
+                    Sign::Positive,
+                    safe_add_u256(step.amount_in, step.fee_amount)?,
+                )
+                .unwrap();
+                state.amount_calculated -=
+                    I256::checked_from_sign_and_abs(Sign::Positive, step.amount_out).unwrap();
+            } else {
+                state.amount_remaining +=
+                    I256::checked_from_sign_and_abs(Sign::Positive, step.amount_out).unwrap();
+                state.amount_calculated += I256::checked_from_sign_and_abs(
+                    Sign::Positive,
+                    safe_add_u256(step.amount_in, step.fee_amount)?,
+                )
+                .unwrap();
+            }
+            if state.sqrt_price == step.sqrt_price_next {
+                if step.initialized {
+                    let liquidity_raw = self
+                        .ticks
+                        .get_tick(step.tick_next)
+                        .unwrap()
+                        .net_liquidity;
+                    let liquidity_net = if zero_for_one { -liquidity_raw } else { liquidity_raw };
+                    state.liquidity =
+                        liquidity_math::add_liquidity_delta(state.liquidity, liquidity_net);
+                }
+                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 = get_tick_at_sqrt_ratio(state.sqrt_price)?;
+            }
+            gas_used = safe_add_u256(gas_used, U256::from(2000))?;
+        }
+        Ok(SwapResults {
+            amount_calculated: state.amount_calculated,
+            sqrt_price: state.sqrt_price,
+            liquidity: state.liquidity,
+            tick: state.tick,
+            gas_used,
+        })
+    }
+
+    fn get_sqrt_ratio_target(
+        sqrt_price_next: U256,
+        sqrt_price_limit: U256,
+        zero_for_one: bool,
+    ) -> U256 {
+        let cond1 = if zero_for_one {
+            sqrt_price_next < sqrt_price_limit
+        } else {
+            sqrt_price_next > sqrt_price_limit
+        };
+
+        if cond1 {
+            sqrt_price_limit
+        } else {
+            sqrt_price_next
+        }
+    }
+}
+
+#[allow(unused_variables)] //TODO: remove when implemented
+impl ProtocolSim for UniswapV4State {
+    fn fee(&self) -> f64 {
+        todo!()
+    }
+
+    fn spot_price(
+        &self,
+        base: &crate::models::Token,
+        quote: &crate::models::Token,
+    ) -> Result<f64, SimulationError> {
+        todo!()
+    }
+
+    fn get_amount_out(
+        &self,
+        amount_in: num_bigint::BigUint,
+        token_in: &crate::models::Token,
+        token_out: &crate::models::Token,
+    ) -> Result<GetAmountOutResult, SimulationError> {
+        todo!()
+    }
+
+    fn delta_transition(
+        &mut self,
+        delta: tycho_core::dto::ProtocolStateDelta,
+        tokens: &std::collections::HashMap<tycho_core::Bytes, crate::models::Token>,
+    ) -> Result<(), crate::protocol::errors::TransitionError<String>> {
+        todo!()
+    }
+
+    fn clone_box(&self) -> Box<dyn ProtocolSim> {
+        todo!()
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        todo!()
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        todo!()
+    }
+
+    fn eq(&self, other: &dyn ProtocolSim) -> bool {
+        todo!()
     }
 }

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

@@ -1,5 +1,5 @@
 // Solidity spec: function addDelta(uint128 x, int128 y) internal pure returns (uint128 z) {
-pub(super) fn add_liquidity_delta(x: u128, y: i128) -> u128 {
+pub(crate) fn add_liquidity_delta(x: u128, y: i128) -> u128 {
     if y < 0 {
         x - ((-y) as u128)
     } else {

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

@@ -1,2 +1,37 @@
+use alloy_primitives::{I256, U256};
+
+pub(crate) mod liquidity_math;
+mod solidity_math;
+pub(crate) mod sqrt_price_math;
+pub(crate) mod swap_math;
 pub(crate) mod tick_list;
 pub(crate) mod tick_math;
+
+#[derive(Debug)]
+pub(crate) struct SwapState {
+    pub(crate) amount_remaining: I256,
+    pub(crate) amount_calculated: I256,
+    pub(crate) sqrt_price: U256,
+    pub(crate) tick: i32,
+    pub(crate) liquidity: u128,
+}
+
+#[derive(Debug)]
+pub(crate) struct StepComputation {
+    pub(crate) sqrt_price_start: U256,
+    pub(crate) tick_next: i32,
+    pub(crate) initialized: bool,
+    pub(crate) sqrt_price_next: U256,
+    pub(crate) amount_in: U256,
+    pub(crate) amount_out: U256,
+    pub(crate) fee_amount: U256,
+}
+
+#[derive(Debug)]
+pub(crate) struct SwapResults {
+    pub(crate) amount_calculated: I256,
+    pub(crate) sqrt_price: U256,
+    pub(crate) liquidity: u128,
+    pub(crate) tick: i32,
+    pub(crate) gas_used: U256,
+}

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

@@ -160,7 +160,7 @@ fn get_next_sqrt_price_from_amount1_rounding_down(
 ///
 /// # Panics
 /// Will panic if the `x` is bigger than U160.
-pub fn sqrt_price_q96_to_f64(x: U256, token_0_decimals: u32, token_1_decimals: u32) -> f64 {
+pub(crate) fn sqrt_price_q96_to_f64(x: U256, token_0_decimals: u32, token_1_decimals: u32) -> f64 {
     assert!(x < U160_MAX);
     let token_correction = 10f64.powi(token_0_decimals as i32 - token_1_decimals as i32);