feat: create ERC20OverwriteFactory utils

src/protocol/vm/utils.rs

@@ -2,12 +2,11 @@
 #![allow(dead_code)]
 use ethabi::{self, decode, ParamType};
 
-use std::path::PathBuf;
 use ethers::{
+    abi::Abi,
     core::utils::keccak256,
     types::{Address, H256, U256},
 };
-use ethers::abi::Abi;
 use hex::FromHex;
 use mini_moka::sync::Cache;
 use reqwest::{blocking::Client, StatusCode};
@@ -131,55 +130,60 @@ fn parse_solidity_error_message(data: &str) -> String {
     format!("Failed to decode: {}", data)
 }
 
-pub fn get_storage_slot_at_key(key: Address, mapping_slot: usize) -> U256 {
+pub fn get_storage_slot_at_key(key: Address, mapping_slot: H256) -> H256 {
     // Convert key to bytes
     let mut key_bytes = key.as_bytes().to_vec();
     key_bytes.resize(32, 0); // Right pad with zeros
 
     // Convert mapping slot to bytes
     let mut mapping_slot_bytes = [0u8; 32];
-    mapping_slot_bytes[32 - std::mem::size_of::<usize>()..]
-        .copy_from_slice(&mapping_slot.to_be_bytes());
+    mapping_slot_bytes.copy_from_slice(mapping_slot.as_bytes());
 
     // Concatenate key and mapping slot bytes, then hash
     let slot_bytes = keccak256([&key_bytes[..], &mapping_slot_bytes[..]].concat());
+    ethers::types::TxHash(slot_bytes)
+}
 
-    // Convert the hashed bytes to a U256
-    U256::from_big_endian(&slot_bytes)
+pub struct GethOverwrite {
+    pub state_diff: HashMap<String, String>, // The formatted overwrites
+    pub code: String,                        // The bytecode as a string
 }
 
 pub struct ERC20OverwriteFactory {
     token_address: Address,
     overwrites: HashMap<H256, U256>,
-    balance_slot: u64,
-    allowance_slot: u64,
-    total_supply_slot: u64,
+    balance_slot: H256,
+    allowance_slot: H256,
+    total_supply_slot: H256,
 }
 
 impl ERC20OverwriteFactory {
-    pub fn new(token_address: Address, token_slots: (u64, u64)) -> Self {
+    pub fn new(token_address: Address, token_slots: (H256, H256)) -> Self {
         ERC20OverwriteFactory {
             token_address,
             overwrites: HashMap::new(),
             balance_slot: token_slots.0,
             allowance_slot: token_slots.1,
-            total_supply_slot: 2,
+            total_supply_slot: H256::from_low_u64_be(2),
         }
     }
 
     pub fn set_balance(&mut self, balance: U256, owner: Address) {
-        let storage_index = get_storage_slot_at_key(owner, self.balance_slot as usize);
-        self.overwrites.insert(H256::from(storage_index), balance);
+        let storage_index = get_storage_slot_at_key(owner, self.balance_slot);
+        self.overwrites
+            .insert(storage_index, balance);
     }
 
     pub fn set_allowance(&mut self, allowance: U256, spender: Address, owner: Address) {
-        let owner_slot = get_storage_slot_at_key(owner, self.allowance_slot as usize);
-        let storage_index = get_storage_slot_at_key(spender, owner_slot.as_usize());
-        self.overwrites.insert(H256::from(storage_index), allowance);
+        let owner_slot = get_storage_slot_at_key(owner, self.allowance_slot);
+        let storage_index = get_storage_slot_at_key(spender, owner_slot);
+        self.overwrites
+            .insert(storage_index, allowance);
     }
 
     pub fn set_total_supply(&mut self, supply: U256) {
-        self.overwrites.insert(H256::from_low_u64_be(self.total_supply_slot), supply);
+        self.overwrites
+            .insert(self.total_supply_slot, supply);
     }
 
     pub fn get_protosim_overwrites(&self) -> HashMap<Address, HashMap<H256, U256>> {
@@ -188,24 +192,32 @@ impl ERC20OverwriteFactory {
         result
     }
 
-    pub fn get_geth_overwrites(&self) -> HashMap<Address, HashMap<String, String>> {
+    pub fn get_geth_overwrites(&self) -> HashMap<Address, GethOverwrite> {
         let mut formatted_overwrites = HashMap::new();
+
         for (key, val) in &self.overwrites {
-            let formatted_key = format!("0x{}", key.to_hex::<String>());
-            let formatted_val = format!("0x{:0>64}", val.to_hex::<String>());
-            formatted_overwrites.insert(formatted_key, formatted_val);
+            let hex_key = hex::encode(key.as_bytes());
+
+            let mut bytes = [0u8; 32];
+            val.to_big_endian(&mut bytes);
+            let hex_val = format!("0x{:0>64}", hex::encode(bytes));
+
+            formatted_overwrites.insert(hex_key, hex_val);
         }
 
-        let code = match get_contract_bytecode(&PathBuf::from(ASSETS_FOLDER).join("ERC20.bin")) {
-            Ok(bytes) => format!("0x{}", hex::encode(bytes)),
-            Err(_) => "0x".to_string(), // Handle error appropriately in your actual code
-        };
+        let erc20_abi_path = Path::new(file!())
+            .parent()
+            .unwrap()
+            .join("assets")
+            .join("ERC20.abi");
+
+        let code = format!(
+            "0x{}",
+            hex::encode(get_contract_bytecode(erc20_abi_path.to_str().unwrap()).unwrap())
+        );
 
         let mut result = HashMap::new();
-        let mut inner = HashMap::new();
-        inner.insert("stateDiff".to_string(), formatted_overwrites);
-        inner.insert("code".to_string(), code);
-        result.insert(self.token_address, inner);
+        result.insert(self.token_address, GethOverwrite { state_diff: formatted_overwrites, code });
 
         result
     }
@@ -523,4 +535,99 @@ mod tests {
         let abi: Abi = result.expect("Failed to retrieve ERC20 ABI result");
         assert!(!abi.functions.is_empty(), "The ERC20 ABI should contain functions.");
     }
+
+    fn setup_factory() -> ERC20OverwriteFactory {
+        let token_address = Address::random();
+        let balance_slot = H256::random();
+        let allowance_slot = H256::random();
+        ERC20OverwriteFactory::new(token_address, (balance_slot, allowance_slot))
+    }
+
+    #[test]
+    fn test_set_balance() {
+        let mut factory = setup_factory();
+        let owner = Address::random();
+        let balance = U256::from(1000);
+
+        factory.set_balance(balance, owner);
+
+        assert_eq!(factory.overwrites.len(), 1);
+        assert!(factory
+            .overwrites
+            .values()
+            .any(|&v| v == balance));
+    }
+
+    #[test]
+    fn test_set_allowance() {
+        let mut factory = setup_factory();
+        let owner = Address::random();
+        let spender = Address::random();
+        let allowance = U256::from(500);
+
+        factory.set_allowance(allowance, spender, owner);
+
+        assert_eq!(factory.overwrites.len(), 1);
+        assert!(factory
+            .overwrites
+            .values()
+            .any(|&v| v == allowance));
+    }
+
+    #[test]
+    fn test_set_total_supply() {
+        let mut factory = setup_factory();
+        let supply = U256::from(1_000_000);
+
+        factory.set_total_supply(supply);
+
+        assert_eq!(factory.overwrites.len(), 1);
+        assert_eq!(factory.overwrites[&factory.total_supply_slot], supply);
+    }
+
+    #[test]
+    fn test_get_protosim_overwrites() {
+        let mut factory = setup_factory();
+        let supply = U256::from(1_000_000);
+        factory.set_total_supply(supply);
+
+        let overwrites = factory.get_protosim_overwrites();
+
+        assert_eq!(overwrites.len(), 1);
+        assert!(overwrites.contains_key(&factory.token_address));
+        assert_eq!(overwrites[&factory.token_address].len(), 1);
+        assert_eq!(overwrites[&factory.token_address][&factory.total_supply_slot], supply);
+    }
+
+    #[test]
+    fn test_get_geth_overwrites() {
+        let mut factory = setup_factory();
+
+        let storage_slot = H256::from_low_u64_be(1);
+        let val = U256::from(123456);
+        factory
+            .overwrites
+            .insert(storage_slot, val);
+
+        let result = factory.get_geth_overwrites();
+
+        assert_eq!(result.len(), 1);
+
+        let geth_overwrite = result
+            .get(&factory.token_address)
+            .expect("Missing token address");
+        assert_eq!(geth_overwrite.state_diff.len(), 1);
+
+        let expected_key =
+            String::from("0000000000000000000000000000000000000000000000000000000000000001");
+        let expected_val =
+            String::from("0x000000000000000000000000000000000000000000000000000000000001e240");
+        assert_eq!(
+            geth_overwrite
+                .state_diff
+                .get(&expected_key),
+            Some(&expected_val)
+        );
+        assert_eq!(geth_overwrite.code.len(), 8752);
+    }
 }