feat: Change ERC20Token address to be alloy's Address instead of H160

Also change the VMPoolState attributes to be Address instead of H160 --- don't change below this line --- ENG-3850 <#DTT#>
propeller-heads · Dec 5, 2024 · b9f55db · b9f55db
1 parent 6adc2a3
commit b9f55db
Show file tree

Hide file tree

Showing 12 changed files with 143 additions and 145 deletions.
diff --git a/examples/explorer/data_feed/tycho.rs b/examples/explorer/data_feed/tycho.rs
@@ -17,7 +17,6 @@ use tycho_client::{
     HttpRPCClient,
 };
 use tycho_core::{dto::Chain, Bytes};
-use tycho_ethereum::BytesCodec;
 use tycho_simulation::{
     evm::{
         engine_db::{
@@ -164,7 +163,7 @@ pub async fn process_messages(
                     .for_each(|(addr, token)| {
                         if token.quality >= 51 {
                             all_tokens
-                                .entry(addr.clone())
+                                .entry(Address::from_slice(addr))
                                 .or_insert_with(|| {
                                     token
                                         .clone()
@@ -185,7 +184,11 @@ pub async fn process_messages(
                         let tokens = comp
                             .tokens
                             .iter()
-                            .flat_map(|addr| all_tokens.get(addr).cloned())
+                            .flat_map(|addr| {
+                                all_tokens
+                                    .get(&Address::from_slice(addr))
+                                    .cloned()
+                            })
                             .collect::<Vec<_>>();
                         let id = Bytes::from_str(id).unwrap_or_else(|_| {
                             panic!("Failed parsing H160 from id string {}", id)
@@ -226,7 +229,7 @@ pub async fn process_messages(
                 let mut skip_pool = false;
 
                 for token in snapshot.component.tokens.clone() {
-                    match all_tokens.get(&token) {
+                    match all_tokens.get(&Address::from_slice(&token)) {
                         Some(token) => pair_tokens.push(token.clone()),
                         None => {
                             debug!(
@@ -337,9 +340,7 @@ pub async fn process_messages(
                                 let tokens: Vec<ERC20Token> = vm_state
                                     .tokens
                                     .iter()
-                                    .filter_map(|token_address| {
-                                        all_tokens.get(&token_address.to_bytes())
-                                    })
+                                    .filter_map(|token_address| all_tokens.get(token_address))
                                     .cloned()
                                     .collect();
 
@@ -361,7 +362,7 @@ pub async fn process_messages(
                                         .downcast_mut::<EVMPoolState<PreCachedDB>>() {
                                         let tokens: Vec<ERC20Token> = vm_state.tokens
                                             .iter()
-                                            .filter_map(|token_address| all_tokens.get(&token_address.to_bytes()))
+                                            .filter_map(|token_address| all_tokens.get(token_address))
                                             .cloned()
                                             .collect();
 
@@ -430,7 +431,7 @@ pub async fn process_messages(
 pub async fn load_all_tokens(
     tycho_url: &str,
     auth_key: Option<&str>,
-) -> HashMap<Bytes, ERC20Token> {
+) -> HashMap<Address, ERC20Token> {
     let rpc_url = format!("https://{tycho_url}");
     let rpc_client = HttpRPCClient::new(rpc_url.as_str(), auth_key).unwrap();
 
@@ -443,7 +444,7 @@ pub async fn load_all_tokens(
         .map(|token| {
             let token_clone = token.clone();
             (
-                token.address.clone(),
+                Address::from_slice(&token.address),
                 token.try_into().unwrap_or_else(|_| {
                     panic!("Couldn't convert {:?} into ERC20 token.", token_clone)
                 }),

diff --git a/src/evm/protocol/uniswap_v2/state.rs b/src/evm/protocol/uniswap_v2/state.rs
@@ -286,7 +286,7 @@ mod tests {
             10_000.to_biguint().unwrap(),
         );
         let weth = ERC20Token::new(
-            "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2 ",
+            "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2",
             18,
             "WETH",
             10_000.to_biguint().unwrap(),

diff --git a/src/evm/protocol/uniswap_v2/tycho_decoder.rs b/src/evm/protocol/uniswap_v2/tycho_decoder.rs
@@ -1,7 +1,6 @@
-use alloy_primitives::U256;
+use alloy_primitives::{Address, U256};
 use std::collections::HashMap;
 use tycho_client::feed::{synchronizer::ComponentWithState, Header};
-use tycho_core::Bytes;
 
 use crate::{
     models::ERC20Token,
@@ -18,7 +17,7 @@ impl TryFromWithBlock<ComponentWithState> for UniswapV2State {
     async fn try_from_with_block(
         snapshot: ComponentWithState,
         _block: Header,
-        _all_tokens: HashMap<Bytes, ERC20Token>,
+        _all_tokens: HashMap<Address, ERC20Token>,
     ) -> Result<Self, Self::Error> {
         let reserve0 = U256::from_be_slice(
             snapshot

diff --git a/src/evm/protocol/uniswap_v3/tycho_decoder.rs b/src/evm/protocol/uniswap_v3/tycho_decoder.rs
@@ -1,4 +1,4 @@
-use alloy_primitives::U256;
+use alloy_primitives::{Address, U256};
 use std::collections::HashMap;
 
 use tycho_client::feed::{synchronizer::ComponentWithState, Header};
@@ -19,7 +19,7 @@ impl TryFromWithBlock<ComponentWithState> for UniswapV3State {
     async fn try_from_with_block(
         snapshot: ComponentWithState,
         _block: Header,
-        _all_tokens: HashMap<Bytes, ERC20Token>,
+        _all_tokens: HashMap<Address, ERC20Token>,
     ) -> Result<Self, Self::Error> {
         let liq = snapshot
             .state

diff --git a/src/evm/protocol/vm/adapter_contract.rs b/src/evm/protocol/vm/adapter_contract.rs
@@ -1,10 +1,11 @@
+use alloy_primitives::Address;
 use std::collections::{HashMap, HashSet};
 
 use ethers::{
-    abi::{Address, Token},
-    types::U256,
+    abi::Token,
+    types::{H160, U256},
 };
-use revm::{primitives::Address as rAddress, DatabaseRef};
+use revm::DatabaseRef;
 
 use crate::{
     evm::{
@@ -51,12 +52,12 @@ where
         buy_token: Address,
         amounts: Vec<U256>,
         block: u64,
-        overwrites: Option<HashMap<rAddress, Overwrites>>,
+        overwrites: Option<HashMap<Address, Overwrites>>,
     ) -> Result<Vec<f64>, SimulationError> {
         let args = vec![
             Token::FixedBytes(pair_id),
-            Token::Address(sell_token),
-            Token::Address(buy_token),
+            Token::Address(H160(sell_token.0 .0)),
+            Token::Address(H160(buy_token.0 .0)),
             Token::Array(
                 amounts
                     .into_iter()
@@ -81,12 +82,12 @@ where
         is_buy: bool,
         amount: U256,
         block: u64,
-        overwrites: Option<HashMap<rAddress, HashMap<U256, U256>>>,
+        overwrites: Option<HashMap<Address, HashMap<U256, U256>>>,
     ) -> Result<(Trade, HashMap<revm::precompile::Address, StateUpdate>), SimulationError> {
         let args = vec![
             Token::FixedBytes(pair_id),
-            Token::Address(sell_token),
-            Token::Address(buy_token),
+            Token::Address(H160(sell_token.0 .0)),
+            Token::Address(H160(buy_token.0 .0)),
             Token::Bool(is_buy),
             Token::Uint(amount),
         ];
@@ -135,10 +136,13 @@ where
         sell_token: Address,
         buy_token: Address,
         block: u64,
-        overwrites: Option<HashMap<rAddress, HashMap<U256, U256>>>,
+        overwrites: Option<HashMap<Address, HashMap<U256, U256>>>,
     ) -> Result<(U256, U256), SimulationError> {
-        let args =
-            vec![Token::FixedBytes(pair_id), Token::Address(sell_token), Token::Address(buy_token)];
+        let args = vec![
+            Token::FixedBytes(pair_id),
+            Token::Address(H160(sell_token.0 .0)),
+            Token::Address(H160(buy_token.0 .0)),
+        ];
 
         let res = self
             .call("getLimits", args, block, None, overwrites, None, U256::zero())?
@@ -164,8 +168,11 @@ where
         sell_token: Address,
         buy_token: Address,
     ) -> Result<HashSet<Capability>, SimulationError> {
-        let args =
-            vec![Token::FixedBytes(pair_id), Token::Address(sell_token), Token::Address(buy_token)];
+        let args = vec![
+            Token::FixedBytes(pair_id),
+            Token::Address(H160(sell_token.0 .0)),
+            Token::Address(H160(buy_token.0 .0)),
+        ];
 
         let res = self
             .call("getCapabilities", args, 1, None, None, None, U256::zero())?

diff --git a/src/evm/protocol/vm/erc20_token.rs b/src/evm/protocol/vm/erc20_token.rs
@@ -1,11 +1,12 @@
+use alloy_primitives::Address;
 use std::{collections::HashMap, str::FromStr};
 
 use ethers::{
-    abi::{Abi, Address, Token},
+    abi::{Abi, Token},
     types::{H160, U256},
 };
 use lazy_static::lazy_static;
-use revm::{primitives::Address as rAddress, DatabaseRef};
+use revm::DatabaseRef;
 use serde_json::from_str;
 
 use crate::{
@@ -40,7 +41,7 @@ impl ERC20Slots {
 pub type Overwrites = HashMap<SlotId, U256>;
 
 pub struct ERC20OverwriteFactory {
-    token_address: rAddress,
+    token_address: Address,
     overwrites: Overwrites,
     balance_slot: SlotId,
     allowance_slot: SlotId,
@@ -50,7 +51,7 @@ pub struct ERC20OverwriteFactory {
 
 impl ERC20OverwriteFactory {
     pub fn new(
-        token_address: rAddress,
+        token_address: Address,
         token_slots: ERC20Slots,
         compiler: ContractCompiler,
     ) -> Self {
@@ -84,7 +85,7 @@ impl ERC20OverwriteFactory {
             .insert(self.total_supply_slot, supply);
     }
 
-    pub fn get_overwrites(&self) -> HashMap<rAddress, Overwrites> {
+    pub fn get_overwrites(&self) -> HashMap<Address, Overwrites> {
         let mut result = HashMap::new();
         result.insert(self.token_address, self.overwrites.clone());
         result
@@ -138,39 +139,33 @@ lazy_static! {
 /// - Once the balance slot is found, it uses the detected compiler to search for the allowance
 ///   slot, which is dependent on the balance slot.
 pub fn brute_force_slots<D: EngineDatabaseInterface + Clone>(
-    token_addr: &H160,
+    token_addr: &Address,
     block: &BlockHeader,
     engine: &SimulationEngine<D>,
 ) -> Result<(ERC20Slots, ContractCompiler), SimulationError>
 where
     <D as DatabaseRef>::Error: std::fmt::Debug,
     <D as EngineDatabaseInterface>::Error: std::fmt::Debug,
 {
-    let token_contract = TychoSimulationContract::new(
-        rAddress::from_slice(token_addr.as_bytes()),
-        engine.clone(),
-        ERC20_ABI.clone(),
-    )
-    .unwrap();
-
-    let external_account = H160::from_slice(&*EXTERNAL_ACCOUNT.0);
+    let token_contract =
+        TychoSimulationContract::new(*token_addr, engine.clone(), ERC20_ABI.clone()).unwrap();
 
     let mut compiler = ContractCompiler::Solidity;
 
     let mut balance_slot = None;
     for i in 0..100 {
         for compiler_flag in [ContractCompiler::Solidity, ContractCompiler::Vyper] {
             let mut overwrite_factory = ERC20OverwriteFactory::new(
-                rAddress::from_slice(token_addr.as_bytes()),
+                *token_addr,
                 ERC20Slots::new(i.into(), 1.into()),
                 compiler_flag,
             );
-            overwrite_factory.set_balance(MARKER_VALUE.into(), external_account);
+            overwrite_factory.set_balance(MARKER_VALUE.into(), *EXTERNAL_ACCOUNT);
 
             let res = token_contract
                 .call(
                     "balanceOf",
-                    vec![Token::Address(external_account)],
+                    vec![Token::Address(H160::from_slice(&*EXTERNAL_ACCOUNT.0))],
                     block.number,
                     Some(block.timestamp),
                     Some(overwrite_factory.get_overwrites()),
@@ -197,23 +192,23 @@ where
     let mut allowance_slot = None;
     for i in 0..100 {
         let mut overwrite_factory = ERC20OverwriteFactory::new(
-            rAddress::from_slice(token_addr.as_bytes()),
+            *token_addr,
             ERC20Slots::new(0.into(), i.into()),
             compiler, /* At this point we know the compiler becase we managed to find the
                        * balance slot */
         );
 
         overwrite_factory.set_allowance(
             MARKER_VALUE.into(),
-            H160::from_str(SPENDER).unwrap(),
-            external_account,
+            Address::from_str(SPENDER).unwrap(),
+            *EXTERNAL_ACCOUNT,
         );
 
         let res = token_contract
             .call(
                 "allowance",
                 vec![
-                    Token::Address(external_account),
+                    Token::Address(H160::from_slice(&*EXTERNAL_ACCOUNT.0)),
                     Token::Address(H160::from_str(SPENDER).unwrap()),
                 ],
                 block.number,
@@ -252,7 +247,7 @@ mod tests {
     use crate::evm::engine_db::simulation_db::SimulationDB;
 
     fn setup_factory() -> ERC20OverwriteFactory {
-        let token_address: rAddress = rAddress::from_slice(
+        let token_address: Address = Address::from_slice(
             &hex::decode("C02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2")
                 .expect("Invalid token address"),
         );
@@ -336,7 +331,7 @@ mod tests {
         };
 
         let (slots, compiler) = brute_force_slots(
-            &H160::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(),
+            &Address::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(),
             &block,
             &eng,
         )
@@ -365,7 +360,7 @@ mod tests {
         };
 
         let (slots, compiler) = brute_force_slots(
-            &H160::from_str("0xa5588f7cdf560811710a2d82d3c9c99769db1dcb").unwrap(),
+            &Address::from_str("0xa5588f7cdf560811710a2d82d3c9c99769db1dcb").unwrap(),
             &block,
             &eng,
         )