feat: create ERC20OverwriteFactory utils

src/protocol/mod.rs

@@ -10,7 +10,7 @@ pub mod state;
 pub mod tycho;
 pub mod uniswap_v2;
 pub mod uniswap_v3;
-mod vm;
+pub mod vm;
 
 /// A trait for converting types to and from `Bytes`.
 ///

src/protocol/vm/mod.rs

@@ -1 +1 @@
-mod utils;
+pub mod 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, H160, 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)]
@@ -30,6 +34,19 @@ pub enum RpcError {
     OutOfGas(String, String),
 }
 
+#[derive(Debug, Error)]
+pub enum FileError {
+    /// Occurs when the ABI file cannot be read
+    #[error("Malformed ABI error: {0}")]
+    MalformedABIError(String),
+    /// Occurs when the parent directory of the current file cannot be retrieved
+    #[error("Structure error {0}")]
+    StructureError(String),
+    /// Occurs when a bad file path was given, which cannot be converted to string.
+    #[error("File path conversion error {0}")]
+    FilePathError(String),
+}
+
 pub fn maybe_coerce_error(
     err: RpcError,
     pool_state: &str,
@@ -126,6 +143,154 @@ fn parse_solidity_error_message(data: &str) -> String {
     format!("Failed to decode: {}", data)
 }
 
+/// Get storage slot index of a value stored at a certain key in a mapping
+///
+/// # Arguments
+///
+/// * `key`: Key in a mapping. Can be any H160 value (such as an address).
+/// * `mapping_slot`: Storage slot at which the mapping itself is stored. See the examples for more
+///   explanation.
+///
+/// # Returns
+///
+/// An `H256` representing the  index of a storage slot where the value at the given
+/// key is stored.
+///
+/// # Examples
+///
+/// If a mapping is declared as a first variable in Solidity code, its storage slot
+/// is 0 (e.g. `balances` in our mocked ERC20 contract). Here's how to compute
+/// a storage slot where balance of a given account is stored:
+///
+/// ```
+/// use protosim::protocol::vm::utils::{get_storage_slot_index_at_key, H256};
+/// use ethers::types::Address;
+/// let address: Address = "0xC63135E4bF73F637AF616DFd64cf701866BB2628".parse().expect("Invalid address");
+/// get_storage_slot_index_at_key(address, H256::from_low_u64_be(0));
+/// ```
+///
+/// For nested mappings, we need to apply the function twice. An example of this is
+/// `allowances` in ERC20. It is a mapping of form:
+/// `HashMap<Owner, HashMap<Spender, U256>>`. In our mocked ERC20 contract, `allowances`
+/// is a second variable, so it is stored at slot 1. Here's how to get a storage slot
+/// where an allowance of `address_spender` to spend `address_owner`'s money is stored:
+///
+/// ```
+/// use protosim::protocol::vm::utils::{get_storage_slot_index_at_key, H256};
+/// use ethers::types::Address;
+/// let address_spender: Address = "0xC63135E4bF73F637AF616DFd64cf701866BB2628".parse().expect("Invalid address");
+/// let address_owner: Address = "0x6F4Feb566b0f29e2edC231aDF88Fe7e1169D7c05".parse().expect("Invalid address");
+/// get_storage_slot_index_at_key(address_spender, get_storage_slot_index_at_key(address_owner, H256::from_low_u64_be(1)));
+/// ```
+///
+/// # See Also
+///
+/// [Solidity Storage Layout documentation](https://docs.soliditylang.org/en/v0.8.13/internals/layout_in_storage.html#mappings-and-dynamic-arrays)
+pub fn get_storage_slot_index_at_key(key: H160, mapping_slot: H256) -> H256 {
+    let mut key_bytes = key.as_bytes().to_vec();
+    key_bytes.resize(32, 0); // Right pad with zeros
+
+    let mut mapping_slot_bytes = [0u8; 32];
+    mapping_slot_bytes.copy_from_slice(mapping_slot.as_bytes());
+
+    let slot_bytes = keccak256([&key_bytes[..], &mapping_slot_bytes[..]].concat());
+    H256::from_slice(&slot_bytes)
+}
+
+pub struct GethOverwrite {
+    /// the formatted overwrites
+    pub state_diff: HashMap<String, String>,
+    /// the bytecode as a string
+    pub code: String,
+}
+
+pub type Overwrites = HashMap<H256, U256>;
+
+pub struct ERC20OverwriteFactory {
+    token_address: Address,
+    overwrites: Overwrites,
+    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_index_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_index_at_key(owner, self.allowance_slot);
+        let storage_index = get_storage_slot_index_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, Overwrites> {
+        let mut result = HashMap::new();
+        result.insert(self.token_address, self.overwrites.clone());
+        result
+    }
+
+    pub fn get_geth_overwrites(&self) -> Result<HashMap<Address, GethOverwrite>, FileError> {
+        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()
+            .ok_or_else(|| {
+                FileError::StructureError(
+                    "Failed to obtain parent directory of current file.".to_string(),
+                )
+            })?
+            .join("assets")
+            .join("ERC20.abi");
+
+        let code = format!(
+            "0x{}",
+            hex::encode(
+                get_contract_bytecode(erc20_abi_path.to_str().ok_or_else(|| {
+                    FileError::FilePathError("Failed to convert file path to string.".to_string())
+                })?)
+                .map_err(|_err| FileError::MalformedABIError(
+                    "Failed to read contract bytecode.".to_string()
+                ))?
+            )
+        );
+
+        let mut result = HashMap::new();
+        result.insert(self.token_address, GethOverwrite { state_diff: formatted_overwrites, code });
+
+        Ok(result)
+    }
+}
+
 fn get_solidity_panic_codes() -> HashMap<u64, String> {
     let mut panic_codes = HashMap::new();
     panic_codes.insert(0, "GenericCompilerPanic".to_string());
@@ -438,4 +603,101 @@ 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()
+            .expect("Failed to 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);
+    }
 }