From 00efa0d5965269149f374ba142fb1c3c7edd6c94 Mon Sep 17 00:00:00 2001 From: grandizzy <38490174+grandizzy@users.noreply.github.com> Date: Fri, 6 Dec 2024 09:03:32 +0200 Subject: [PATCH] feat(cheatcodes): add `vm.getStateDiff` to get state diffs as string (#9435) * feat(cheatcodes): add vm.getStateDiff() to get state diffs as string * Nit arrow * Add json output * Better json format * Rename to original and dirty * Changes after review: split in 2 cheatcodes, rename to prev/newValues * Slots as hex strings, add balance diffs, cleanup * Record balance diffs only if changed. Add nonce diff placeholder * Backoff nonce placeholder --- Cargo.lock | 1 + crates/cheatcodes/Cargo.toml | 1 + crates/cheatcodes/assets/cheatcodes.json | 40 +++++ crates/cheatcodes/spec/src/vm.rs | 8 + crates/cheatcodes/src/evm.rs | 143 +++++++++++++++++- testdata/cheats/Vm.sol | 2 + .../cheats/RecordAccountAccesses.t.sol | 37 +++++ 7 files changed, 231 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 37369e34ebae..857f0a91536b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3732,6 +3732,7 @@ dependencies = [ "revm", "revm-inspectors", "semver 1.0.23", + "serde", "serde_json", "thiserror 2.0.3", "toml 0.8.19", diff --git a/crates/cheatcodes/Cargo.toml b/crates/cheatcodes/Cargo.toml index a06416160111..e358fdce737b 100644 --- a/crates/cheatcodes/Cargo.toml +++ b/crates/cheatcodes/Cargo.toml @@ -68,3 +68,4 @@ toml = { workspace = true, features = ["preserve_order"] } tracing.workspace = true walkdir.workspace = true proptest.workspace = true +serde.workspace = true diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index 54556043ac9d..82a7de2aa509 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -5791,6 +5791,46 @@ "status": "stable", "safety": "safe" }, + { + "func": { + "id": "getStateDiff", + "description": "Returns state diffs from current `vm.startStateDiffRecording` session.", + "declaration": "function getStateDiff() external view returns (string memory diff);", + "visibility": "external", + "mutability": "view", + "signature": "getStateDiff()", + "selector": "0x80df01cc", + "selectorBytes": [ + 128, + 223, + 1, + 204 + ] + }, + "group": "evm", + "status": "stable", + "safety": "safe" + }, + { + "func": { + "id": "getStateDiffJson", + "description": "Returns state diffs from current `vm.startStateDiffRecording` session, in json format.", + "declaration": "function getStateDiffJson() external view returns (string memory diff);", + "visibility": "external", + "mutability": "view", + "signature": "getStateDiffJson()", + "selector": "0xf54fe009", + "selectorBytes": [ + 245, + 79, + 224, + 9 + ] + }, + "group": "evm", + "status": "stable", + "safety": "safe" + }, { "func": { "id": "getWallets", diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index 6954fd1e526d..4bc8c9b0344e 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -384,6 +384,14 @@ interface Vm { #[cheatcode(group = Evm, safety = Safe)] function stopAndReturnStateDiff() external returns (AccountAccess[] memory accountAccesses); + /// Returns state diffs from current `vm.startStateDiffRecording` session. + #[cheatcode(group = Evm, safety = Safe)] + function getStateDiff() external view returns (string memory diff); + + /// Returns state diffs from current `vm.startStateDiffRecording` session, in json format. + #[cheatcode(group = Evm, safety = Safe)] + function getStateDiffJson() external view returns (string memory diff); + // -------- Recording Map Writes -------- /// Starts recording all map SSTOREs for later retrieval. diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index 723beb1cfc8a..6c1a4718512d 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -18,10 +18,15 @@ use foundry_evm_core::{ use foundry_evm_traces::StackSnapshotType; use rand::Rng; use revm::primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY}; -use std::{collections::BTreeMap, path::Path}; +use std::{ + collections::{btree_map::Entry, BTreeMap}, + fmt::Display, + path::Path, +}; mod record_debug_step; use record_debug_step::{convert_call_trace_to_debug_step, flatten_call_trace}; +use serde::Serialize; mod fork; pub(crate) mod mapping; @@ -76,6 +81,70 @@ pub struct DealRecord { pub new_balance: U256, } +/// Storage slot diff info. +#[derive(Serialize, Default)] +#[serde(rename_all = "camelCase")] +struct SlotStateDiff { + /// Initial storage value. + previous_value: B256, + /// Current storage value. + new_value: B256, +} + +/// Balance diff info. +#[derive(Serialize, Default)] +#[serde(rename_all = "camelCase")] +struct BalanceDiff { + /// Initial storage value. + previous_value: U256, + /// Current storage value. + new_value: U256, +} + +/// Account state diff info. +#[derive(Serialize, Default)] +#[serde(rename_all = "camelCase")] +struct AccountStateDiffs { + /// Address label, if any set. + label: Option, + /// Account balance changes. + balance_diff: Option, + /// State changes, per slot. + state_diff: BTreeMap, +} + +impl Display for AccountStateDiffs { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> eyre::Result<(), std::fmt::Error> { + // Print changed account. + if let Some(label) = &self.label { + writeln!(f, "label: {label}")?; + } + // Print balance diff if changed. + if let Some(balance_diff) = &self.balance_diff { + if balance_diff.previous_value != balance_diff.new_value { + writeln!( + f, + "- balance diff: {} → {}", + balance_diff.previous_value, balance_diff.new_value + )?; + } + } + // Print state diff if any. + if !&self.state_diff.is_empty() { + writeln!(f, "- state diff:")?; + for (slot, slot_changes) in &self.state_diff { + writeln!( + f, + "@ {slot}: {} → {}", + slot_changes.previous_value, slot_changes.new_value + )?; + } + } + + Ok(()) + } +} + impl Cheatcode for addrCall { fn apply(&self, _state: &mut Cheatcodes) -> Result { let Self { privateKey } = self; @@ -683,6 +752,25 @@ impl Cheatcode for stopAndReturnStateDiffCall { } } +impl Cheatcode for getStateDiffCall { + fn apply(&self, state: &mut Cheatcodes) -> Result { + let mut diffs = String::new(); + let state_diffs = get_recorded_state_diffs(state); + for (address, state_diffs) in state_diffs { + diffs.push_str(&format!("{address}\n")); + diffs.push_str(&format!("{state_diffs}\n")); + } + Ok(diffs.abi_encode()) + } +} + +impl Cheatcode for getStateDiffJsonCall { + fn apply(&self, state: &mut Cheatcodes) -> Result { + let state_diffs = get_recorded_state_diffs(state); + Ok(serde_json::to_string(&state_diffs)?.abi_encode()) + } +} + impl Cheatcode for broadcastRawTransactionCall { fn apply_full(&self, ccx: &mut CheatsCtxt, executor: &mut dyn CheatcodesExecutor) -> Result { let tx = TxEnvelope::decode(&mut self.data.as_ref()) @@ -1044,3 +1132,56 @@ fn genesis_account(account: &Account) -> GenesisAccount { private_key: None, } } + +/// Helper function to returns state diffs recorded for each changed account. +fn get_recorded_state_diffs(state: &mut Cheatcodes) -> BTreeMap { + let mut state_diffs: BTreeMap = BTreeMap::default(); + if let Some(records) = &state.recorded_account_diffs_stack { + records + .iter() + .flatten() + .filter(|account_access| { + !account_access.storageAccesses.is_empty() || + account_access.oldBalance != account_access.newBalance + }) + .for_each(|account_access| { + let account_diff = + state_diffs.entry(account_access.account).or_insert(AccountStateDiffs { + label: state.labels.get(&account_access.account).cloned(), + ..Default::default() + }); + + // Record account balance diffs. + if account_access.oldBalance != account_access.newBalance { + // Update balance diff. Do not overwrite the initial balance if already set. + if let Some(diff) = &mut account_diff.balance_diff { + diff.new_value = account_access.newBalance; + } else { + account_diff.balance_diff = Some(BalanceDiff { + previous_value: account_access.oldBalance, + new_value: account_access.newBalance, + }); + } + } + + // Record account state diffs. + for storage_access in &account_access.storageAccesses { + if storage_access.isWrite && !storage_access.reverted { + // 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) => { + slot_state_diff.insert(SlotStateDiff { + previous_value: storage_access.previousValue, + new_value: storage_access.newValue, + }); + } + Entry::Occupied(mut slot_state_diff) => { + slot_state_diff.get_mut().new_value = storage_access.newValue; + } + } + } + } + }); + } + state_diffs +} diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index b3746b6bc319..bdbb68e372e8 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -283,6 +283,8 @@ interface Vm { function getNonce(address account) external view returns (uint64 nonce); function getNonce(Wallet calldata wallet) external returns (uint64 nonce); function getRecordedLogs() external returns (Log[] memory logs); + function getStateDiff() external view returns (string memory diff); + function getStateDiffJson() external view returns (string memory diff); function getWallets() external returns (address[] memory wallets); function indexOf(string calldata input, string calldata key) external pure returns (uint256); function isContext(ForgeContext context) external view returns (bool result); diff --git a/testdata/default/cheats/RecordAccountAccesses.t.sol b/testdata/default/cheats/RecordAccountAccesses.t.sol index 98b5843b2a7f..8de7bcdc5bd3 100644 --- a/testdata/default/cheats/RecordAccountAccesses.t.sol +++ b/testdata/default/cheats/RecordAccountAccesses.t.sol @@ -3,6 +3,7 @@ pragma solidity ^0.8.18; import "ds-test/test.sol"; import "cheats/Vm.sol"; +import "../logs/console.sol"; /// @notice Helper contract with a construction that makes a call to itself then /// optionally reverts if zero-length data is passed @@ -261,6 +262,16 @@ contract RecordAccountAccessesTest is DSTest { two.write(bytes32(uint256(5678)), bytes32(uint256(123469))); two.write(bytes32(uint256(5678)), bytes32(uint256(1234))); + string memory diffs = cheats.getStateDiff(); + assertEq( + "0x5991A2dF15A8F6A256D3Ec51E99254Cd3fb576A9\n- state diff:\n@ 0x00000000000000000000000000000000000000000000000000000000000004d3: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x000000000000000000000000000000000000000000000000000000000000162e\n\n0xc7183455a4C133Ae270771860664b6B7ec320bB1\n- state diff:\n@ 0x000000000000000000000000000000000000000000000000000000000000162e: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x00000000000000000000000000000000000000000000000000000000000004d2\n\n", + diffs + ); + string memory diffsJson = cheats.getStateDiffJson(); + assertEq( + "{\"0x5991a2df15a8f6a256d3ec51e99254cd3fb576a9\":{\"label\":null,\"balanceDiff\":null,\"stateDiff\":{\"0x00000000000000000000000000000000000000000000000000000000000004d3\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x000000000000000000000000000000000000000000000000000000000000162e\"}}},\"0xc7183455a4c133ae270771860664b6b7ec320bb1\":{\"label\":null,\"balanceDiff\":null,\"stateDiff\":{\"0x000000000000000000000000000000000000000000000000000000000000162e\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x00000000000000000000000000000000000000000000000000000000000004d2\"}}}}", + diffsJson + ); Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff()); assertEq(called.length, 4, "incorrect length"); @@ -332,6 +343,15 @@ contract RecordAccountAccessesTest is DSTest { // contract calls to self in constructor SelfCaller caller = new SelfCaller{value: 2 ether}("hello2 world2"); + assertEq( + "0x000000000000000000000000000000000000162e\n- balance diff: 0 \xE2\x86\x92 1000000000000000000\n\n0x1d1499e622D69689cdf9004d05Ec547d650Ff211\n- balance diff: 0 \xE2\x86\x92 2000000000000000000\n\n", + cheats.getStateDiff() + ); + assertEq( + "{\"0x000000000000000000000000000000000000162e\":{\"label\":null,\"balanceDiff\":{\"previousValue\":\"0x0\",\"newValue\":\"0xde0b6b3a7640000\"},\"stateDiff\":{}},\"0x1d1499e622d69689cdf9004d05ec547d650ff211\":{\"label\":null,\"balanceDiff\":{\"previousValue\":\"0x0\",\"newValue\":\"0x1bc16d674ec80000\"},\"stateDiff\":{}}}", + cheats.getStateDiffJson() + ); + Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff()); assertEq(called.length, 6); assertEq( @@ -451,6 +471,14 @@ contract RecordAccountAccessesTest is DSTest { uint256 initBalance = address(this).balance; cheats.startStateDiffRecording(); try this.revertingCall{value: 1 ether}(address(1234), "") {} catch {} + assertEq( + "0x00000000000000000000000000000000000004d2\n- balance diff: 0 \xE2\x86\x92 100000000000000000\n\n", + cheats.getStateDiff() + ); + assertEq( + "{\"0x00000000000000000000000000000000000004d2\":{\"label\":null,\"balanceDiff\":{\"previousValue\":\"0x0\",\"newValue\":\"0x16345785d8a0000\"},\"stateDiff\":{}}}", + cheats.getStateDiffJson() + ); Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff()); assertEq(called.length, 2); assertEq( @@ -768,6 +796,15 @@ contract RecordAccountAccessesTest is DSTest { function testNestedStorage() public { cheats.startStateDiffRecording(); nestedStorer.run(); + cheats.label(address(nestedStorer), "NestedStorer"); + assertEq( + "0x2e234DAe75C793f67A35089C9d99245E1C58470b\nlabel: NestedStorer\n- state diff:\n@ 0x4566fa0cd03218c55bba914d793f5e6b9113172c1f684bb5f464c08c867e8977: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x0000000000000000000000000000000000000000000000000000000000000001\n@ 0xbf57896b60daefa2c41de2feffecfc11debd98ea8c913a5170f60e53959ac00a: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x0000000000000000000000000000000000000000000000000000000000000001\n@ 0xc664893a982d78bbeab379feef216ff517b7ea73626b280723be1ace370364cd: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x0000000000000000000000000000000000000000000000000000000000000001\n@ 0xdc5330afa9872081253545dca3f448752688ff1b098b38c1abe4c4cdff4b0b0e: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x0000000000000000000000000000000000000000000000000000000000000001\n\n", + cheats.getStateDiff() + ); + assertEq( + "{\"0x2e234dae75c793f67a35089c9d99245e1c58470b\":{\"label\":\"NestedStorer\",\"balanceDiff\":null,\"stateDiff\":{\"0x4566fa0cd03218c55bba914d793f5e6b9113172c1f684bb5f464c08c867e8977\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x0000000000000000000000000000000000000000000000000000000000000001\"},\"0xbf57896b60daefa2c41de2feffecfc11debd98ea8c913a5170f60e53959ac00a\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x0000000000000000000000000000000000000000000000000000000000000001\"},\"0xc664893a982d78bbeab379feef216ff517b7ea73626b280723be1ace370364cd\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x0000000000000000000000000000000000000000000000000000000000000001\"},\"0xdc5330afa9872081253545dca3f448752688ff1b098b38c1abe4c4cdff4b0b0e\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x0000000000000000000000000000000000000000000000000000000000000001\"}}}}", + cheats.getStateDiffJson() + ); Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff()); assertEq(called.length, 3, "incorrect account access length");