fix(cheatcodes): identify common proxies in state diffs (#11404)

* fix(cheatcodes): identify proxies in state diffs

* dedup get_contract_name and get_contract_data

* clippy

Author: yash-atreya

crates/cheatcodes/src/evm.rs

@@ -1346,8 +1346,8 @@ fn get_recorded_state_diffs(ccx: &mut CheatsCtxt) -> BTreeMap<Address, AccountSt
     let mut contract_names = HashMap::new();
     let mut storage_layouts = HashMap::new();
     for address in addresses_to_lookup {
-        if let Some(name) = get_contract_name(ccx, address) {
-            contract_names.insert(address, name);
+        if let Some((artifact_id, _)) = get_contract_data(ccx, address) {
+            contract_names.insert(address, artifact_id.identifier());
         }
 
         // Also get storage layout if available
@@ -1471,33 +1471,12 @@ fn get_recorded_state_diffs(ccx: &mut CheatsCtxt) -> BTreeMap<Address, AccountSt
     state_diffs
 }
 
-/// Helper function to get the contract name from the deployed code.
-fn get_contract_name(ccx: &mut CheatsCtxt, address: Address) -> Option<String> {
-    // Check if we have available artifacts to match against
-    let artifacts = ccx.state.config.available_artifacts.as_ref()?;
-
-    // Try to load the account and get its code
-    let account = ccx.ecx.journaled_state.load_account(address).ok()?;
-    let code = account.info.code.as_ref()?;
-
-    // Skip if code is empty
-    if code.is_empty() {
-        return None;
-    }
-
-    // Try to find the artifact by deployed code
-    let code_bytes = code.original_bytes();
-    if let Some((artifact_id, _)) = artifacts.find_by_deployed_code_exact(&code_bytes) {
-        return Some(artifact_id.identifier());
-    }
+/// EIP-1967 implementation storage slot
+const EIP1967_IMPL_SLOT: &str = "360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc";
 
-    // Fallback to fuzzy matching if exact match fails
-    if let Some((artifact_id, _)) = artifacts.find_by_deployed_code(&code_bytes) {
-        return Some(artifact_id.identifier());
-    }
-
-    None
-}
+/// EIP-1822 UUPS implementation storage slot: keccak256("PROXIABLE")
+const EIP1822_PROXIABLE_SLOT: &str =
+    "c5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7";
 
 /// Helper function to get the contract data from the deployed code at an address.
 fn get_contract_data<'a>(
@@ -1518,6 +1497,22 @@ fn get_contract_data<'a>(
 
     // Try to find the artifact by deployed code
     let code_bytes = code.original_bytes();
+    // First check for proxy patterns
+    let hex_str = hex::encode(&code_bytes);
+    let find_by_suffix =
+        |suffix: &str| artifacts.iter().find(|(a, _)| a.identifier().ends_with(suffix));
+    // Simple proxy detection based on storage slot patterns
+    if hex_str.contains(EIP1967_IMPL_SLOT)
+        && let Some(result) = find_by_suffix(":TransparentUpgradeableProxy")
+    {
+        return Some(result);
+    } else if hex_str.contains(EIP1822_PROXIABLE_SLOT)
+        && let Some(result) = find_by_suffix(":UUPSUpgradeable")
+    {
+        return Some(result);
+    }
+
+    // Try exact match
     if let Some(result) = artifacts.find_by_deployed_code_exact(&code_bytes) {
         return Some(result);
     }