feat: create ERC20OverwriteFactory utils

src/protocol/vm/utils.rs

@@ -1,7 +1,12 @@
 // TODO: remove skip for clippy dead_code check
 #![allow(dead_code)]
 use ethabi::{self, decode, ParamType};
-use ethers::abi::Abi;
+
+use ethers::{
+    abi::Abi,
+    core::utils::keccak256,
+    types::{Address, H256, U256},
+};
 use hex::FromHex;
 use mini_moka::sync::Cache;
 use reqwest::{blocking::Client, StatusCode};
@@ -15,7 +20,6 @@ use std::{
     sync::{Arc, LazyLock},
     time::Duration,
 };
-
 use thiserror::Error;
 
 #[derive(Debug, Error)]
@@ -126,6 +130,99 @@ fn parse_solidity_error_message(data: &str) -> String {
     format!("Failed to decode: {}", data)
 }
 
+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.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)
+}
+
+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: H256,
+    allowance_slot: H256,
+    total_supply_slot: H256,
+}
+
+impl ERC20OverwriteFactory {
+    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: 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);
+        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);
+        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(self.total_supply_slot, supply);
+    }
+
+    pub fn get_protosim_overwrites(&self) -> HashMap<Address, HashMap<H256, U256>> {
+        let mut result = HashMap::new();
+        result.insert(self.token_address, self.overwrites.clone());
+        result
+    }
+
+    pub fn get_geth_overwrites(&self) -> HashMap<Address, GethOverwrite> {
+        let mut formatted_overwrites = HashMap::new();
+
+        for (key, val) in &self.overwrites {
+            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 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();
+        result.insert(self.token_address, GethOverwrite { state_diff: formatted_overwrites, code });
+
+        result
+    }
+}
+
 fn get_solidity_panic_codes() -> HashMap<u64, String> {
     let mut panic_codes = HashMap::new();
     panic_codes.insert(0, "GenericCompilerPanic".to_string());
@@ -438,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);
+    }
 }