Skip to content

Commit 1b3b9b4

Browse files
feat(common): identify bytes and string slots using storageLayout (#11476)
* feat(common): identify bytes and string slots using storageLayout * fmt * clippy * aggregate and decode bytes and strings * cleanup * nits * clippy * fmt * nit * cleanup * identify the slots using the length of data in base slot value * nit * fix * fix mapping identification * nit * fmt * nit * Update crates/common/src/slot_identifier.rs * Update crates/common/src/slot_identifier.rs * cleanup: slot_identifier * cleanup * nit --------- Co-authored-by: zerosnacks <[email protected]>
1 parent 124f134 commit 1b3b9b4

File tree

3 files changed

+585
-17
lines changed

3 files changed

+585
-17
lines changed

crates/cheatcodes/src/evm.rs

Lines changed: 48 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@ use crate::{
88
use alloy_consensus::TxEnvelope;
99
use alloy_genesis::{Genesis, GenesisAccount};
1010
use alloy_network::eip2718::EIP4844_TX_TYPE_ID;
11-
use alloy_primitives::{Address, B256, U256, hex, map::HashMap};
11+
use alloy_primitives::{
12+
Address, B256, U256, hex,
13+
map::{B256Map, HashMap},
14+
};
1215
use alloy_rlp::Decodable;
1316
use alloy_sol_types::SolValue;
1417
use foundry_common::{
@@ -1376,6 +1379,16 @@ fn get_recorded_state_diffs(ccx: &mut CheatsCtxt) -> BTreeMap<Address, AccountSt
13761379
}
13771380
}
13781381

1382+
// Collect all storage accesses for this account
1383+
let raw_changes_by_slot = account_access
1384+
.storageAccesses
1385+
.iter()
1386+
.filter_map(|access| {
1387+
(access.isWrite && !access.reverted)
1388+
.then_some((access.slot, (access.previousValue, access.newValue)))
1389+
})
1390+
.collect::<BTreeMap<_, _>>();
1391+
13791392
// Record account state diffs.
13801393
for storage_access in &account_access.storageAccesses {
13811394
if storage_access.isWrite && !storage_access.reverted {
@@ -1398,19 +1411,43 @@ fn get_recorded_state_diffs(ccx: &mut CheatsCtxt) -> BTreeMap<Address, AccountSt
13981411
.as_ref()
13991412
.and_then(|slots| slots.get(&storage_access.account));
14001413

1401-
let mut slot_info = layout.and_then(|layout| {
1414+
let slot_info = layout.and_then(|layout| {
14021415
let decoder = SlotIdentifier::new(layout.clone());
1403-
decoder.identify(&storage_access.slot, mapping_slots)
1416+
decoder
1417+
.identify(&storage_access.slot, mapping_slots)
1418+
.or_else(|| {
1419+
// Create a map of new values for bytes/string
1420+
// identification. These values are used to determine
1421+
// the length of the data which helps determine how many
1422+
// slots to search
1423+
let current_base_slot_values = raw_changes_by_slot
1424+
.iter()
1425+
.map(|(slot, (_, new_val))| (*slot, *new_val))
1426+
.collect::<B256Map<_>>();
1427+
decoder.identify_bytes_or_string(
1428+
&storage_access.slot,
1429+
&current_base_slot_values,
1430+
)
1431+
})
1432+
.map(|mut info| {
1433+
// Always decode values first
1434+
info.decode_values(
1435+
storage_access.previousValue,
1436+
storage_access.newValue,
1437+
);
1438+
1439+
// Then handle long bytes/strings if applicable
1440+
if info.is_bytes_or_string() {
1441+
info.decode_bytes_or_string_values(
1442+
&storage_access.slot,
1443+
&raw_changes_by_slot,
1444+
);
1445+
}
1446+
1447+
info
1448+
})
14041449
});
14051450

1406-
// Decode values if we have slot info
1407-
if let Some(ref mut info) = slot_info {
1408-
info.decode_values(
1409-
storage_access.previousValue,
1410-
storage_access.newValue,
1411-
);
1412-
}
1413-
14141451
slot_state_diff.insert(SlotStateDiff {
14151452
previous_value: storage_access.previousValue,
14161453
new_value: storage_access.newValue,

0 commit comments

Comments
 (0)