Merge pull request #26 from propeller-heads/tnl/ENG-3756-capabilities

feat: (VMPoolState) set and ensure capabilities

Cargo.toml

@@ -35,6 +35,7 @@ foundry-evm = { git = "https://github.com/foundry-rs/foundry", rev = "2544793" }
 revm-inspectors = { version = "0.5", features = ["serde"] }
 mini-moka = "0.10"
 lazy_static = "1.4.0"
+itertools = "0.10.5"
 
 [dev-dependencies]
 mockito = "1.1.1"

src/evm/simulation.rs

@@ -1,8 +1,5 @@
-// Necessary for the init_account method to be in scope
-#![allow(unused_imports)]
 use std::collections::HashMap;
 
-use crate::evm::engine_db_interface::EngineDatabaseInterface;
 use ethers::types::{Bytes, U256};
 use foundry_config::{Chain, Config};
 use foundry_evm::traces::TraceKind;
@@ -21,6 +18,9 @@ use strum_macros::Display;
 use tokio::runtime::Runtime;
 use tracing::debug;
 
+// Necessary for the init_account method to be in scope
+#[allow(unused_imports)]
+use crate::evm::engine_db_interface::EngineDatabaseInterface;
 use crate::evm::simulation_db::OverriddenSimulationDB;
 
 use super::{

src/protocol/vm/errors.rs

@@ -1,11 +1,13 @@
 // TODO: remove skips for clippy
 
-use crate::evm::simulation::SimulationError;
+use std::io;
+
 use ethers::prelude::ProviderError;
 use serde_json::Error as SerdeError;
-use std::io;
 use thiserror::Error;
 
+use crate::evm::simulation::SimulationError;
+
 /// Represents the outer-level, user-facing errors of the Protosim package.
 ///
 /// `ProtosimError` encompasses all possible errors that can occur in the package,
@@ -17,6 +19,12 @@ use thiserror::Error;
 /// - `SimulationFailure`: Wraps errors that occur during simulation, containing a
 ///   `SimulationError`.
 /// - `DecodingError`: Indicates an error in decoding data.
+/// - `RPCError`: Indicates an error related to RPC interaction.
+/// - `UnsupportedCapability`: Denotes an error when a pool state does not support a necessary
+///   capability.
+/// - `UninitializedAdapter`: Indicates an error when trying to set capabilities before initializing
+///   the adapter.
+/// - `CapabilityRetrievalFailure`: Indicates an error when trying to retrieve capabilities.
 #[derive(Error, Debug)]
 pub enum ProtosimError {
     #[error("ABI loading error: {0}")]
@@ -29,6 +37,10 @@ pub enum ProtosimError {
     DecodingError(String),
     #[error("RPC related error {0}")]
     RpcError(RpcError),
+    #[error("Unsupported Capability: {0}")]
+    UnsupportedCapability(String),
+    #[error("Adapter not initialized: {0}")]
+    UninitializedAdapter(String),
 }
 
 #[derive(Debug, Error)]

src/protocol/vm/models.rs

@@ -1,9 +1,10 @@
 // TODO: remove skip for clippy dead_code check
 use crate::protocol::vm::errors::ProtosimError;
 use ethers::abi::Uint;
+use strum_macros::Display;
 
 #[allow(dead_code)]
-#[derive(Eq, PartialEq, Hash, Debug)]
+#[derive(Eq, PartialEq, Hash, Debug, Display, Clone)]
 pub enum Capability {
     SellSide = 1,
     BuySide = 2,

src/protocol/vm/state.rs

@@ -1,11 +1,28 @@
-// Necessary for the init_account method to be in scope
-#![allow(unused_imports)]
 // TODO: remove skip for clippy dead_code check
 #![allow(dead_code)]
 
+use std::collections::{HashMap, HashSet};
+
+use chrono::Utc;
+use ethers::{
+    abi::{decode, ParamType},
+    prelude::U256,
+    types::H160,
+    utils::to_checksum,
+};
+use revm::{
+    precompile::{Address, Bytes},
+    primitives::{
+        alloy_primitives::Keccak256, keccak256, AccountInfo, Bytecode, B256, KECCAK_EMPTY,
+    },
+    DatabaseRef,
+};
+use tracing::warn;
+
+use itertools::Itertools;
+
 use crate::{
     evm::{
-        engine_db_interface::EngineDatabaseInterface,
         simulation::{SimulationEngine, SimulationParameters},
         simulation_db::BlockHeader,
         tycho_db::PreCachedDB,
@@ -15,26 +32,15 @@ use crate::{
         constants::{ADAPTER_ADDRESS, EXTERNAL_ACCOUNT, MAX_BALANCE},
         engine::{create_engine, SHARED_TYCHO_DB},
         errors::ProtosimError,
+        models::Capability,
         protosim_contract::ProtosimContract,
         utils::{get_code_for_address, get_contract_bytecode},
     },
 };
-use chrono::Utc;
-use ethers::{
-    abi::{decode, Address as EthAddress, ParamType},
-    prelude::U256,
-    types::{Res, H160},
-    utils::to_checksum,
-};
-use revm::{
-    precompile::{Address, Bytes},
-    primitives::{
-        alloy_primitives::Keccak256, keccak256, AccountInfo, Bytecode, B256, KECCAK_EMPTY,
-    },
-    DatabaseRef,
-};
-use std::{collections::HashMap, fmt::Debug, sync::Arc};
-use tokio::sync::RwLock;
+// Necessary for the init_account method to be in scope
+#[allow(unused_imports)]
+use crate::evm::engine_db_interface::EngineDatabaseInterface;
+
 #[derive(Clone)]
 pub struct VMPoolState<D: DatabaseRef + EngineDatabaseInterface + Clone> {
     /// The pool's identifier
@@ -54,6 +60,8 @@ pub struct VMPoolState<D: DatabaseRef + EngineDatabaseInterface + Clone> {
     pub stateless_contracts: HashMap<String, Option<Vec<u8>>>,
     /// If set, vm will emit detailed traces about the execution
     pub trace: bool,
+    /// The supported capabilities of this pool
+    pub capabilities: HashSet<Capability>,
     engine: Option<SimulationEngine<D>>,
     /// The adapter contract. This is used to run simulations
     adapter_contract: Option<ProtosimContract<D>>,
@@ -81,6 +89,7 @@ impl VMPoolState<PreCachedDB> {
             trace,
             engine: None,
             adapter_contract: None,
+            capabilities: HashSet::new(),
         };
         state
             .set_engine()
@@ -93,6 +102,7 @@ impl VMPoolState<PreCachedDB> {
                 .clone()
                 .expect("Engine not set"),
         )?);
+        state.set_capabilities().await?;
         Ok(state)
     }
 
@@ -259,38 +269,85 @@ impl VMPoolState<PreCachedDB> {
             .parse()
             .map_err(|_| ProtosimError::DecodingError("Expected an Address".into()))
     }
+
+    /// Ensures the pool supports the given capability
+    fn ensure_capability(&self, capability: Capability) -> Result<(), ProtosimError> {
+        if !self.capabilities.contains(&capability) {
+            return Err(ProtosimError::UnsupportedCapability(capability.to_string()));
+        }
+        Ok(())
+    }
+
+    async fn set_capabilities(&mut self) -> Result<(), ProtosimError> {
+        let mut capabilities = Vec::new();
+
+        // Generate all permutations of tokens and retrieve capabilities
+        for tokens_pair in self.tokens.iter().permutations(2) {
+            // Manually unpack the inner vector
+            if let [t0, t1] = &tokens_pair[..] {
+                let caps = self
+                    .adapter_contract
+                    .clone()
+                    .ok_or_else(|| {
+                        ProtosimError::UninitializedAdapter(
+                            "Adapter contract must be initialized before setting capabilities"
+                                .to_string(),
+                        )
+                    })?
+                    .get_capabilities(self.id.clone()[2..].to_string(), t0.address, t1.address)
+                    .await?;
+                capabilities.push(caps);
+            }
+        }
+
+        // Find the maximum capabilities length
+        let max_capabilities = capabilities
+            .iter()
+            .map(|c| c.len())
+            .max()
+            .unwrap_or(0);
+
+        // Intersect all capability sets
+        let common_capabilities: HashSet<_> = capabilities
+            .iter()
+            .fold(capabilities[0].clone(), |acc, cap| acc.intersection(cap).cloned().collect());
+
+        self.capabilities = common_capabilities;
+
+        // Check for mismatches in capabilities
+        if self.capabilities.len() < max_capabilities {
+            warn!(
+                "Warning: Pool {} has different capabilities depending on the token pair!",
+                self.id
+            );
+        }
+        Ok(())
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use super::*;
-    use crate::{
-        evm::{simulation_db::BlockHeader, tycho_models::AccountUpdate},
-        models::ERC20Token,
-        protocol::{
-            vm::{models::Capability, utils::maybe_coerce_error},
-            BytesConvertible,
-        },
-    };
-    use ethers::{
-        prelude::{H256, U256},
-        types::Address as EthAddress,
-        utils::hex::traits::FromHex,
-    };
-    use serde_json::Value;
     use std::{
         collections::{HashMap, HashSet},
         fs::File,
-        io::Read,
         path::Path,
         str::FromStr,
     };
-    use tokio::runtime::Runtime;
-    use tycho_core::{dto::ChangeType, Bytes};
+
+    use ethers::prelude::{H256, U256};
+    use serde_json::Value;
+
+    use crate::{
+        evm::{simulation_db::BlockHeader, tycho_models::AccountUpdate},
+        models::ERC20Token,
+        protocol::vm::models::Capability,
+    };
+
+    use super::*;
 
     #[tokio::test]
     async fn test_set_engine_initialization() {
-        let id = "test_pool".to_string();
+        let id = "0x4315fd1afc25cc2ebc72029c543293f9fd833eeb305e2e30159459c827733b1b".to_string();
         let tokens = vec![
             ERC20Token::new(
                 "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48",
@@ -314,7 +371,7 @@ mod tests {
             HashMap::new(),
             "src/protocol/vm/assets/BalancerV2SwapAdapter.evm.runtime".to_string(),
             HashMap::new(),
-            true,
+            false,
         )
         .await;
 
@@ -408,26 +465,45 @@ mod tests {
         )
         .await;
 
-        let capabilities = pool_state
-            .unwrap()
+        let pool_state_initialized = pool_state.unwrap();
+
+        let expected_capabilities = vec![
+            Capability::SellSide,
+            Capability::BuySide,
+            Capability::PriceFunction,
+            Capability::HardLimits,
+        ]
+        .into_iter()
+        .collect::<HashSet<_>>();
+
+        let capabilities_adapter_contract = pool_state_initialized
             .clone()
             .adapter_contract
             .unwrap()
             .get_capabilities(pool_id[2..].to_string(), dai.address, bal.address)
             .await
             .unwrap();
 
-        assert_eq!(
-            capabilities,
-            vec![
-                Capability::SellSide,
-                Capability::BuySide,
-                Capability::PriceFunction,
-                Capability::HardLimits,
-            ]
-            .into_iter()
-            .collect::<HashSet<_>>()
-        );
+        assert_eq!(capabilities_adapter_contract, expected_capabilities.clone());
+
+        let capabilities_state = pool_state_initialized
+            .clone()
+            .capabilities;
+
+        assert_eq!(capabilities_state, expected_capabilities.clone());
+
+        for capability in expected_capabilities.clone() {
+            assert!(pool_state_initialized
+                .clone()
+                .ensure_capability(capability)
+                .is_ok());
+        }
+
+        assert!(pool_state_initialized
+            .clone()
+            .ensure_capability(Capability::MarginalPrice)
+            .is_err());
+
         // // Assert spot prices TODO: in 3757
         // assert_eq!(
         //     pool.spot_prices,