feat: vm.dumpState

crates/cheatcodes/spec/src/vm.rs

@@ -227,6 +227,10 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Safe)]
     function addr(uint256 privateKey) external pure returns (address keyAddr);
 
+    /// Dump a genesis JSON file's `allocs` to disk.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function dumpState(string calldata pathToStateJson) external;
+
     /// Gets the nonce of an account.
     #[cheatcode(group = Evm, safety = Safe)]
     function getNonce(address account) external view returns (uint64 nonce);

crates/cheatcodes/src/evm.rs

@@ -3,10 +3,19 @@
 use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*};
 use alloy_primitives::{Address, Bytes, U256};
 use alloy_sol_types::SolValue;
-use ethers_core::utils::{Genesis, GenesisAccount};
+use ethers_core::{
+    types::H256,
+    utils::{Genesis, GenesisAccount},
+};
 use ethers_signers::Signer;
-use foundry_common::{fs::read_json_file, types::ToAlloy};
-use foundry_evm_core::backend::{DatabaseExt, RevertSnapshotAction};
+use foundry_common::{
+    fs::{read_json_file, write_json_file},
+    types::{ToAlloy, ToEthers},
+};
+use foundry_evm_core::{
+    backend::{DatabaseExt, RevertSnapshotAction},
+    constants::{CALLER, CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, TEST_CONTRACT_ADDRESS},
+};
 use revm::{
     primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY},
     EVMData,
@@ -86,6 +95,59 @@ impl Cheatcode for loadAllocsCall {
     }
 }
 
