Skip to content

Commit

Permalink
Slots as hex strings, add balance diffs, cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@grandizzy
grandizzy committed Dec 4, 2024
1 parent fa63595 commit a188fd2
Show file tree
Hide file tree
Showing 2 changed files with 101 additions and 63 deletions.
152 changes: 95 additions & 57 deletions crates/cheatcodes/src/evm.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,11 @@ 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};
Expand Down Expand Up @@ -77,27 +81,70 @@ pub struct DealRecord {
pub new_balance: U256,
}

/// Recorded state diffs, for each changed address.
type StateDiffs = BTreeMap<Address, AccountStateDiffs>;
/// 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)]
#[derive(Serialize, Default)]
#[serde(rename_all = "camelCase")]
struct AccountStateDiffs {
/// Address label, if any set.
#[serde(skip_serializing_if = "Option::is_none")]
label: Option<String>,
/// Account balance changes.
#[serde(skip_serializing_if = "Option::is_none")]
balance_diff: Option<BalanceDiff>,
/// State changes, per slot.
changes: BTreeMap<String, SlotStateDiff>,
state_diff: BTreeMap<B256, SlotStateDiff>,
}

/// Storage slot diff info.
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct SlotStateDiff {
/// Initial storage value.
previous_value: B256,
/// Current storage value.
new_value: B256,
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 {
Expand Down Expand Up @@ -712,20 +759,9 @@ impl Cheatcode for getStateDiffCall {
let mut diffs = String::new();
let state_diffs = get_recorded_state_diffs(state);
for (address, state_diffs) in state_diffs {
// Print changed account.
if let Some(label) = state_diffs.label {
diffs.push_str(&format!("{address} ({label}):\n"));
} else {
diffs.push_str(&format!("{address}:\n"));
}
for (slot, slot_changes) in state_diffs.changes {
diffs.push_str(&format!(
"@ {slot}: {} → {}\n",
slot_changes.previous_value, slot_changes.new_value
));
}
diffs.push_str(&format!("{address}\n"));
diffs.push_str(&format!("{state_diffs}\n"));
}

Ok(diffs.abi_encode())
}
}
Expand Down Expand Up @@ -1100,47 +1136,49 @@ fn genesis_account(account: &Account) -> GenesisAccount {
}

/// Helper function to returns state diffs recorded for each changed account.
fn get_recorded_state_diffs(state: &mut Cheatcodes) -> StateDiffs {
let mut state_diffs: StateDiffs = BTreeMap::default();
fn get_recorded_state_diffs(state: &mut Cheatcodes) -> BTreeMap<Address, AccountStateDiffs> {
let mut state_diffs: BTreeMap<Address, AccountStateDiffs> = BTreeMap::default();
if let Some(records) = &state.recorded_account_diffs_stack {
let mut recorded_changes: BTreeMap<Address, BTreeMap<U256, Vec<(B256, B256)>>> =
BTreeMap::default();
records
.iter()
.flatten()
.filter(|account_access| !account_access.storageAccesses.is_empty())
.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()
});
// 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,
});
}

for storage_access in &account_access.storageAccesses {
// Retain only records with storage accesses writes that didn't revert.
if storage_access.isWrite && !storage_access.reverted {
recorded_changes
.entry(storage_access.account)
.or_default()
.entry(storage_access.slot.into())
.or_default()
.push((storage_access.previousValue, storage_access.newValue));
// 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;
}
}
}
}
});

