fix: Return new state in GetAmountOutResult

Bonus fix:
- Combine `get_amount_out and` `get_amount_out_no_limit_check` in preparation for moving `get_amount_out` to the ProtocolSim trait. This is also nicer since we are no longer checking for the limit amount twice.

src/protocol/dodo/state.rs

@@ -173,7 +173,11 @@ where
             .pool_abi
             .decode_output::<U256, _>("querySellBaseToken", simulation_result.result)
             .expect("DODO: Failed decoding get_amount_out result!");
-        Ok(GetAmountOutResult { amount: amount_out, gas: U256::from(simulation_result.gas_used) })
+        Ok(GetAmountOutResult {
+            amount: amount_out,
+            gas: U256::from(simulation_result.gas_used),
+            new_state: None,
+        })
     }
 
     #[allow(unused_variables)]

src/protocol/models.rs

@@ -75,12 +75,13 @@ impl PartialEq for Pair {
 pub struct GetAmountOutResult {
     pub amount: U256,
     pub gas: U256,
+    pub new_state: Option<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: Option<Box<dyn ProtocolSim>>) -> Self {
+        GetAmountOutResult { amount, gas, new_state }
     }
 
     /// Aggregates the given GetAmountOutResult struct to the current one.

src/protocol/uniswap_v2/state.rs

@@ -118,7 +118,7 @@ impl ProtocolSim for UniswapV2State {
 
         let amount_out = safe_div_u256(numerator, denominator)?;
 
-        Ok(GetAmountOutResult::new(amount_out, U256::from(120_000)))
+        Ok(GetAmountOutResult::new(amount_out, U256::from(120_000), None))
     }
 
     fn delta_transition(

src/protocol/uniswap_v3/state.rs

@@ -154,6 +154,7 @@ impl UniswapV3State {
                             Some(GetAmountOutResult::new(
                                 state.amount_calculated.abs().into_raw(),
                                 gas_used,
+                                None,
                             )),
                         ))
                     }
@@ -281,6 +282,7 @@ impl ProtocolSim for UniswapV3State {
                 .abs()
                 .into_raw(),
             result.gas_used,
+            None,
         ))
     }
 

src/protocol/vm/state.rs

@@ -558,45 +558,23 @@ impl VMPoolState<PreCachedDB> {
         sell_amount: U256,
         buy_token: ERC20Token,
     ) -> Result<(U256, U256, VMPoolState<PreCachedDB>), TychoSimulationError> {
-        if self
-            .capabilities
-            .contains(&Capability::HardLimits)
-        {
-            let sell_amount_limit = self
-                .clone()
-                .get_sell_amount_limit(vec![sell_token.clone(), buy_token.clone()])
-                .await?;
-            if sell_amount_limit < sell_amount {
-                let (partial_buy_amount, partial_gas_used, new_state) = self
-                    .get_amount_out_no_limit_check(
-                        sell_token.clone(),
-                        sell_amount_limit,
-                        buy_token.clone(),
-                    )
-                    .await?;
-                return Err(TychoSimulationError::SellAmountTooHigh(
-                    partial_buy_amount,
-                    partial_gas_used,
-                    new_state,
-                    sell_amount_limit,
-                ));
-            }
-        }
-        self.get_amount_out_no_limit_check(sell_token, sell_amount, buy_token)
-            .await
-    }
+        let mut sell_amount_respecting_limit = sell_amount;
+        let mut sell_amount_exceeds_limit = false;
 
-    async fn get_amount_out_no_limit_check(
-        &self,
-        sell_token: ERC20Token,
-        sell_amount: U256,
-        buy_token: ERC20Token,
-    ) -> Result<(U256, U256, VMPoolState<PreCachedDB>), TychoSimulationError> {
         let sell_amount_limit = self
             .clone()
             .get_sell_amount_limit(vec![sell_token.clone(), buy_token.clone()])
             .await?;
 
+        if self
+            .capabilities
+            .contains(&Capability::HardLimits) &&
+            sell_amount_limit < sell_amount
+        {
+            sell_amount_respecting_limit = sell_amount_limit;
+            sell_amount_exceeds_limit = true;
+        }
+
         let overwrites = self
             .clone()
             .get_overwrites(vec![sell_token.clone(), buy_token.clone()], sell_amount_limit)
@@ -617,7 +595,7 @@ impl VMPoolState<PreCachedDB> {
                 sell_token.clone().address,
                 buy_token.clone().address,
                 false,
-                sell_amount,
+                sell_amount_respecting_limit,
                 self.block.number,
                 Some(overwrites),
             )
@@ -664,6 +642,16 @@ impl VMPoolState<PreCachedDB> {
         }
 
         let buy_amount = trade.received_amount;
+
+        if sell_amount_exceeds_limit {
+            return Err(TychoSimulationError::SellAmountTooHigh(
+                // Partial buy amount and gas used
+                buy_amount,
+                trade.gas_used,
+                new_state,
+                sell_amount_limit,
+            ));
+        }
         Ok((buy_amount, trade.gas_used, new_state))
     }
 }