feat: increase coverage of spot prices function

examples/price_printer/main.rs

@@ -14,7 +14,9 @@ use tycho_simulation::{
     evm::{
         engine_db::tycho_db::PreCachedDB,
         protocol::{
-            filters::{balancer_pool_filter, curve_pool_filter, uniswap_v4_pool_with_hook_filter},
+            filters::{
+                balancer_v2_pool_filter, curve_pool_filter, uniswap_v4_pool_with_hook_filter,
+            },
             uniswap_v2::state::UniswapV2State,
             uniswap_v3::state::UniswapV3State,
             uniswap_v4::state::UniswapV4State,
@@ -69,7 +71,7 @@ async fn main() {
             .exchange::<EVMPoolState<PreCachedDB>>(
                 "vm:balancer_v2",
                 tvl_filter.clone(),
-                Some(balancer_pool_filter),
+                Some(balancer_v2_pool_filter),
             )
             .exchange::<EVMPoolState<PreCachedDB>>(
                 "vm:curve",

examples/quickstart/main.rs

@@ -7,7 +7,7 @@ use tycho_simulation::{
     evm::{
         engine_db::tycho_db::PreCachedDB,
         protocol::{
-            filters::{balancer_pool_filter, uniswap_v4_pool_with_hook_filter},
+            filters::{balancer_v2_pool_filter, uniswap_v4_pool_with_hook_filter},
             uniswap_v2::state::UniswapV2State,
             uniswap_v4::state::UniswapV4State,
             vm::state::EVMPoolState,
@@ -52,7 +52,7 @@ async fn main() {
         .exchange::<EVMPoolState<PreCachedDB>>(
             "vm:balancer_v2",
             tvl_filter.clone(),
-            Some(balancer_pool_filter),
+            Some(balancer_v2_pool_filter),
         )
         .exchange::<UniswapV4State>(
             "uniswap_v4",

src/evm/protocol/filters.rs

@@ -9,7 +9,7 @@ use crate::evm::protocol::vm::utils::json_deserialize_be_bigint_list;
 const ZERO_ADDRESS: &str = "0x0000000000000000000000000000000000000000";
 const ZERO_ADDRESS_ARR: [u8; 20] = [0u8; 20];
 
-pub fn balancer_pool_filter(component: &ComponentWithState) -> bool {
+pub fn balancer_v2_pool_filter(component: &ComponentWithState) -> bool {
     // Check for rate_providers in static_attributes
     info!("Checking Balancer pool {}", component.component.id);
     if let Some(rate_providers_data) = component
@@ -140,3 +140,21 @@ pub fn uniswap_v4_pool_with_hook_filter(component: &ComponentWithState) -> bool
     }
     true
 }
+
+/// Filters out pool that have rely on ERC4626 in Balancer V3
+pub fn balancer_v3_pool_filter(component: &ComponentWithState) -> bool {
+    if let Some(erc4626) = component
+        .component
+        .static_attributes
+        .get("erc4626")
+    {
+        if erc4626.to_vec() == [1u8] {
+            info!(
+                "Filtering out Balancer V3 pool {} because it uses ERC4626",
+                component.component.id
+            );
+            return false;
+        }
+    }
+    true
+}

src/evm/protocol/vm/constants.rs

@@ -13,10 +13,12 @@ lazy_static! {
 
 pub const ERC20_BYTECODE: &[u8] = include_bytes!("assets/ERC20.bin");
 pub const BALANCER_V2: &[u8] = include_bytes!("assets/BalancerV2SwapAdapter.evm.runtime");
+pub const BALANCER_V3: &[u8] = include_bytes!("assets/BalancerV3SwapAdapter.evm.runtime");
 pub const CURVE: &[u8] = include_bytes!("assets/CurveSwapAdapter.evm.runtime");
 pub fn get_adapter_file(protocol: &str) -> Result<&'static [u8], SimulationError> {
     match protocol {
         "balancer_v2" => Ok(BALANCER_V2),
+        "balancer_v3" => Ok(BALANCER_V3),
         "curve" => Ok(CURVE),
         _ => {
             Err(SimulationError::FatalError(format!("Adapter for protocol {} not found", protocol)))

src/evm/protocol/vm/state.rs

@@ -24,7 +24,10 @@ use crate::{
             engine_db_interface::EngineDatabaseInterface, simulation_db::BlockHeader,
             tycho_db::PreCachedDB,
         },
-        protocol::{u256_num::u256_to_biguint, utils::bytes_to_address},
+        protocol::{
+            u256_num::{u256_to_biguint, u256_to_f64},
+            utils::bytes_to_address,
+        },
         ContractCompiler, SlotId,
     },
     models::{Balances, Token},
@@ -180,51 +183,96 @@ where
         &mut self,
         tokens: &HashMap<Bytes, Token>,
     ) -> Result<(), SimulationError> {
-        self.ensure_capability(Capability::PriceFunction)?;
-        for [sell_token_address, buy_token_address] in self
-            .tokens
-            .iter()
-            .permutations(2)
-            .map(|p| [p[0], p[1]])
-        {
-            let sell_token_address = bytes_to_address(sell_token_address)?;
-            let buy_token_address = bytes_to_address(buy_token_address)?;
-            let overwrites = Some(self.get_overwrites(
-                vec![sell_token_address, buy_token_address],
-                *MAX_BALANCE / U256::from(100),
-            )?);
-            let sell_amount_limit = self.get_sell_amount_limit(
-                vec![sell_token_address, buy_token_address],
-                overwrites.clone(),
-            )?;
-            let price_result = self.adapter_contract.price(
-                &self.id,
-                sell_token_address,
-                buy_token_address,
-                vec![sell_amount_limit / U256::from(100)],
-                self.block.number,
-                overwrites,
-            )?;
-
-            let price = if self
-                .capabilities
-                .contains(&Capability::ScaledPrice)
-            {
-                *price_result.first().ok_or_else(|| {
-                    SimulationError::FatalError("Calculated price array is empty".to_string())
-                })?
-            } else {
-                let unscaled_price = price_result.first().ok_or_else(|| {
-                    SimulationError::FatalError("Calculated price array is empty".to_string())
-                })?;
-                let sell_token_decimals = self.get_decimals(tokens, &sell_token_address)?;
-                let buy_token_decimals = self.get_decimals(tokens, &buy_token_address)?;
-                *unscaled_price * 10f64.powi(sell_token_decimals as i32) /
-                    10f64.powi(buy_token_decimals as i32)
-            };
-
-            self.spot_prices
-                .insert((sell_token_address, buy_token_address), price);
+        match self.ensure_capability(Capability::PriceFunction) {
+            Ok(_) => {
+                for [sell_token_address, buy_token_address] in self
+                    .tokens
+                    .iter()
+                    .permutations(2)
+                    .map(|p| [p[0], p[1]])
+                {
+                    let sell_token_address = bytes_to_address(sell_token_address)?;
+                    let buy_token_address = bytes_to_address(buy_token_address)?;
+                    let overwrites = Some(self.get_overwrites(
+                        vec![sell_token_address, buy_token_address],
+                        *MAX_BALANCE / U256::from(100),
+                    )?);
+                    let sell_amount_limit = self.get_sell_amount_limit(
+                        vec![sell_token_address, buy_token_address],
+                        overwrites.clone(),
+                    )?;
+                    let price_result = self.adapter_contract.price(
+                        &self.id,
+                        sell_token_address,
+                        buy_token_address,
+                        vec![sell_amount_limit / U256::from(100)],
+                        self.block.number,
+                        overwrites,
+                    )?;
+
+                    let price = if self
+                        .capabilities
+                        .contains(&Capability::ScaledPrice)
+                    {
+                        *price_result.first().ok_or_else(|| {
+                            SimulationError::FatalError(
+                                "Calculated price array is empty".to_string(),
+                            )
+                        })?
+                    } else {
+                        let unscaled_price = price_result.first().ok_or_else(|| {
+                            SimulationError::FatalError(
+                                "Calculated price array is empty".to_string(),
+                            )
+                        })?;
+                        let sell_token_decimals = self.get_decimals(tokens, &sell_token_address)?;
+                        let buy_token_decimals = self.get_decimals(tokens, &buy_token_address)?;
+                        *unscaled_price * 10f64.powi(sell_token_decimals as i32) /
+                            10f64.powi(buy_token_decimals as i32)
+                    };
+
+                    self.spot_prices
+                        .insert((sell_token_address, buy_token_address), price);
+                }
+            }
+            Err(_) => {
+                for iter_tokens in self.tokens.iter().permutations(2) {
+                    let t0 = bytes_to_address(iter_tokens[0])?;
+                    let t1 = bytes_to_address(iter_tokens[1])?;
+
+                    let overwrites =
+                        Some(self.get_overwrites(vec![t0, t1], *MAX_BALANCE / U256::from(100))?);
+
+                    let x1 = self.get_sell_amount_limit(vec![t0, t1], overwrites.clone())? /
+                        U256::from(100);
+                    let x2 = x1 + (x1 / U256::from(100));
+
+                    let y1 = self
+                        .adapter_contract
+                        .swap(&self.id, t0, t1, false, x1, self.block.number, overwrites.clone())?
+                        .0
+                        .received_amount;
+
+                    let y2 = self
+                        .adapter_contract
+                        .swap(&self.id, t0, t1, false, x2, self.block.number, overwrites)?
+                        .0
+                        .received_amount;
+
+                    let sell_token_decimals = self.get_decimals(tokens, &t0)?;
+                    let buy_token_decimals = self.get_decimals(tokens, &t1)?;
+
+                    let num = y2 - y1;
+                    let den = x2 - x1;
+
+                    let token_correction =
+                        10f64.powi(sell_token_decimals as i32 - buy_token_decimals as i32);
+                    let marginal_price = u256_to_f64(num) / u256_to_f64(den) * token_correction;
+
+                    self.spot_prices
+                        .insert((t0, t1), marginal_price);
+                }
+            }
         }
         Ok(())
     }
@@ -995,6 +1043,42 @@ mod tests {
         assert_eq!(bal_dai_spot_price, &7.071_503_245_428_246);
     }
 
+    #[tokio::test]
+    async fn test_set_spot_prices_without_capability() {
+        // Tests set Spot Prices functions when the pool doesn't have PriceFunction capability
+        let mut pool_state = setup_pool_state().await;
+
+        pool_state
+            .capabilities
+            .remove(&Capability::PriceFunction);
+
+        pool_state
+            .set_spot_prices(
+                &vec![bal(), dai()]
+                    .into_iter()
+                    .map(|t| (t.address.clone(), t))
+                    .collect(),
+            )
+            .unwrap();
+
+        let dai_bal_spot_price = pool_state
+            .spot_prices
+            .get(&(
+                bytes_to_address(&pool_state.tokens[0]).unwrap(),
+                bytes_to_address(&pool_state.tokens[1]).unwrap(),
+            ))
+            .unwrap();
+        let bal_dai_spot_price = pool_state
+            .spot_prices
+            .get(&(
+                bytes_to_address(&pool_state.tokens[1]).unwrap(),
+                bytes_to_address(&pool_state.tokens[0]).unwrap(),
+            ))
+            .unwrap();
+        assert_eq!(dai_bal_spot_price, &0.13736685496467538);
+        assert_eq!(bal_dai_spot_price, &7.050354297665408);
+    }
+
     #[tokio::test]
     async fn test_get_balance_overwrites_with_component_balances() {
         let pool_state: EVMPoolState<PreCachedDB> = setup_pool_state().await;

src/evm/simulation.rs

@@ -15,7 +15,7 @@ use revm::{
 use revm_inspectors::tracing::{TracingInspector, TracingInspectorConfig};
 use strum_macros::Display;
 use tokio::runtime::{Handle, Runtime};
-use tracing::{debug, info};
+use tracing::debug;
 
 use super::{
     account_storage::StateUpdate,
@@ -240,7 +240,6 @@ fn interpret_evm_result<DBError: std::fmt::Debug>(
                 gas_used: None,
             }),
             EVMError::Database(db_error) => {
-                info!("Are we at database error? {:?}", &db_error);
                 Err(SimulationEngineError::StorageError(format!("Storage error: {:?}", db_error)))
             }
             EVMError::Custom(err) => Err(SimulationEngineError::TransactionError {