+impl Cheatcode for dumpStateCall {
+    fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { pathToStateJson } = self;
+        let path = Path::new(pathToStateJson);
+
+        // Do not include system accounts in the dump.
+        let skip = |key: &Address| {
+            key == &CHEATCODE_ADDRESS ||
+                key == &CALLER ||
+                key == &HARDHAT_CONSOLE_ADDRESS ||
+                key == &TEST_CONTRACT_ADDRESS ||
+                key == &ccx.caller ||
+                key == &ccx.state.config.evm_opts.sender
+        };
+
+        let accounts = ccx
+            .data
+            .journaled_state
+            .state()
+            .iter()
+            .filter(|(key, _)| !skip(key))
+            .collect::<HashMap<_, _>>();
+
+        let mut alloc: HashMap<Address, GenesisAccount> = HashMap::with_capacity(accounts.len());
+
+        accounts.into_iter().for_each(|(key, val)| {
+            alloc.insert(
+                *key,
+                GenesisAccount {
+                    nonce: Some(val.info.nonce),
+                    balance: val.info.balance.to_ethers(),
+                    code: Some(
+                        val.info.code.clone().unwrap_or_default().original_bytes().to_ethers(),
+                    ),
+                    storage: Some(
+                        val.storage
+                            .iter()
+                            .map(|(k, v)| {
+                                let key = k.to_be_bytes::<32>();
+                                let val = v.present_value().to_be_bytes::<32>();
+                                (H256::from_slice(&key), H256::from_slice(&val))
+                            })
+                            .collect(),
+                    ),
+                },
+            );
+        });
+
+        write_json_file(path, &alloc)?;
+        Ok(Default::default())
+    }
+}
+
 impl Cheatcode for sign_0Call {
     fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { privateKey, digest } = self;

testdata/cheats/dumpState.t.sol

@@ -0,0 +1,123 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity 0.8.18;
+
+import "ds-test/test.sol";
+import "./Vm.sol";
+
+contract SimpleContract {
+    constructor() {
+        assembly {
+            sstore(1, 2)
+        }
+    }
+}
+
+contract DumpStateTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    function testDumpStateCheatAccount() public {
+        // Path to temporary file that is deleted after the test
+        string memory path = string.concat(vm.projectRoot(), "/fixtures/Json/test_dump_state_cheat.json");
+
+        // Define some values to set in the state using cheatcodes
+        address target = address(1001);
+        bytes memory bytecode = hex"11223344";
+        uint256 balance = 1.2 ether;
+        uint64 nonce = 45;
+
+        vm.etch(target, bytecode);
+        vm.deal(target, balance);
+        vm.setNonce(target, nonce);
+        vm.store(target, bytes32(uint256(0x20)), bytes32(uint256(0x40)));
+        vm.store(target, bytes32(uint256(0x40)), bytes32(uint256(0x60)));
+
+        // Write the state to disk
+        vm.dumpState(path);
+
+        string memory json = vm.readFile(path);
+        string[] memory keys = vm.parseJsonKeys(json, "");
+        assertEq(keys.length, 1);
+
+        string memory key = keys[0];
+        assertEq(nonce, vm.parseJsonUint(json, string.concat(".", key, ".nonce")));
+        assertEq(balance, vm.parseJsonUint(json, string.concat(".", key, ".balance")));
+        assertEq(bytecode, vm.parseJsonBytes(json, string.concat(".", key, ".code")));
+
+        string[] memory slots = vm.parseJsonKeys(json, string.concat(".", key, ".storage"));
+        assertEq(slots.length, 2);
+
+        assertEq(
+            bytes32(uint256(0x40)),
+            vm.parseJsonBytes32(json, string.concat(".", key, ".storage.", vm.toString(bytes32(uint256(0x20)))))
+        );
+        assertEq(
+            bytes32(uint256(0x60)),
+            vm.parseJsonBytes32(json, string.concat(".", key, ".storage.", vm.toString(bytes32(uint256(0x40)))))
+        );
+
+        vm.removeFile(path);
+    }
+
+    function testDumpStateMultipleAccounts() public {
+        string memory path = string.concat(vm.projectRoot(), "/fixtures/Json/test_dump_state_multiple_accounts.json");
+
+        vm.setNonce(address(0x100), 1);
+        vm.deal(address(0x200), 1 ether);
+        vm.store(address(0x300), bytes32(uint256(1)), bytes32(uint256(2)));
+        vm.etch(address(0x400), hex"af");
+
+        vm.dumpState(path);
+
+        string memory json = vm.readFile(path);
+        string[] memory keys = vm.parseJsonKeys(json, "");
+        assertEq(keys.length, 4);
+
+        assertEq(4, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x100)))).length);
+        assertEq(1, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x100)), ".nonce")));
+        assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x100)), ".balance")));
+        assertEq(hex"", vm.parseJsonBytes(json, string.concat(".", vm.toString(address(0x100)), ".code")));
+        assertEq(0, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x100)), ".storage")).length);
+
+        assertEq(4, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x200)))).length);
+        assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x200)), ".nonce")));
+        assertEq(1 ether, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x200)), ".balance")));
+        assertEq(hex"", vm.parseJsonBytes(json, string.concat(".", vm.toString(address(0x200)), ".code")));
+        assertEq(0, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x200)), ".storage")).length);
+
+        assertEq(4, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x300)))).length);
+        assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x300)), ".nonce")));
+        assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x300)), ".balance")));
+        assertEq(hex"", vm.parseJsonBytes(json, string.concat(".", vm.toString(address(0x300)), ".code")));
+        assertEq(1, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x300)), ".storage")).length);
+        assertEq(
+            2,
+            vm.parseJsonUint(
+                json, string.concat(".", vm.toString(address(0x300)), ".storage.", vm.toString(bytes32(uint256(1))))
+            )
+        );
+
+        assertEq(4, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x400)))).length);
+        assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x400)), ".nonce")));
+        assertEq(0, vm.parseJsonUint(json, string.concat(".", vm.toString(address(0x400)), ".balance")));
+        assertEq(hex"af", vm.parseJsonBytes(json, string.concat(".", vm.toString(address(0x400)), ".code")));
+        assertEq(0, vm.parseJsonKeys(json, string.concat(".", vm.toString(address(0x400)), ".storage")).length);
+
+        vm.removeFile(path);
+    }
+
+    function testDumpStateDeployment() public {
+        string memory path = string.concat(vm.projectRoot(), "/fixtures/Json/test_dump_state_deployment.json");
+
+        SimpleContract s = new SimpleContract();
+        vm.dumpState(path);
+
+        string memory json = vm.readFile(path);
+        string[] memory keys = vm.parseJsonKeys(json, "");
+        assertEq(keys.length, 1);
+        assertEq(address(s), vm.parseAddress(keys[0]));
+        assertEq(1, vm.parseJsonKeys(json, string.concat(".", keys[0], ".storage")).length);
+        assertEq(2, vm.parseJsonUint(json, string.concat(".", keys[0], ".storage.", vm.toString(bytes32(uint256(1))))));
+
+        vm.removeFile(path);
+    }
+}