Post rebase fixes

propeller-heads · Dec 12, 2024 · 049d505 · 049d505
1 parent 3c9e182
commit 049d505
Show file tree

Hide file tree

Showing 9 changed files with 62 additions and 66 deletions.
diff --git a/examples/explorer/utils.rs b/examples/explorer/utils.rs
@@ -3,12 +3,12 @@ use std::collections::HashMap;
 use tracing_subscriber::{fmt, EnvFilter};
 use tycho_client::{rpc::RPCClient, HttpRPCClient};
 use tycho_core::{dto::Chain, Bytes};
-use tycho_simulation::models::ERC20Token;
+use tycho_simulation::models::Token;
 
 pub async fn load_all_tokens(
     tycho_url: &str,
     auth_key: Option<&str>,
-) -> HashMap<Bytes, ERC20Token> {
+) -> HashMap<Bytes, Token> {
     let rpc_url = format!("https://{tycho_url}");
     let rpc_client = HttpRPCClient::new(rpc_url.as_str(), auth_key).unwrap();
 
@@ -27,7 +27,7 @@ pub async fn load_all_tokens(
                 }),
             )
         })
-        .collect::<HashMap<_, ERC20Token>>()
+        .collect::<HashMap<_, Token>>()
 }
 
 pub fn setup_tracing() {

diff --git a/src/evm/protocol/uniswap_v3/state.rs b/src/evm/protocol/uniswap_v3/state.rs
@@ -285,7 +285,7 @@ impl ProtocolSim for UniswapV3State {
     fn delta_transition(
         &mut self,
         delta: ProtocolStateDelta,
-        _tokens: Vec<Token>,
+        _tokens: &HashMap<Bytes, Token>,
     ) -> Result<(), TransitionError<String>> {
         // apply attribute changes
         if let Some(liquidity) = delta

diff --git a/src/evm/protocol/vm/state.rs b/src/evm/protocol/vm/state.rs
@@ -135,12 +135,11 @@ where
         for [sell_token_address, buy_token_address] in self
             .tokens
             .iter()
-            .copied()
             .permutations(2)
             .map(|p| [p[0], p[1]])
         {
-            let sell_token_address = bytes_to_erc20_address(&sell_token.address)?;
-            let buy_token_address = bytes_to_erc20_address(&buy_token.address)?;
+            let sell_token_address = bytes_to_erc20_address(sell_token_address)?;
+            let buy_token_address = bytes_to_erc20_address(buy_token_address)?;
             let overwrites = Some(self.get_overwrites(
                 vec![sell_token_address, buy_token_address],
                 *MAX_BALANCE / U256::from(100),
@@ -183,7 +182,7 @@ where
 
     fn get_decimals(
         &self,
-        tokens: &HashMap<Bytes, ERC20Token>,
+        tokens: &HashMap<Bytes, Token>,
         sell_token_address: &Address,
     ) -> Result<usize, SimulationError> {
         tokens
@@ -216,10 +215,7 @@ where
         Ok(limits?.0)
     }
 
-    fn clear_all_cache(
-        &mut self,
-        tokens: &HashMap<Bytes, Token>,
-    ) -> Result<(), SimulationError> {
+    fn clear_all_cache(&mut self, tokens: &HashMap<Bytes, Token>) -> Result<(), SimulationError> {
         self.adapter_contract
             .engine
             .clear_temp_storage();
@@ -847,7 +843,7 @@ mod tests {
             .set_spot_prices(
                 &vec![bal(), dai()]
                     .into_iter()
-                    .map(|t| (Bytes::from(t.address.as_slice()), t))
+                    .map(|t| (t.address.clone(), t))
                     .collect(),
             )
             .unwrap();

diff --git a/src/evm/protocol/vm/tycho_decoder.rs b/src/evm/protocol/vm/tycho_decoder.rs
@@ -48,12 +48,7 @@ impl TryFromWithBlock<ComponentWithState> for EVMPoolState<PreCachedDB> {
         all_tokens: &HashMap<Bytes, Token>,
     ) -> Result<Self, Self::Error> {
         let id = snapshot.component.id.clone();
-        let tokens: Vec<Bytes> = snapshot
-            .component
-            .tokens
-            .iter()
-            .map(|x| Address::from_slice(x.as_ref()))
-            .collect();
+        let tokens = snapshot.component.tokens.clone();
 
         let block = BlockHeader::from(block);
         let balances = snapshot
@@ -263,21 +258,21 @@ mod tests {
         .into_iter()
         .collect();
         let tokens = [
-            ERC20Token::new(
+            Token::new(
                 "0x6b175474e89094c44da98b954eedeac495271d0f",
                 18,
                 "DAI",
                 10_000.to_biguint().unwrap(),
             ),
-            ERC20Token::new(
+            Token::new(
                 "0xba100000625a3754423978a60c9317c58a424e3d",
                 18,
                 "BAL",
                 10_000.to_biguint().unwrap(),
             ),
         ]
         .into_iter()
-        .map(|t| (Bytes::from(t.address.as_slice()), t))
+        .map(|t| (t.address.clone(), t))
         .collect::<HashMap<_, _>>();
         let snapshot = ComponentWithState {
             state: ResponseProtocolState {

diff --git a/src/evm/protocol/vm/utils.rs b/src/evm/protocol/vm/utils.rs
@@ -10,7 +10,7 @@ use std::{
 
 use alloy::{
     providers::{Provider, ProviderBuilder},
-    transports::RpcError,
+    transports::{RpcError, TransportErrorKind},
 };
 use alloy_primitives::{Address, U256};
 use alloy_sol_types::SolValue;
@@ -245,36 +245,46 @@ pub(crate) async fn get_code_for_contract(
         }
     };
 
-    // Create a provider with the URL
-    let provider = ProviderBuilder::new()
-        .on_builtin(&connection_string)
-        .await
-        .unwrap();
-
     let addr = Address::from_str(address).expect("Failed to parse address to get code for");
-
     // Call eth_getCode to get the bytecode of the contract
-    match provider.get_code_at(addr, None).await {
-        Ok(code) if code.is_empty() => Err(SimulationError::FatalError("Empty code response from RPC".to_string())),
+    match sync_get_code(&connection_string, addr) {
+        Ok(code) if code.is_empty() => {
+            Err(SimulationError::FatalError("Empty code response from RPC".to_string()))
+        }
         Ok(code) => {
             let bytecode = Bytecode::new_raw(Bytes::from(code.to_vec()));
             Ok(bytecode)
         }
-        Err(e) => {
-            match e {
-                RpcError::Transport(err) => Err(SimulationError::RecoverableError(format!(
-                    "Failed to get code for contract due to internal RPC error: {:?}",
-                    err
-                ))),
-                _ => Err(SimulationError::FatalError(format!(
-                    "Failed to get code for contract. Invalid response from RPC: {:?}",
-                    e
-                ))),
-            }
-        }
+        Err(e) => match e {
+            RpcError::Transport(err) => Err(SimulationError::RecoverableError(format!(
+                "Failed to get code for contract due to internal RPC error: {:?}",
+                err
+            ))),
+            _ => Err(SimulationError::FatalError(format!(
+                "Failed to get code for contract. Invalid response from RPC: {:?}",
+                e
+            ))),
+        },
     }
 }
 
+fn sync_get_code(
+    connection_string: &str,
+    addr: Address,
+) -> Result<Bytes, RpcError<TransportErrorKind>> {
+    tokio::task::block_in_place(|| {
+        tokio::runtime::Handle::current().block_on(async {
+            // Create a provider with the URL
+            let provider = ProviderBuilder::new()
+                .on_builtin(connection_string)
+                .await
+                .unwrap();
+
+            provider.get_code_at(addr).await
+        })
+    })
+}
+
 static BYTECODE_CACHE: LazyLock<Cache<Arc<String>, Bytecode>> = LazyLock::new(|| Cache::new(1_000));
 
 pub(crate) fn get_contract_bytecode(path: &PathBuf) -> Result<Bytecode, FileError> {

diff --git a/src/protocol/decoder.rs b/src/protocol/decoder.rs
@@ -3,7 +3,7 @@ use crate::{
         engine_db::{update_engine, SHARED_TYCHO_DB},
         tycho_models::{AccountUpdate, ResponseAccount},
     },
-    models::ERC20Token,
+    models::Token,
     protocol::{
         errors::InvalidSnapshotError,
         models::{BlockUpdate, ProtocolComponent, TryFromWithBlock},
@@ -32,7 +32,7 @@ pub enum StreamDecodeError {
 
 #[derive(Default)]
 struct DecoderState {
-    tokens: HashMap<Bytes, ERC20Token>,
+    tokens: HashMap<Bytes, Token>,
     states: HashMap<String, Box<dyn ProtocolSim>>,
 }
 
@@ -61,7 +61,7 @@ impl TychoStreamDecoder {
         }
     }
 
-    pub fn set_tokens(&self, tokens: HashMap<Bytes, ERC20Token>) {
+    pub fn set_tokens(&self, tokens: HashMap<Bytes, Token>) {
         let state = self.state.clone();
         // We prefer this to be sync since it is used mainly in builders.
         tokio::task::spawn_blocking(move || {
@@ -127,8 +127,7 @@ impl TychoStreamDecoder {
                             return None;
                         }
 
-                        let token: Result<ERC20Token, std::num::TryFromIntError> =
-                            t.clone().try_into();
+                        let token: Result<Token, std::num::TryFromIntError> = t.clone().try_into();
                         let result = match token {
                             Ok(t) => Ok((addr.clone(), t)),
                             Err(e) => Err(StreamDecodeError::Fatal(format!(
@@ -137,7 +136,7 @@ impl TychoStreamDecoder {
                         };
                         Some(result)
                     })
-                    .collect::<Result<HashMap<Bytes, ERC20Token>, StreamDecodeError>>()?;
+                    .collect::<Result<HashMap<Bytes, Token>, StreamDecodeError>>()?;
 
                 if !res.is_empty() {
                     debug!(n = res.len(), "NewTokens");
@@ -207,7 +206,10 @@ impl TychoStreamDecoder {
                 .clone()
             {
                 // Skip any unsupported pools
-                if let Some(predicate) = self.inclusion_filters.get(protocol.as_str()) {
+                if let Some(predicate) = self
+                    .inclusion_filters
+                    .get(protocol.as_str())
+                {
                     if !predicate(&snapshot) {
                         continue
                     }
@@ -331,7 +333,7 @@ impl TychoStreamDecoder {
 mod tests {
     use crate::{
         evm::protocol::uniswap_v2::state::UniswapV2State,
-        models::ERC20Token,
+        models::Token,
         protocol::decoder::{StreamDecodeError, TychoStreamDecoder},
     };
     use num_bigint::ToBigUint;
@@ -351,10 +353,7 @@ mod tests {
             .iter()
             .map(|addr| {
                 let addr_str = format!("{:x}", addr);
-                (
-                    addr.clone(),
-                    ERC20Token::new(&addr_str, 18, &addr_str, 100_000.to_biguint().unwrap()),
-                )
+                (addr.clone(), Token::new(&addr_str, 18, &addr_str, 100_000.to_biguint().unwrap()))
             })
             .collect();
             decoder.set_tokens(tokens);
@@ -396,10 +395,7 @@ mod tests {
             .iter()
             .map(|addr| {
                 let addr_str = format!("{:x}", addr);
-                (
-                    addr.clone(),
-                    ERC20Token::new(&addr_str, 18, &addr_str, 100_000.to_biguint().unwrap()),
-                )
+                (addr.clone(), Token::new(&addr_str, 18, &addr_str, 100_000.to_biguint().unwrap()))
             })
             .collect();
         decoder.set_tokens(tokens);

diff --git a/src/protocol/models.rs b/src/protocol/models.rs
@@ -24,8 +24,7 @@
 //! It's worth emphasizin that although the term "pair" used in this
 //! module refers to a trading pair, it does not necessarily imply two
 //! tokens only. Some pairs might have more than two tokens.
-use std::collections::HashMap;
-use std::future::Future;
+use std::{collections::HashMap, future::Future};
 
 use num_bigint::BigUint;
 
@@ -62,7 +61,7 @@ pub trait TryFromWithBlock<T> {
     fn try_from_with_block(
         value: T,
         block: Header,
-        all_tokens: HashMap<Bytes, Token>,
+        all_tokens: &HashMap<Bytes, Token>,
     ) -> impl Future<Output = Result<Self, Self::Error>> + Send + Sync
     where
         Self: Sized;

diff --git a/src/protocol/state.rs b/src/protocol/state.rs
@@ -117,7 +117,7 @@ pub trait ProtocolSim: std::fmt::Debug + Send + Sync + 'static {
     fn delta_transition(
         &mut self,
         delta: ProtocolStateDelta,
-        tokens: Vec<Token>,
+        tokens: &HashMap<Bytes, Token>,
     ) -> Result<(), TransitionError<String>>;
 
     /// Clones the protocol state as a trait object.

diff --git a/src/protocol/stream.rs b/src/protocol/stream.rs
@@ -1,5 +1,5 @@
 use crate::{
-    models::ERC20Token,
+    models::Token,
     protocol::{
         decoder::{StreamDecodeError, TychoStreamDecoder},
         errors::InvalidSnapshotError,
@@ -79,7 +79,7 @@ impl ProtocolStreamBuilder {
         self
     }
 
-    pub fn set_tokens(self, tokens: HashMap<Bytes, ERC20Token>) -> Self {
+    pub fn set_tokens(self, tokens: HashMap<Bytes, Token>) -> Self {
         self.decoder.set_tokens(tokens);
         self
     }