fix(cheatcode): use storage access address instead account access (#9646)

* fix(cheatcode): use storage access address instead account access

* Update crates/cheatcodes/src/evm.rs

Co-authored-by: DaniPopes <[email protected]>

* Fix fmt

---------

Co-authored-by: DaniPopes <[email protected]>

Author: grandizzy

crates/cheatcodes/src/evm.rs

@@ -1149,11 +1149,12 @@ fn get_recorded_state_diffs(state: &mut Cheatcodes) -> BTreeMap<Address, Account
                     account_access.oldBalance != account_access.newBalance
             })
             .for_each(|account_access| {
-                let account_diff =
-                    state_diffs.entry(account_access.account).or_insert(AccountStateDiffs {
+                let account_diff = state_diffs.entry(account_access.account).or_insert_with(|| {
+                    AccountStateDiffs {
                         label: state.labels.get(&account_access.account).cloned(),
                         ..Default::default()
-                    });
+                    }
+                });
 
                 // Record account balance diffs.
                 if account_access.oldBalance != account_access.newBalance {
@@ -1171,6 +1172,12 @@ fn get_recorded_state_diffs(state: &mut Cheatcodes) -> BTreeMap<Address, Account
                 // Record account state diffs.
                 for storage_access in &account_access.storageAccesses {
                     if storage_access.isWrite && !storage_access.reverted {
+                        let account_diff = state_diffs
+                            .entry(storage_access.account)
+                            .or_insert_with(|| AccountStateDiffs {
+                                label: state.labels.get(&storage_access.account).cloned(),
+                                ..Default::default()
+                            });
                         // Update state diff. Do not overwrite the initial value if already set.
                         match account_diff.state_diff.entry(storage_access.slot) {
                             Entry::Vacant(slot_state_diff) => {

crates/forge/tests/it/repros.rs

@@ -392,3 +392,6 @@ test_repro!(8639);
 
 // https://github.com/foundry-rs/foundry/issues/8566
 test_repro!(8566);
+
+// https://github.com/foundry-rs/foundry/issues/9643
+test_repro!(9643);

testdata/default/repros/Issue9643.t.sol

@@ -0,0 +1,49 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity ^0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+
+contract Mock {
+    uint256 private counter;
+
+    function setCounter(uint256 _counter) external {
+        counter = _counter;
+    }
+}
+
+contract DelegateProxy {
+    address internal implementation;
+
+    constructor(address mock) {
+        implementation = mock;
+    }
+
+    fallback() external payable {
+        address addr = implementation;
+
+        assembly {
+            calldatacopy(0, 0, calldatasize())
+            let result := delegatecall(gas(), addr, 0, calldatasize(), 0, 0)
+            returndatacopy(0, 0, returndatasize())
+            switch result
+            case 0 { revert(0, returndatasize()) }
+            default { return(0, returndatasize()) }
+        }
+    }
+}
+
+contract Issue9643Test is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    function test_storage_json_diff() public {
+        vm.startStateDiffRecording();
+        Mock proxied = Mock(address(new DelegateProxy(address(new Mock()))));
+        proxied.setCounter(42);
+        string memory rawDiff = vm.getStateDiffJson();
+        assertEq(
+            "{\"0x2e234dae75c793f67a35089c9d99245e1c58470b\":{\"label\":null,\"balanceDiff\":null,\"stateDiff\":{\"0x0000000000000000000000000000000000000000000000000000000000000000\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x000000000000000000000000000000000000000000000000000000000000002a\"}}},\"0x5615deb798bb3e4dfa0139dfa1b3d433cc23b72f\":{\"label\":null,\"balanceDiff\":null,\"stateDiff\":{}}}",
+            rawDiff
+        );
+    }
+}