for (address, recorded_changes) in recorded_changes {
let mut changes = BTreeMap::default();
for (slot, recorded_slot_changes) in recorded_changes {
// For each slot we retain the initial value from first state change and the new
// value from last state change.
changes.insert(
slot.to_string(),
SlotStateDiff {
previous_value: recorded_slot_changes.first().unwrap().0,
new_value: recorded_slot_changes.last().unwrap().1,
},
);
}
state_diffs.insert(
address,
AccountStateDiffs { label: state.labels.get(&address).cloned(), changes },
);
}
}
state_diffs
}
12 changes: 6 additions & 6 deletions testdata/default/cheats/RecordAccountAccesses.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -264,12 +264,12 @@ contract RecordAccountAccessesTest is DSTest {

string memory diffs = cheats.getStateDiff();
assertEq(
"0x5991A2dF15A8F6A256D3Ec51E99254Cd3fb576A9:\n@ 1235: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x000000000000000000000000000000000000000000000000000000000000162e\n0xc7183455a4C133Ae270771860664b6B7ec320bB1:\n@ 5678: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x00000000000000000000000000000000000000000000000000000000000004d2\n",
"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\":{\"changes\":{\"1235\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x000000000000000000000000000000000000000000000000000000000000162e\"}}},\"0xc7183455a4c133ae270771860664b6b7ec320bb1\":{\"changes\":{\"5678\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x00000000000000000000000000000000000000000000000000000000000004d2\"}}}}",
"{\"0x5991a2df15a8f6a256d3ec51e99254cd3fb576a9\":{\"balanceDiff\":{\"previousValue\":\"0x0\",\"newValue\":\"0x0\"},\"stateDiff\":{\"0x00000000000000000000000000000000000000000000000000000000000004d3\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x000000000000000000000000000000000000000000000000000000000000162e\"}}},\"0xc7183455a4c133ae270771860664b6b7ec320bb1\":{\"balanceDiff\":{\"previousValue\":\"0x0\",\"newValue\":\"0x0\"},\"stateDiff\":{\"0x000000000000000000000000000000000000000000000000000000000000162e\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x00000000000000000000000000000000000000000000000000000000000004d2\"}}}}",
diffsJson
);
Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
Expand Down Expand Up @@ -343,7 +343,7 @@ contract RecordAccountAccessesTest is DSTest {
// contract calls to self in constructor
SelfCaller caller = new SelfCaller{value: 2 ether}("hello2 world2");

assertEq("", cheats.getStateDiff());
assertEq("0x000000000000000000000000000000000000162e\n- balance diff: 0 \xE2\x86\x92 1000000000000000000\n\n0x1d1499e622D69689cdf9004d05Ec547d650Ff211\n- balance diff: 0 \xE2\x86\x92 2000000000000000000\n\n", cheats.getStateDiff());
Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
assertEq(called.length, 6);
assertEq(
Expand Down Expand Up @@ -463,7 +463,7 @@ contract RecordAccountAccessesTest is DSTest {
uint256 initBalance = address(this).balance;
cheats.startStateDiffRecording();
try this.revertingCall{value: 1 ether}(address(1234), "") {} catch {}
assertEq("", cheats.getStateDiff());
assertEq("0x00000000000000000000000000000000000004d2\n- balance diff: 0 \xE2\x86\x92 100000000000000000\n\n", cheats.getStateDiff());
Vm.AccountAccess[] memory called = filterExtcodesizeForLegacyTests(cheats.stopAndReturnStateDiff());
assertEq(called.length, 2);
assertEq(
Expand Down Expand Up @@ -783,11 +783,11 @@ contract RecordAccountAccessesTest is DSTest {
nestedStorer.run();
cheats.label(address(nestedStorer), "NestedStorer");
assertEq(
"0x2e234DAe75C793f67A35089C9d99245E1C58470b (NestedStorer):\n@ 31391530734884398925509096751136955997235046655136458338700630915422204365175: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x0000000000000000000000000000000000000000000000000000000000000001\n@ 86546418208203448386783321347074308435724792809315873744194221534962779865098: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x0000000000000000000000000000000000000000000000000000000000000001\n@ 89735575844917174604881245405098157398514761457822262993733937076486162048205: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x0000000000000000000000000000000000000000000000000000000000000001\n@ 99655811014363889343382125167956395016210879868288374279890486979400290732814: 0x0000000000000000000000000000000000000000000000000000000000000000 \xE2\x86\x92 0x0000000000000000000000000000000000000000000000000000000000000001\n",
"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\",\"changes\":{\"31391530734884398925509096751136955997235046655136458338700630915422204365175\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x0000000000000000000000000000000000000000000000000000000000000001\"},\"86546418208203448386783321347074308435724792809315873744194221534962779865098\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x0000000000000000000000000000000000000000000000000000000000000001\"},\"89735575844917174604881245405098157398514761457822262993733937076486162048205\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x0000000000000000000000000000000000000000000000000000000000000001\"},\"99655811014363889343382125167956395016210879868288374279890486979400290732814\":{\"previousValue\":\"0x0000000000000000000000000000000000000000000000000000000000000000\",\"newValue\":\"0x0000000000000000000000000000000000000000000000000000000000000001\"}}}}",
"{\"0x2e234dae75c793f67a35089c9d99245e1c58470b\":{\"label\":\"NestedStorer\",\"balanceDiff\":{\"previousValue\":\"0x0\",\"newValue\":\"0x0\"},\"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());
Expand Down

0 comments on commit a188fd2

Please sign in to comment.