Merge pull request #27 from propeller-heads/tnl/ENG-3759-get-amount-out

feat: (VMPoolState) implement get_amount_out

src/evm/simulation.rs

@@ -29,7 +29,7 @@ use super::{
 };
 
 /// An error representing any transaction simulation result other than successful execution
-#[derive(Debug, Display, Clone)]
+#[derive(Debug, Display, Clone, PartialEq)]
 pub enum SimulationEngineError {
     /// Something went wrong while getting storage; might be caused by network issues.
     /// Retrying may help.
@@ -298,6 +298,7 @@ fn interpret_evm_success(
         gas_used: gas_used - gas_refunded,
     }
 }
+
 #[derive(Debug)]
 /// Data needed to invoke a transaction simulation
 pub struct SimulationParameters {
@@ -635,9 +636,9 @@ mod tests {
         let caller = Address::from_str("0x0000000000000000000000000000000000000000")?;
         let router_addr = Address::from_str("0x7a250d5630B4cF539739dF2C5dAcb4c659F2488D")?;
         let router_abi = BaseContract::from(
-        parse_abi(&[
-            "function getAmountsOut(uint amountIn, address[] memory path) public view returns (uint[] memory amounts)",
-        ])?
+            parse_abi(&[
+                "function getAmountsOut(uint amountIn, address[] memory path) public view returns (uint[] memory amounts)",
+            ])?
         );
         let weth_addr = Address::from_str("0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2")?;
         let usdc_addr = Address::from_str("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48")?;

src/protocol/errors.rs

@@ -57,6 +57,7 @@ impl From<SimulationError> for InvalidSnapshotError {
 /// - `InsufficientAmount`: Error indicating that the amount provided for the trade is too low.
 /// - `ArithmeticOverflow`: Error indicating that an arithmetic operation got an U256 to overflow
 /// - `Unknown`: Error indicating that an unknown error occurred during the simulation.
+/// - `SellAmountTooHigh`: Indicates an error when the sell amount is higher than the sell limit.
 #[derive(Error, Debug)]
 pub enum SimulationError {
     #[error("ABI loading error: {0}")]
@@ -83,4 +84,6 @@ pub enum SimulationError {
     ArithmeticOverflow(),
     #[error("Unknown error")]
     Unknown(),
+    #[error("Sell amount is higher than sell limit")]
+    SellAmountTooHigh(), // TODO: Make it recoverable
 }

src/protocol/models.rs

@@ -75,12 +75,13 @@ impl PartialEq for Pair {
 pub struct GetAmountOutResult {
     pub amount: U256,
     pub gas: U256,
+    pub new_state: Box<dyn ProtocolSim>,
 }
 
 impl GetAmountOutResult {
     /// Constructs a new GetAmountOutResult struct with the given amount and gas
-    pub fn new(amount: U256, gas: U256) -> Self {
-        GetAmountOutResult { amount, gas }
+    pub fn new(amount: U256, gas: U256, new_state: Box<dyn ProtocolSim>) -> Self {
+        GetAmountOutResult { amount, gas, new_state }
     }
 
     /// Aggregates the given GetAmountOutResult struct to the current one.

src/protocol/uniswap_v2/state.rs

@@ -9,14 +9,14 @@ use crate::{
         state::{ProtocolEvent, ProtocolSim},
         BytesConvertible,
     },
-    safe_math::{safe_add_u256, safe_div_u256, safe_mul_u256},
+    safe_math::{safe_add_u256, safe_div_u256, safe_mul_u256, safe_sub_u256},
 };
 use ethers::types::U256;
 use tycho_core::dto::ProtocolStateDelta;
 
 use super::{events::UniswapV2Sync, reserve_price::spot_price_from_reserves};
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct UniswapV2State {
     pub reserve0: U256,
     pub reserve1: U256,
@@ -110,8 +110,15 @@ impl ProtocolSim for UniswapV2State {
             safe_add_u256(safe_mul_u256(reserve_sell, U256::from(1000))?, amount_in_with_fee)?;
 
         let amount_out = safe_div_u256(numerator, denominator)?;
-
-        Ok(GetAmountOutResult::new(amount_out, U256::from(120_000)))
+        let mut new_state = self.clone();
+        if zero2one {
+            new_state.reserve0 = safe_add_u256(self.reserve0, amount_in)?;
+            new_state.reserve1 = safe_sub_u256(self.reserve1, amount_out)?;
+        } else {
+            new_state.reserve0 = safe_sub_u256(self.reserve0, amount_out)?;
+            new_state.reserve1 = safe_add_u256(self.reserve1, amount_in)?;
+        };
+        Ok(GetAmountOutResult::new(amount_out, U256::from(120_000), Box::new(new_state)))
     }
 
     fn delta_transition(
@@ -154,7 +161,7 @@ impl ProtocolSim for UniswapV2State {
     }
 
     fn clone_box(&self) -> Box<dyn ProtocolSim> {
-        Box::new(*self)
+        Box::new(self.clone())
     }
 
     fn as_any(&self) -> &dyn Any {
@@ -211,20 +218,20 @@ mod tests {
     fn test_get_amount_out(
         #[case] r0: U256,
         #[case] r1: U256,
-        #[case] t0d: usize,
-        #[case] t1d: usize,
+        #[case] token_0_decimals: usize,
+        #[case] token_1_decimals: usize,
         #[case] amount_in: U256,
         #[case] exp: U256,
     ) {
         let t0 = ERC20Token::new(
             "0x0000000000000000000000000000000000000000",
-            t0d,
+            token_0_decimals,
             "T0",
             U256::from(10_000),
         );
         let t1 = ERC20Token::new(
             "0x0000000000000000000000000000000000000001",
-            t1d,
+            token_1_decimals,
             "T0",
             U256::from(10_000),
         );
@@ -235,6 +242,16 @@ mod tests {
             .unwrap();
 
         assert_eq!(res.amount, exp);
+        let new_state = res
+            .new_state
+            .as_any()
+            .downcast_ref::<UniswapV2State>()
+            .unwrap();
+        assert_eq!(new_state.reserve0, r0 + amount_in);
+        assert_eq!(new_state.reserve1, r1 - exp);
+        // Assert that the old state is unchanged
+        assert_eq!(state.reserve0, r0);
+        assert_eq!(state.reserve1, r1);
     }
 
     #[test]

src/protocol/uniswap_v3/state.rs

@@ -152,6 +152,7 @@ impl UniswapV3State {
                         return Err(SimulationError::InsufficientData(GetAmountOutResult::new(
                             state.amount_calculated.abs().into_raw(),
                             gas_used,
+                            self.clone_box(),
                         )))
                     }
                     _ => return Err(SimulationError::Unknown()),
@@ -267,13 +268,18 @@ impl ProtocolSim for UniswapV3State {
         let result = self.swap(zero_for_one, amount_specified, None)?;
 
         trace!(?amount_in, ?token_a, ?token_b, ?zero_for_one, ?result, "V3 SWAP");
+        let mut new_state = self.clone();
+        new_state.liquidity = result.liquidity;
+        new_state.tick = result.tick;
+        new_state.sqrt_price = result.sqrt_price;
 
         Ok(GetAmountOutResult::new(
             result
                 .amount_calculated
                 .abs()
                 .into_raw(),
             result.gas_used,
+            Box::new(new_state),
         ))
     }
 

src/protocol/vm/adapter_contract.rs

@@ -1,6 +1,14 @@
 // TODO: remove skip for clippy dead_code check
 #![allow(dead_code)]
 
+use std::collections::{HashMap, HashSet};
+
+use ethers::{
+    abi::{Address, Token},
+    types::U256,
+};
+use revm::{primitives::Address as rAddress, DatabaseRef};
+
 use crate::{
     evm::account_storage::StateUpdate,
     protocol::{
@@ -11,18 +19,12 @@ use crate::{
         },
     },
 };
-use ethers::{
-    abi::{Address, Token},
-    types::U256,
-};
-use revm::{primitives::Address as rAddress, DatabaseRef};
-use std::collections::{HashMap, HashSet};
 
 #[derive(Debug)]
 pub struct Trade {
-    received_amount: U256,
-    gas_used: U256,
-    price: f64,
+    pub received_amount: U256,
+    pub gas_used: U256,
+    pub price: f64,
 }
 
 /// An implementation of `TychoSimulationContract` specific to the `AdapterContract` ABI interface,
@@ -42,7 +44,7 @@ impl<D: DatabaseRef + std::clone::Clone> TychoSimulationContract<D>
 where
     D::Error: std::fmt::Debug,
 {
-    pub async fn price(
+    pub fn price(
         &self,
         pair_id: String,
         sell_token: Address,
@@ -64,15 +66,14 @@ where
         ];
 
         let res = self
-            .call("price", args, block, None, overwrites, None, U256::zero())
-            .await?
+            .call("price", args, block, None, overwrites, None, U256::zero())?
             .return_value;
         let price = self.calculate_price(res[0].clone())?;
         Ok(price)
     }
 
     #[allow(clippy::too_many_arguments)]
-    pub async fn swap(
+    pub fn swap(
         &self,
         pair_id: String,
         sell_token: Address,
@@ -90,25 +91,37 @@ where
             Token::Uint(amount),
         ];
 
-        let res = self
-            .call("swap", args, block, None, overwrites, None, U256::zero())
-            .await?;
-        let received_amount = res.return_value[0]
-            .clone()
-            .into_uint()
-            .unwrap();
-        let gas_used = res.return_value[1]
-            .clone()
-            .into_uint()
-            .unwrap();
-        let price = self
-            .calculate_price(res.return_value[2].clone())
-            .unwrap()[0];
+        let res = self.call("swap", args, block, None, overwrites, None, U256::zero())?;
+
+        let (received_amount, gas_used, price) = if let Token::Tuple(ref return_value) =
+            res.return_value[0]
+        {
+            match &return_value[..] {
+                [Token::Uint(amount), Token::Uint(gas), Token::Tuple(price_elements)] => {
+                    let received_amount = *amount;
+                    let gas_used = *gas;
+
+                    let price_token = Token::Array(vec![Token::Tuple(price_elements.clone())]);
+                    let price = self
+                        .calculate_price(price_token)?
+                        .first()
+                        .cloned()
+                        .ok_or_else(|| {
+                            SimulationError::DecodingError("There wasn't a calculated price".into())
+                        })?;
+
+                    Ok((received_amount, gas_used, price))
+                }
+                _ => Err(SimulationError::DecodingError("Incorrect types found for price".into())),
+            }
+        } else {
+            Err(SimulationError::DecodingError("return_value is not a Token::Tuple".into()))
+        }?;
 
         Ok((Trade { received_amount, gas_used, price }, res.simulation_result.state_updates))
     }
 
-    pub async fn get_limits(
+    pub fn get_limits(
         &self,
         pair_id: String,
         sell_token: Address,
@@ -123,8 +136,7 @@ where
         ];
 
         let res = self
-            .call("getLimits", args, block, None, overwrites, None, U256::zero())
-            .await?
+            .call("getLimits", args, block, None, overwrites, None, U256::zero())?
             .return_value;
 
         if let Some(Token::Array(inner)) = res.first() {
@@ -138,7 +150,7 @@ where
         Err(SimulationError::DecodingError("Unexpected response format".into()))
     }
 
-    pub async fn get_capabilities(
+    pub fn get_capabilities(
         &self,
         pair_id: String,
         sell_token: Address,
@@ -151,8 +163,7 @@ where
         ];
 
         let res = self
-            .call("getCapabilities", args, 1, None, None, None, U256::zero())
-            .await?
+            .call("getCapabilities", args, 1, None, None, None, U256::zero())?
             .return_value;
         let capabilities: HashSet<Capability> = match res.first() {
             Some(Token::Array(inner_tokens)) => inner_tokens
@@ -168,10 +179,9 @@ where
         Ok(capabilities)
     }
 
-    pub async fn min_gas_usage(&self) -> Result<u64, SimulationError> {
+    pub fn min_gas_usage(&self) -> Result<u64, SimulationError> {
         let res = self
-            .call("minGasUsage", vec![], 1, None, None, None, U256::zero())
-            .await?
+            .call("minGasUsage", vec![], 1, None, None, None, U256::zero())?
             .return_value;
         Ok(res[0]
             .clone()
@@ -213,7 +223,7 @@ where
                 })
                 .collect()
         } else {
-            Err(SimulationError::DecodingError("Expected Token::Array".to_string()))
+            Err(SimulationError::DecodingError("Price is not a Token::Array".to_string()))
         }
     }
 }