foundry_cheatcodes/
evm.rs

1//! Implementations of [`Evm`](spec::Group::Evm) cheatcodes.
2
3use crate::{
4    BroadcastableTransaction, Cheatcode, Cheatcodes, CheatcodesExecutor, CheatsCtxt, Error, Result,
5    Vm::*,
6    inspector::{Ecx, RecordDebugStepInfo},
7};
8use alloy_consensus::TxEnvelope;
9use alloy_genesis::{Genesis, GenesisAccount};
10use alloy_network::eip2718::EIP4844_TX_TYPE_ID;
11use alloy_primitives::{Address, B256, U256, hex, map::HashMap};
12use alloy_rlp::Decodable;
13use alloy_sol_types::SolValue;
14use foundry_common::{
15    fs::{read_json_file, write_json_file},
16    slot_identifier::{SlotIdentifier, SlotInfo},
17};
18use foundry_evm_core::{
19    ContextExt,
20    backend::{DatabaseExt, RevertStateSnapshotAction},
21    constants::{CALLER, CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, TEST_CONTRACT_ADDRESS},
22    utils::get_blob_base_fee_update_fraction_by_spec_id,
23};
24use foundry_evm_traces::StackSnapshotType;
25use itertools::Itertools;
26use rand::Rng;
27use revm::{
28    bytecode::Bytecode,
29    context::{Block, JournalTr},
30    primitives::{KECCAK_EMPTY, hardfork::SpecId},
31    state::Account,
32};
33use std::{
34    collections::{BTreeMap, HashSet, btree_map::Entry},
35    fmt::Display,
36    path::Path,
37};
38
39mod record_debug_step;
40use foundry_common::fmt::format_token_raw;
41use record_debug_step::{convert_call_trace_to_debug_step, flatten_call_trace};
42use serde::Serialize;
43
44mod fork;
45pub(crate) mod mapping;
46pub(crate) mod mock;
47pub(crate) mod prank;
48
49/// Records storage slots reads and writes.
50#[derive(Clone, Debug, Default)]
51pub struct RecordAccess {
52    /// Storage slots reads.
53    pub reads: HashMap<Address, Vec<U256>>,
54    /// Storage slots writes.
55    pub writes: HashMap<Address, Vec<U256>>,
56}
57
58impl RecordAccess {
59    /// Records a read access to a storage slot.
60    pub fn record_read(&mut self, target: Address, slot: U256) {
61        self.reads.entry(target).or_default().push(slot);
62    }
63
64    /// Records a write access to a storage slot.
65    ///
66    /// This also records a read internally as `SSTORE` does an implicit `SLOAD`.
67    pub fn record_write(&mut self, target: Address, slot: U256) {
68        self.record_read(target, slot);
69        self.writes.entry(target).or_default().push(slot);
70    }
71
72    /// Clears the recorded reads and writes.
73    pub fn clear(&mut self) {
74        // Also frees memory.
75        *self = Default::default();
76    }
77}
78
79/// Records the `snapshotGas*` cheatcodes.
80#[derive(Clone, Debug)]
81pub struct GasRecord {
82    /// The group name of the gas snapshot.
83    pub group: String,
84    /// The name of the gas snapshot.
85    pub name: String,
86    /// The total gas used in the gas snapshot.
87    pub gas_used: u64,
88    /// Depth at which the gas snapshot was taken.
89    pub depth: usize,
90}
91
92/// Records `deal` cheatcodes
93#[derive(Clone, Debug)]
94pub struct DealRecord {
95    /// Target of the deal.
96    pub address: Address,
97    /// The balance of the address before deal was applied
98    pub old_balance: U256,
99    /// Balance after deal was applied
100    pub new_balance: U256,
101}
102
103/// Storage slot diff info.
104#[derive(Serialize, Default)]
105#[serde(rename_all = "camelCase")]
106struct SlotStateDiff {
107    /// Initial storage value.
108    previous_value: B256,
109    /// Current storage value.
110    new_value: B256,
111    /// Storage layout metadata (variable name, type, offset).
112    /// Only present when contract has storage layout output.
113    /// This includes decoded values when available.
114    #[serde(skip_serializing_if = "Option::is_none", flatten)]
115    slot_info: Option<SlotInfo>,
116}
117
118/// Balance diff info.
119#[derive(Serialize, Default)]
120#[serde(rename_all = "camelCase")]
121struct BalanceDiff {
122    /// Initial storage value.
123    previous_value: U256,
124    /// Current storage value.
125    new_value: U256,
126}
127
128/// Nonce diff info.
129#[derive(Serialize, Default)]
130#[serde(rename_all = "camelCase")]
131struct NonceDiff {
132    /// Initial nonce value.
133    previous_value: u64,
134    /// Current nonce value.
135    new_value: u64,
136}
137
138/// Account state diff info.
139#[derive(Serialize, Default)]
140#[serde(rename_all = "camelCase")]
141struct AccountStateDiffs {
142    /// Address label, if any set.
143    label: Option<String>,
144    /// Contract identifier from artifact. e.g "src/Counter.sol:Counter"
145    contract: Option<String>,
146    /// Account balance changes.
147    balance_diff: Option<BalanceDiff>,
148    /// Account nonce changes.
149    nonce_diff: Option<NonceDiff>,
150    /// State changes, per slot.
151    state_diff: BTreeMap<B256, SlotStateDiff>,
152}
153
154impl Display for AccountStateDiffs {
155    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> eyre::Result<(), std::fmt::Error> {
156        // Print changed account.
157        if let Some(label) = &self.label {
158            writeln!(f, "label: {label}")?;
159        }
160        if let Some(contract) = &self.contract {
161            writeln!(f, "contract: {contract}")?;
162        }
163        // Print balance diff if changed.
164        if let Some(balance_diff) = &self.balance_diff
165            && balance_diff.previous_value != balance_diff.new_value
166        {
167            writeln!(
168                f,
169                "- balance diff: {} → {}",
170                balance_diff.previous_value, balance_diff.new_value
171            )?;
172        }
173        // Print nonce diff if changed.
174        if let Some(nonce_diff) = &self.nonce_diff
175            && nonce_diff.previous_value != nonce_diff.new_value
176        {
177            writeln!(f, "- nonce diff: {} → {}", nonce_diff.previous_value, nonce_diff.new_value)?;
178        }
179        // Print state diff if any.
180        if !&self.state_diff.is_empty() {
181            writeln!(f, "- state diff:")?;
182            for (slot, slot_changes) in &self.state_diff {
183                match &slot_changes.slot_info {
184                    Some(slot_info) => {
185                        if slot_info.decoded.is_some() {
186                            // Have slot info with decoded values - show decoded values
187                            let decoded = slot_info.decoded.as_ref().unwrap();
188                            writeln!(
189                                f,
190                                "@ {slot} ({}, {}): {} → {}",
191                                slot_info.label,
192                                slot_info.slot_type.dyn_sol_type,
193                                format_token_raw(&decoded.previous_value),
194                                format_token_raw(&decoded.new_value)
195                            )?;
196                        } else {
197                            // Have slot info but no decoded values - show raw hex values
198                            writeln!(
199                                f,
200                                "@ {slot} ({}, {}): {} → {}",
201                                slot_info.label,
202                                slot_info.slot_type.dyn_sol_type,
203                                slot_changes.previous_value,
204                                slot_changes.new_value
205                            )?;
206                        }
207                    }
208                    None => {
209                        // No slot info - show raw hex values
210                        writeln!(
211                            f,
212                            "@ {slot}: {} → {}",
213                            slot_changes.previous_value, slot_changes.new_value
214                        )?;
215                    }
216                }
217            }
218        }
219
220        Ok(())
221    }
222}
223
224impl Cheatcode for addrCall {
225    fn apply(&self, _state: &mut Cheatcodes) -> Result {
226        let Self { privateKey } = self;
227        let wallet = super::crypto::parse_wallet(privateKey)?;
228        Ok(wallet.address().abi_encode())
229    }
230}
231
232impl Cheatcode for getNonce_0Call {
233    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
234        let Self { account } = self;
235        get_nonce(ccx, account)
236    }
237}
238
239impl Cheatcode for getNonce_1Call {
240    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
241        let Self { wallet } = self;
242        get_nonce(ccx, &wallet.addr)
243    }
244}
245
246impl Cheatcode for loadCall {
247    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
248        let Self { target, slot } = *self;
249        ccx.ensure_not_precompile(&target)?;
250        ccx.ecx.journaled_state.load_account(target)?;
251        let mut val = ccx.ecx.journaled_state.sload(target, slot.into())?;
252
253        if val.is_cold && val.data.is_zero() {
254            if ccx.state.has_arbitrary_storage(&target) {
255                // If storage slot is untouched and load from a target with arbitrary storage,
256                // then set random value for current slot.
257                let rand_value = ccx.state.rng().random();
258                ccx.state.arbitrary_storage.as_mut().unwrap().save(
259                    ccx.ecx,
260                    target,
261                    slot.into(),
262                    rand_value,
263                );
264                val.data = rand_value;
265            } else if ccx.state.is_arbitrary_storage_copy(&target) {
266                // If storage slot is untouched and load from a target that copies storage from
267                // a source address with arbitrary storage, then copy existing arbitrary value.
268                // If no arbitrary value generated yet, then the random one is saved and set.
269                let rand_value = ccx.state.rng().random();
270                val.data = ccx.state.arbitrary_storage.as_mut().unwrap().copy(
271                    ccx.ecx,
272                    target,
273                    slot.into(),
274                    rand_value,
275                );
276            }
277        }
278
279        Ok(val.abi_encode())
280    }
281}
282
283impl Cheatcode for loadAllocsCall {
284    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
285        let Self { pathToAllocsJson } = self;
286
287        let path = Path::new(pathToAllocsJson);
288        ensure!(path.exists(), "allocs file does not exist: {pathToAllocsJson}");
289
290        // Let's first assume we're reading a file with only the allocs.
291        let allocs: BTreeMap<Address, GenesisAccount> = match read_json_file(path) {
292            Ok(allocs) => allocs,
293            Err(_) => {
294                // Let's try and read from a genesis file, and extract allocs.
295                let genesis = read_json_file::<Genesis>(path)?;
296                genesis.alloc
297            }
298        };
299
300        // Then, load the allocs into the database.
301        let (db, journal, _) = ccx.ecx.as_db_env_and_journal();
302        db.load_allocs(&allocs, journal)
303            .map(|()| Vec::default())
304            .map_err(|e| fmt_err!("failed to load allocs: {e}"))
305    }
306}
307
308impl Cheatcode for cloneAccountCall {
309    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
310        let Self { source, target } = self;
311
312        let (db, journal, _) = ccx.ecx.as_db_env_and_journal();
313        let account = journal.load_account(db, *source)?;
314        let genesis = &genesis_account(account.data);
315        db.clone_account(genesis, target, journal)?;
316        // Cloned account should persist in forked envs.
317        ccx.ecx.journaled_state.database.add_persistent_account(*target);
318        Ok(Default::default())
319    }
320}
321
322impl Cheatcode for dumpStateCall {
323    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
324        let Self { pathToStateJson } = self;
325        let path = Path::new(pathToStateJson);
326
327        // Do not include system account or empty accounts in the dump.
328        let skip = |key: &Address, val: &Account| {
329            key == &CHEATCODE_ADDRESS
330                || key == &CALLER
331                || key == &HARDHAT_CONSOLE_ADDRESS
332                || key == &TEST_CONTRACT_ADDRESS
333                || key == &ccx.caller
334                || key == &ccx.state.config.evm_opts.sender
335                || val.is_empty()
336        };
337
338        let alloc = ccx
339            .ecx
340            .journaled_state
341            .state()
342            .iter_mut()
343            .filter(|(key, val)| !skip(key, val))
344            .map(|(key, val)| (key, genesis_account(val)))
345            .collect::<BTreeMap<_, _>>();
346
347        write_json_file(path, &alloc)?;
348        Ok(Default::default())
349    }
350}
351
352impl Cheatcode for recordCall {
353    fn apply(&self, state: &mut Cheatcodes) -> Result {
354        let Self {} = self;
355        state.recording_accesses = true;
356        state.accesses.clear();
357        Ok(Default::default())
358    }
359}
360
361impl Cheatcode for stopRecordCall {
362    fn apply(&self, state: &mut Cheatcodes) -> Result {
363        state.recording_accesses = false;
364        Ok(Default::default())
365    }
366}
367
368impl Cheatcode for accessesCall {
369    fn apply(&self, state: &mut Cheatcodes) -> Result {
370        let Self { target } = *self;
371        let result = (
372            state.accesses.reads.entry(target).or_default().as_slice(),
373            state.accesses.writes.entry(target).or_default().as_slice(),
374        );
375        Ok(result.abi_encode_params())
376    }
377}
378
379impl Cheatcode for recordLogsCall {
380    fn apply(&self, state: &mut Cheatcodes) -> Result {
381        let Self {} = self;
382        state.recorded_logs = Some(Default::default());
383        Ok(Default::default())
384    }
385}
386
387impl Cheatcode for getRecordedLogsCall {
388    fn apply(&self, state: &mut Cheatcodes) -> Result {
389        let Self {} = self;
390        Ok(state.recorded_logs.replace(Default::default()).unwrap_or_default().abi_encode())
391    }
392}
393
394impl Cheatcode for pauseGasMeteringCall {
395    fn apply(&self, state: &mut Cheatcodes) -> Result {
396        let Self {} = self;
397        state.gas_metering.paused = true;
398        Ok(Default::default())
399    }
400}
401
402impl Cheatcode for resumeGasMeteringCall {
403    fn apply(&self, state: &mut Cheatcodes) -> Result {
404        let Self {} = self;
405        state.gas_metering.resume();
406        Ok(Default::default())
407    }
408}
409
410impl Cheatcode for resetGasMeteringCall {
411    fn apply(&self, state: &mut Cheatcodes) -> Result {
412        let Self {} = self;
413        state.gas_metering.reset();
414        Ok(Default::default())
415    }
416}
417
418impl Cheatcode for lastCallGasCall {
419    fn apply(&self, state: &mut Cheatcodes) -> Result {
420        let Self {} = self;
421        let Some(last_call_gas) = &state.gas_metering.last_call_gas else {
422            bail!("no external call was made yet");
423        };
424        Ok(last_call_gas.abi_encode())
425    }
426}
427
428impl Cheatcode for getChainIdCall {
429    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
430        let Self {} = self;
431        Ok(U256::from(ccx.ecx.cfg.chain_id).abi_encode())
432    }
433}
434
435impl Cheatcode for chainIdCall {
436    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
437        let Self { newChainId } = self;
438        ensure!(*newChainId <= U256::from(u64::MAX), "chain ID must be less than 2^64");
439        ccx.ecx.cfg.chain_id = newChainId.to();
440        Ok(Default::default())
441    }
442}
443
444impl Cheatcode for coinbaseCall {
445    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
446        let Self { newCoinbase } = self;
447        ccx.ecx.block.beneficiary = *newCoinbase;
448        Ok(Default::default())
449    }
450}
451
452impl Cheatcode for difficultyCall {
453    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
454        let Self { newDifficulty } = self;
455        ensure!(
456            ccx.ecx.cfg.spec < SpecId::MERGE,
457            "`difficulty` is not supported after the Paris hard fork, use `prevrandao` instead; \
458             see EIP-4399: https://eips.ethereum.org/EIPS/eip-4399"
459        );
460        ccx.ecx.block.difficulty = *newDifficulty;
461        Ok(Default::default())
462    }
463}
464
465impl Cheatcode for feeCall {
466    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
467        let Self { newBasefee } = self;
468        ensure!(*newBasefee <= U256::from(u64::MAX), "base fee must be less than 2^64");
469        ccx.ecx.block.basefee = newBasefee.saturating_to();
470        Ok(Default::default())
471    }
472}
473
474impl Cheatcode for prevrandao_0Call {
475    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
476        let Self { newPrevrandao } = self;
477        ensure!(
478            ccx.ecx.cfg.spec >= SpecId::MERGE,
479            "`prevrandao` is not supported before the Paris hard fork, use `difficulty` instead; \
480             see EIP-4399: https://eips.ethereum.org/EIPS/eip-4399"
481        );
482        ccx.ecx.block.prevrandao = Some(*newPrevrandao);
483        Ok(Default::default())
484    }
485}
486
487impl Cheatcode for prevrandao_1Call {
488    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
489        let Self { newPrevrandao } = self;
490        ensure!(
491            ccx.ecx.cfg.spec >= SpecId::MERGE,
492            "`prevrandao` is not supported before the Paris hard fork, use `difficulty` instead; \
493             see EIP-4399: https://eips.ethereum.org/EIPS/eip-4399"
494        );
495        ccx.ecx.block.prevrandao = Some((*newPrevrandao).into());
496        Ok(Default::default())
497    }
498}
499
500impl Cheatcode for blobhashesCall {
501    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
502        let Self { hashes } = self;
503        ensure!(
504            ccx.ecx.cfg.spec >= SpecId::CANCUN,
505            "`blobhashes` is not supported before the Cancun hard fork; \
506             see EIP-4844: https://eips.ethereum.org/EIPS/eip-4844"
507        );
508        ccx.ecx.tx.blob_hashes.clone_from(hashes);
509        // force this as 4844 txtype
510        ccx.ecx.tx.tx_type = EIP4844_TX_TYPE_ID;
511        Ok(Default::default())
512    }
513}
514
515impl Cheatcode for getBlobhashesCall {
516    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
517        let Self {} = self;
518        ensure!(
519            ccx.ecx.cfg.spec >= SpecId::CANCUN,
520            "`getBlobhashes` is not supported before the Cancun hard fork; \
521             see EIP-4844: https://eips.ethereum.org/EIPS/eip-4844"
522        );
523        Ok(ccx.ecx.tx.blob_hashes.clone().abi_encode())
524    }
525}
526
527impl Cheatcode for rollCall {
528    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
529        let Self { newHeight } = self;
530        ccx.ecx.block.number = *newHeight;
531        Ok(Default::default())
532    }
533}
534
535impl Cheatcode for getBlockNumberCall {
536    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
537        let Self {} = self;
538        Ok(ccx.ecx.block.number.abi_encode())
539    }
540}
541
542impl Cheatcode for txGasPriceCall {
543    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
544        let Self { newGasPrice } = self;
545        ensure!(*newGasPrice <= U256::from(u64::MAX), "gas price must be less than 2^64");
546        ccx.ecx.tx.gas_price = newGasPrice.saturating_to();
547        Ok(Default::default())
548    }
549}
550
551impl Cheatcode for warpCall {
552    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
553        let Self { newTimestamp } = self;
554        ccx.ecx.block.timestamp = *newTimestamp;
555        Ok(Default::default())
556    }
557}
558
559impl Cheatcode for getBlockTimestampCall {
560    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
561        let Self {} = self;
562        Ok(ccx.ecx.block.timestamp.abi_encode())
563    }
564}
565
566impl Cheatcode for blobBaseFeeCall {
567    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
568        let Self { newBlobBaseFee } = self;
569        ensure!(
570            ccx.ecx.cfg.spec >= SpecId::CANCUN,
571            "`blobBaseFee` is not supported before the Cancun hard fork; \
572             see EIP-4844: https://eips.ethereum.org/EIPS/eip-4844"
573        );
574
575        ccx.ecx.block.set_blob_excess_gas_and_price(
576            (*newBlobBaseFee).to(),
577            get_blob_base_fee_update_fraction_by_spec_id(ccx.ecx.cfg.spec),
578        );
579        Ok(Default::default())
580    }
581}
582
583impl Cheatcode for getBlobBaseFeeCall {
584    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
585        let Self {} = self;
586        Ok(ccx.ecx.block.blob_excess_gas().unwrap_or(0).abi_encode())
587    }
588}
589
590impl Cheatcode for dealCall {
591    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
592        let Self { account: address, newBalance: new_balance } = *self;
593        let account = journaled_account(ccx.ecx, address)?;
594        let old_balance = std::mem::replace(&mut account.info.balance, new_balance);
595        let record = DealRecord { address, old_balance, new_balance };
596        ccx.state.eth_deals.push(record);
597        Ok(Default::default())
598    }
599}
600
601impl Cheatcode for etchCall {
602    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
603        let Self { target, newRuntimeBytecode } = self;
604        ccx.ensure_not_precompile(target)?;
605        ccx.ecx.journaled_state.load_account(*target)?;
606        let bytecode = Bytecode::new_raw_checked(newRuntimeBytecode.clone())
607            .map_err(|e| fmt_err!("failed to create bytecode: {e}"))?;
608        ccx.ecx.journaled_state.set_code(*target, bytecode);
609        Ok(Default::default())
610    }
611}
612
613impl Cheatcode for resetNonceCall {
614    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
615        let Self { account } = self;
616        let account = journaled_account(ccx.ecx, *account)?;
617        // Per EIP-161, EOA nonces start at 0, but contract nonces
618        // start at 1. Comparing by code_hash instead of code
619        // to avoid hitting the case where account's code is None.
620        let empty = account.info.code_hash == KECCAK_EMPTY;
621        let nonce = if empty { 0 } else { 1 };
622        account.info.nonce = nonce;
623        debug!(target: "cheatcodes", nonce, "reset");
624        Ok(Default::default())
625    }
626}
627
628impl Cheatcode for setNonceCall {
629    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
630        let Self { account, newNonce } = *self;
631        let account = journaled_account(ccx.ecx, account)?;
632        // nonce must increment only
633        let current = account.info.nonce;
634        ensure!(
635            newNonce >= current,
636            "new nonce ({newNonce}) must be strictly equal to or higher than the \
637             account's current nonce ({current})"
638        );
639        account.info.nonce = newNonce;
640        Ok(Default::default())
641    }
642}
643
644impl Cheatcode for setNonceUnsafeCall {
645    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
646        let Self { account, newNonce } = *self;
647        let account = journaled_account(ccx.ecx, account)?;
648        account.info.nonce = newNonce;
649        Ok(Default::default())
650    }
651}
652
653impl Cheatcode for storeCall {
654    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
655        let Self { target, slot, value } = *self;
656        ccx.ensure_not_precompile(&target)?;
657        ensure_loaded_account(ccx.ecx, target)?;
658        ccx.ecx.journaled_state.sstore(target, slot.into(), value.into())?;
659        Ok(Default::default())
660    }
661}
662
663impl Cheatcode for coolCall {
664    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
665        let Self { target } = self;
666        if let Some(account) = ccx.ecx.journaled_state.state.get_mut(target) {
667            account.unmark_touch();
668            account.storage.values_mut().for_each(|slot| slot.mark_cold());
669        }
670        Ok(Default::default())
671    }
672}
673
674impl Cheatcode for accessListCall {
675    fn apply(&self, state: &mut Cheatcodes) -> Result {
676        let Self { access } = self;
677        let access_list = access
678            .iter()
679            .map(|item| {
680                let keys = item.storageKeys.iter().map(|key| B256::from(*key)).collect_vec();
681                alloy_rpc_types::AccessListItem { address: item.target, storage_keys: keys }
682            })
683            .collect_vec();
684        state.access_list = Some(alloy_rpc_types::AccessList::from(access_list));
685        Ok(Default::default())
686    }
687}
688
689impl Cheatcode for noAccessListCall {
690    fn apply(&self, state: &mut Cheatcodes) -> Result {
691        let Self {} = self;
692        // Set to empty option in order to override previous applied access list.
693        if state.access_list.is_some() {
694            state.access_list = Some(alloy_rpc_types::AccessList::default());
695        }
696        Ok(Default::default())
697    }
698}
699
700impl Cheatcode for warmSlotCall {
701    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
702        let Self { target, slot } = *self;
703        set_cold_slot(ccx, target, slot.into(), false);
704        Ok(Default::default())
705    }
706}
707
708impl Cheatcode for coolSlotCall {
709    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
710        let Self { target, slot } = *self;
711        set_cold_slot(ccx, target, slot.into(), true);
712        Ok(Default::default())
713    }
714}
715
716impl Cheatcode for readCallersCall {
717    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
718        let Self {} = self;
719        read_callers(ccx.state, &ccx.ecx.tx.caller, ccx.ecx.journaled_state.depth())
720    }
721}
722
723impl Cheatcode for snapshotValue_0Call {
724    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
725        let Self { name, value } = self;
726        inner_value_snapshot(ccx, None, Some(name.clone()), value.to_string())
727    }
728}
729
730impl Cheatcode for snapshotValue_1Call {
731    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
732        let Self { group, name, value } = self;
733        inner_value_snapshot(ccx, Some(group.clone()), Some(name.clone()), value.to_string())
734    }
735}
736
737impl Cheatcode for snapshotGasLastCall_0Call {
738    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
739        let Self { name } = self;
740        let Some(last_call_gas) = &ccx.state.gas_metering.last_call_gas else {
741            bail!("no external call was made yet");
742        };
743        inner_last_gas_snapshot(ccx, None, Some(name.clone()), last_call_gas.gasTotalUsed)
744    }
745}
746
747impl Cheatcode for snapshotGasLastCall_1Call {
748    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
749        let Self { name, group } = self;
750        let Some(last_call_gas) = &ccx.state.gas_metering.last_call_gas else {
751            bail!("no external call was made yet");
752        };
753        inner_last_gas_snapshot(
754            ccx,
755            Some(group.clone()),
756            Some(name.clone()),
757            last_call_gas.gasTotalUsed,
758        )
759    }
760}
761
762impl Cheatcode for startSnapshotGas_0Call {
763    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
764        let Self { name } = self;
765        inner_start_gas_snapshot(ccx, None, Some(name.clone()))
766    }
767}
768
769impl Cheatcode for startSnapshotGas_1Call {
770    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
771        let Self { group, name } = self;
772        inner_start_gas_snapshot(ccx, Some(group.clone()), Some(name.clone()))
773    }
774}
775
776impl Cheatcode for stopSnapshotGas_0Call {
777    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
778        let Self {} = self;
779        inner_stop_gas_snapshot(ccx, None, None)
780    }
781}
782
783impl Cheatcode for stopSnapshotGas_1Call {
784    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
785        let Self { name } = self;
786        inner_stop_gas_snapshot(ccx, None, Some(name.clone()))
787    }
788}
789
790impl Cheatcode for stopSnapshotGas_2Call {
791    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
792        let Self { group, name } = self;
793        inner_stop_gas_snapshot(ccx, Some(group.clone()), Some(name.clone()))
794    }
795}
796
797// Deprecated in favor of `snapshotStateCall`
798impl Cheatcode for snapshotCall {
799    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
800        let Self {} = self;
801        inner_snapshot_state(ccx)
802    }
803}
804
805impl Cheatcode for snapshotStateCall {
806    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
807        let Self {} = self;
808        inner_snapshot_state(ccx)
809    }
810}
811
812// Deprecated in favor of `revertToStateCall`
813impl Cheatcode for revertToCall {
814    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
815        let Self { snapshotId } = self;
816        inner_revert_to_state(ccx, *snapshotId)
817    }
818}
819
820impl Cheatcode for revertToStateCall {
821    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
822        let Self { snapshotId } = self;
823        inner_revert_to_state(ccx, *snapshotId)
824    }
825}
826
827// Deprecated in favor of `revertToStateAndDeleteCall`
828impl Cheatcode for revertToAndDeleteCall {
829    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
830        let Self { snapshotId } = self;
831        inner_revert_to_state_and_delete(ccx, *snapshotId)
832    }
833}
834
835impl Cheatcode for revertToStateAndDeleteCall {
836    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
837        let Self { snapshotId } = self;
838        inner_revert_to_state_and_delete(ccx, *snapshotId)
839    }
840}
841
842// Deprecated in favor of `deleteStateSnapshotCall`
843impl Cheatcode for deleteSnapshotCall {
844    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
845        let Self { snapshotId } = self;
846        inner_delete_state_snapshot(ccx, *snapshotId)
847    }
848}
849
850impl Cheatcode for deleteStateSnapshotCall {
851    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
852        let Self { snapshotId } = self;
853        inner_delete_state_snapshot(ccx, *snapshotId)
854    }
855}
856
857// Deprecated in favor of `deleteStateSnapshotsCall`
858impl Cheatcode for deleteSnapshotsCall {
859    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
860        let Self {} = self;
861        inner_delete_state_snapshots(ccx)
862    }
863}
864
865impl Cheatcode for deleteStateSnapshotsCall {
866    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
867        let Self {} = self;
868        inner_delete_state_snapshots(ccx)
869    }
870}
871
872impl Cheatcode for startStateDiffRecordingCall {
873    fn apply(&self, state: &mut Cheatcodes) -> Result {
874        let Self {} = self;
875        state.recorded_account_diffs_stack = Some(Default::default());
876        // Enable mapping recording to track mapping slot accesses
877        state.mapping_slots.get_or_insert_default();
878        Ok(Default::default())
879    }
880}
881
882impl Cheatcode for stopAndReturnStateDiffCall {
883    fn apply(&self, state: &mut Cheatcodes) -> Result {
884        let Self {} = self;
885        get_state_diff(state)
886    }
887}
888
889impl Cheatcode for getStateDiffCall {
890    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
891        let mut diffs = String::new();
892        let state_diffs = get_recorded_state_diffs(ccx);
893        for (address, state_diffs) in state_diffs {
894            diffs.push_str(&format!("{address}\n"));
895            diffs.push_str(&format!("{state_diffs}\n"));
896        }
897        Ok(diffs.abi_encode())
898    }
899}
900
901impl Cheatcode for getStateDiffJsonCall {
902    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
903        let state_diffs = get_recorded_state_diffs(ccx);
904        Ok(serde_json::to_string(&state_diffs)?.abi_encode())
905    }
906}
907
908impl Cheatcode for getStorageAccessesCall {
909    fn apply(&self, state: &mut Cheatcodes) -> Result {
910        let mut storage_accesses = Vec::new();
911
912        if let Some(recorded_diffs) = &state.recorded_account_diffs_stack {
913            for account_accesses in recorded_diffs.iter().flatten() {
914                storage_accesses.extend(account_accesses.storageAccesses.clone());
915            }
916        }
917
918        Ok(storage_accesses.abi_encode())
919    }
920}
921
922impl Cheatcode for broadcastRawTransactionCall {
923    fn apply_full(&self, ccx: &mut CheatsCtxt, executor: &mut dyn CheatcodesExecutor) -> Result {
924        let tx = TxEnvelope::decode(&mut self.data.as_ref())
925            .map_err(|err| fmt_err!("failed to decode RLP-encoded transaction: {err}"))?;
926
927        let (db, journal, env) = ccx.ecx.as_db_env_and_journal();
928        db.transact_from_tx(
929            &tx.clone().into(),
930            env.to_owned(),
931            journal,
932            &mut *executor.get_inspector(ccx.state),
933        )?;
934
935        if ccx.state.broadcast.is_some() {
936            ccx.state.broadcastable_transactions.push_back(BroadcastableTransaction {
937                rpc: ccx.ecx.journaled_state.database.active_fork_url(),
938                transaction: tx.try_into()?,
939            });
940        }
941
942        Ok(Default::default())
943    }
944}
945
946impl Cheatcode for setBlockhashCall {
947    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
948        let Self { blockNumber, blockHash } = *self;
949        ensure!(blockNumber <= U256::from(u64::MAX), "blockNumber must be less than 2^64");
950        ensure!(
951            blockNumber <= U256::from(ccx.ecx.block.number),
952            "block number must be less than or equal to the current block number"
953        );
954
955        ccx.ecx.journaled_state.database.set_blockhash(blockNumber, blockHash);
956
957        Ok(Default::default())
958    }
959}
960
961impl Cheatcode for startDebugTraceRecordingCall {
962    fn apply_full(&self, ccx: &mut CheatsCtxt, executor: &mut dyn CheatcodesExecutor) -> Result {
963        let Some(tracer) = executor.tracing_inspector() else {
964            return Err(Error::from("no tracer initiated, consider adding -vvv flag"));
965        };
966
967        let mut info = RecordDebugStepInfo {
968            // will be updated later
969            start_node_idx: 0,
970            // keep the original config to revert back later
971            original_tracer_config: *tracer.config(),
972        };
973
974        // turn on tracer configuration for recording
975        tracer.update_config(|config| {
976            config
977                .set_steps(true)
978                .set_memory_snapshots(true)
979                .set_stack_snapshots(StackSnapshotType::Full)
980        });
981
982        // track where the recording starts
983        if let Some(last_node) = tracer.traces().nodes().last() {
984            info.start_node_idx = last_node.idx;
985        }
986
987        ccx.state.record_debug_steps_info = Some(info);
988        Ok(Default::default())
989    }
990}
991
992impl Cheatcode for stopAndReturnDebugTraceRecordingCall {
993    fn apply_full(&self, ccx: &mut CheatsCtxt, executor: &mut dyn CheatcodesExecutor) -> Result {
994        let Some(tracer) = executor.tracing_inspector() else {
995            return Err(Error::from("no tracer initiated, consider adding -vvv flag"));
996        };
997
998        let Some(record_info) = ccx.state.record_debug_steps_info else {
999            return Err(Error::from("nothing recorded"));
1000        };
1001
1002        // Use the trace nodes to flatten the call trace
1003        let root = tracer.traces();
1004        let steps = flatten_call_trace(0, root, record_info.start_node_idx);
1005
1006        let debug_steps: Vec<DebugStep> =
1007            steps.iter().map(|&step| convert_call_trace_to_debug_step(step)).collect();
1008        // Free up memory by clearing the steps if they are not recorded outside of cheatcode usage.
1009        if !record_info.original_tracer_config.record_steps {
1010            tracer.traces_mut().nodes_mut().iter_mut().for_each(|node| {
1011                node.trace.steps = Vec::new();
1012                node.logs = Vec::new();
1013                node.ordering = Vec::new();
1014            });
1015        }
1016
1017        // Revert the tracer config to the one before recording
1018        tracer.update_config(|_config| record_info.original_tracer_config);
1019
1020        // Clean up the recording info
1021        ccx.state.record_debug_steps_info = None;
1022
1023        Ok(debug_steps.abi_encode())
1024    }
1025}
1026
1027pub(super) fn get_nonce(ccx: &mut CheatsCtxt, address: &Address) -> Result {
1028    let account = ccx.ecx.journaled_state.load_account(*address)?;
1029    Ok(account.info.nonce.abi_encode())
1030}
1031
1032fn inner_snapshot_state(ccx: &mut CheatsCtxt) -> Result {
1033    let (db, journal, mut env) = ccx.ecx.as_db_env_and_journal();
1034    Ok(db.snapshot_state(journal, &mut env).abi_encode())
1035}
1036
1037fn inner_revert_to_state(ccx: &mut CheatsCtxt, snapshot_id: U256) -> Result {
1038    let (db, journal, mut env) = ccx.ecx.as_db_env_and_journal();
1039    let result = if let Some(journaled_state) =
1040        db.revert_state(snapshot_id, &*journal, &mut env, RevertStateSnapshotAction::RevertKeep)
1041    {
1042        // we reset the evm's journaled_state to the state of the snapshot previous state
1043        ccx.ecx.journaled_state.inner = journaled_state;
1044        true
1045    } else {
1046        false
1047    };
1048    Ok(result.abi_encode())
1049}
1050
1051fn inner_revert_to_state_and_delete(ccx: &mut CheatsCtxt, snapshot_id: U256) -> Result {
1052    let (db, journal, mut env) = ccx.ecx.as_db_env_and_journal();
1053
1054    let result = if let Some(journaled_state) =
1055        db.revert_state(snapshot_id, &*journal, &mut env, RevertStateSnapshotAction::RevertRemove)
1056    {
1057        // we reset the evm's journaled_state to the state of the snapshot previous state
1058        ccx.ecx.journaled_state.inner = journaled_state;
1059        true
1060    } else {
1061        false
1062    };
1063    Ok(result.abi_encode())
1064}
1065
1066fn inner_delete_state_snapshot(ccx: &mut CheatsCtxt, snapshot_id: U256) -> Result {
1067    let result = ccx.ecx.journaled_state.database.delete_state_snapshot(snapshot_id);
1068    Ok(result.abi_encode())
1069}
1070
1071fn inner_delete_state_snapshots(ccx: &mut CheatsCtxt) -> Result {
1072    ccx.ecx.journaled_state.database.delete_state_snapshots();
1073    Ok(Default::default())
1074}
1075
1076fn inner_value_snapshot(
1077    ccx: &mut CheatsCtxt,
1078    group: Option<String>,
1079    name: Option<String>,
1080    value: String,
1081) -> Result {
1082    let (group, name) = derive_snapshot_name(ccx, group, name);
1083
1084    ccx.state.gas_snapshots.entry(group).or_default().insert(name, value);
1085
1086    Ok(Default::default())
1087}
1088
1089fn inner_last_gas_snapshot(
1090    ccx: &mut CheatsCtxt,
1091    group: Option<String>,
1092    name: Option<String>,
1093    value: u64,
1094) -> Result {
1095    let (group, name) = derive_snapshot_name(ccx, group, name);
1096
1097    ccx.state.gas_snapshots.entry(group).or_default().insert(name, value.to_string());
1098
1099    Ok(value.abi_encode())
1100}
1101
1102fn inner_start_gas_snapshot(
1103    ccx: &mut CheatsCtxt,
1104    group: Option<String>,
1105    name: Option<String>,
1106) -> Result {
1107    // Revert if there is an active gas snapshot as we can only have one active snapshot at a time.
1108    if ccx.state.gas_metering.active_gas_snapshot.is_some() {
1109        let (group, name) = ccx.state.gas_metering.active_gas_snapshot.as_ref().unwrap().clone();
1110        bail!("gas snapshot was already started with group: {group} and name: {name}");
1111    }
1112
1113    let (group, name) = derive_snapshot_name(ccx, group, name);
1114
1115    ccx.state.gas_metering.gas_records.push(GasRecord {
1116        group: group.clone(),
1117        name: name.clone(),
1118        gas_used: 0,
1119        depth: ccx.ecx.journaled_state.depth(),
1120    });
1121
1122    ccx.state.gas_metering.active_gas_snapshot = Some((group, name));
1123
1124    ccx.state.gas_metering.start();
1125
1126    Ok(Default::default())
1127}
1128
1129fn inner_stop_gas_snapshot(
1130    ccx: &mut CheatsCtxt,
1131    group: Option<String>,
1132    name: Option<String>,
1133) -> Result {
1134    // If group and name are not provided, use the last snapshot group and name.
1135    let (group, name) = group.zip(name).unwrap_or_else(|| {
1136        let (group, name) = ccx.state.gas_metering.active_gas_snapshot.as_ref().unwrap().clone();
1137        (group, name)
1138    });
1139
1140    if let Some(record) = ccx
1141        .state
1142        .gas_metering
1143        .gas_records
1144        .iter_mut()
1145        .find(|record| record.group == group && record.name == name)
1146    {
1147        // Calculate the gas used since the snapshot was started.
1148        // We subtract 171 from the gas used to account for gas used by the snapshot itself.
1149        let value = record.gas_used.saturating_sub(171);
1150
1151        ccx.state
1152            .gas_snapshots
1153            .entry(group.clone())
1154            .or_default()
1155            .insert(name.clone(), value.to_string());
1156
1157        // Stop the gas metering.
1158        ccx.state.gas_metering.stop();
1159
1160        // Remove the gas record.
1161        ccx.state
1162            .gas_metering
1163            .gas_records
1164            .retain(|record| record.group != group && record.name != name);
1165
1166        // Clear last snapshot cache if we have an exact match.
1167        if let Some((snapshot_group, snapshot_name)) = &ccx.state.gas_metering.active_gas_snapshot
1168            && snapshot_group == &group
1169            && snapshot_name == &name
1170        {
1171            ccx.state.gas_metering.active_gas_snapshot = None;
1172        }
1173
1174        Ok(value.abi_encode())
1175    } else {
1176        bail!("no gas snapshot was started with the name: {name} in group: {group}");
1177    }
1178}
1179
1180// Derives the snapshot group and name from the provided group and name or the running contract.
1181fn derive_snapshot_name(
1182    ccx: &CheatsCtxt,
1183    group: Option<String>,
1184    name: Option<String>,
1185) -> (String, String) {
1186    let group = group.unwrap_or_else(|| {
1187        ccx.state.config.running_artifact.clone().expect("expected running contract").name
1188    });
1189    let name = name.unwrap_or_else(|| "default".to_string());
1190    (group, name)
1191}
1192
1193/// Reads the current caller information and returns the current [CallerMode], `msg.sender` and
1194/// `tx.origin`.
1195///
1196/// Depending on the current caller mode, one of the following results will be returned:
1197/// - If there is an active prank:
1198///     - caller_mode will be equal to:
1199///         - [CallerMode::Prank] if the prank has been set with `vm.prank(..)`.
1200///         - [CallerMode::RecurrentPrank] if the prank has been set with `vm.startPrank(..)`.
1201///     - `msg.sender` will be equal to the address set for the prank.
1202///     - `tx.origin` will be equal to the default sender address unless an alternative one has been
1203///       set when configuring the prank.
1204///
1205/// - If there is an active broadcast:
1206///     - caller_mode will be equal to:
1207///         - [CallerMode::Broadcast] if the broadcast has been set with `vm.broadcast(..)`.
1208///         - [CallerMode::RecurrentBroadcast] if the broadcast has been set with
1209///           `vm.startBroadcast(..)`.
1210///     - `msg.sender` and `tx.origin` will be equal to the address provided when setting the
1211///       broadcast.
1212///
1213/// - If no caller modification is active:
1214///     - caller_mode will be equal to [CallerMode::None],
1215///     - `msg.sender` and `tx.origin` will be equal to the default sender address.
1216fn read_callers(state: &Cheatcodes, default_sender: &Address, call_depth: usize) -> Result {
1217    let mut mode = CallerMode::None;
1218    let mut new_caller = default_sender;
1219    let mut new_origin = default_sender;
1220    if let Some(prank) = state.get_prank(call_depth) {
1221        mode = if prank.single_call { CallerMode::Prank } else { CallerMode::RecurrentPrank };
1222        new_caller = &prank.new_caller;
1223        if let Some(new) = &prank.new_origin {
1224            new_origin = new;
1225        }
1226    } else if let Some(broadcast) = &state.broadcast {
1227        mode = if broadcast.single_call {
1228            CallerMode::Broadcast
1229        } else {
1230            CallerMode::RecurrentBroadcast
1231        };
1232        new_caller = &broadcast.new_origin;
1233        new_origin = &broadcast.new_origin;
1234    }
1235
1236    Ok((mode, new_caller, new_origin).abi_encode_params())
1237}
1238
1239/// Ensures the `Account` is loaded and touched.
1240pub(super) fn journaled_account<'a>(
1241    ecx: Ecx<'a, '_, '_>,
1242    addr: Address,
1243) -> Result<&'a mut Account> {
1244    ensure_loaded_account(ecx, addr)?;
1245    Ok(ecx.journaled_state.state.get_mut(&addr).expect("account is loaded"))
1246}
1247
1248pub(super) fn ensure_loaded_account(ecx: Ecx, addr: Address) -> Result<()> {
1249    ecx.journaled_state.load_account(addr)?;
1250    ecx.journaled_state.touch(addr);
1251    Ok(())
1252}
1253
1254/// Consumes recorded account accesses and returns them as an abi encoded
1255/// array of [AccountAccess]. If there are no accounts were
1256/// recorded as accessed, an abi encoded empty array is returned.
1257///
1258/// In the case where `stopAndReturnStateDiff` is called at a lower
1259/// depth than `startStateDiffRecording`, multiple `Vec<RecordedAccountAccesses>`
1260/// will be flattened, preserving the order of the accesses.
1261fn get_state_diff(state: &mut Cheatcodes) -> Result {
1262    let res = state
1263        .recorded_account_diffs_stack
1264        .replace(Default::default())
1265        .unwrap_or_default()
1266        .into_iter()
1267        .flatten()
1268        .collect::<Vec<_>>();
1269    Ok(res.abi_encode())
1270}
1271
1272/// Helper function that creates a `GenesisAccount` from a regular `Account`.
1273fn genesis_account(account: &Account) -> GenesisAccount {
1274    GenesisAccount {
1275        nonce: Some(account.info.nonce),
1276        balance: account.info.balance,
1277        code: account.info.code.as_ref().map(|o| o.original_bytes()),
1278        storage: Some(
1279            account
1280                .storage
1281                .iter()
1282                .map(|(k, v)| (B256::from(*k), B256::from(v.present_value())))
1283                .collect(),
1284        ),
1285        private_key: None,
1286    }
1287}
1288
1289/// Helper function to returns state diffs recorded for each changed account.
1290fn get_recorded_state_diffs(ccx: &mut CheatsCtxt) -> BTreeMap<Address, AccountStateDiffs> {
1291    let mut state_diffs: BTreeMap<Address, AccountStateDiffs> = BTreeMap::default();
1292
1293    // First, collect all unique addresses we need to look up
1294    let mut addresses_to_lookup = HashSet::new();
1295    if let Some(records) = &ccx.state.recorded_account_diffs_stack {
1296        for account_access in records.iter().flatten() {
1297            if !account_access.storageAccesses.is_empty()
1298                || account_access.oldBalance != account_access.newBalance
1299            {
1300                addresses_to_lookup.insert(account_access.account);
1301                for storage_access in &account_access.storageAccesses {
1302                    if storage_access.isWrite && !storage_access.reverted {
1303                        addresses_to_lookup.insert(storage_access.account);
1304                    }
1305                }
1306            }
1307        }
1308    }
1309
1310    // Look up contract names and storage layouts for all addresses
1311    let mut contract_names = HashMap::new();
1312    let mut storage_layouts = HashMap::new();
1313    for address in addresses_to_lookup {
1314        if let Some((artifact_id, _)) = get_contract_data(ccx, address) {
1315            contract_names.insert(address, artifact_id.identifier());
1316        }
1317
1318        // Also get storage layout if available
1319        if let Some((_artifact_id, contract_data)) = get_contract_data(ccx, address)
1320            && let Some(storage_layout) = &contract_data.storage_layout
1321        {
1322            storage_layouts.insert(address, storage_layout.clone());
1323        }
1324    }
1325
1326    // Now process the records
1327    if let Some(records) = &ccx.state.recorded_account_diffs_stack {
1328        records
1329            .iter()
1330            .flatten()
1331            .filter(|account_access| {
1332                !account_access.storageAccesses.is_empty()
1333                    || account_access.oldBalance != account_access.newBalance
1334                    || account_access.oldNonce != account_access.newNonce
1335            })
1336            .for_each(|account_access| {
1337                // Record account balance diffs.
1338                if account_access.oldBalance != account_access.newBalance {
1339                    let account_diff =
1340                        state_diffs.entry(account_access.account).or_insert_with(|| {
1341                            AccountStateDiffs {
1342                                label: ccx.state.labels.get(&account_access.account).cloned(),
1343                                contract: contract_names.get(&account_access.account).cloned(),
1344                                ..Default::default()
1345                            }
1346                        });
1347                    // Update balance diff. Do not overwrite the initial balance if already set.
1348                    if let Some(diff) = &mut account_diff.balance_diff {
1349                        diff.new_value = account_access.newBalance;
1350                    } else {
1351                        account_diff.balance_diff = Some(BalanceDiff {
1352                            previous_value: account_access.oldBalance,
1353                            new_value: account_access.newBalance,
1354                        });
1355                    }
1356                }
1357
1358                // Record account nonce diffs.
1359                if account_access.oldNonce != account_access.newNonce {
1360                    let account_diff =
1361                        state_diffs.entry(account_access.account).or_insert_with(|| {
1362                            AccountStateDiffs {
1363                                label: ccx.state.labels.get(&account_access.account).cloned(),
1364                                contract: contract_names.get(&account_access.account).cloned(),
1365                                ..Default::default()
1366                            }
1367                        });
1368                    // Update nonce diff. Do not overwrite the initial nonce if already set.
1369                    if let Some(diff) = &mut account_diff.nonce_diff {
1370                        diff.new_value = account_access.newNonce;
1371                    } else {
1372                        account_diff.nonce_diff = Some(NonceDiff {
1373                            previous_value: account_access.oldNonce,
1374                            new_value: account_access.newNonce,
1375                        });
1376                    }
1377                }
1378
1379                // Record account state diffs.
1380                for storage_access in &account_access.storageAccesses {
1381                    if storage_access.isWrite && !storage_access.reverted {
1382                        let account_diff = state_diffs
1383                            .entry(storage_access.account)
1384                            .or_insert_with(|| AccountStateDiffs {
1385                                label: ccx.state.labels.get(&storage_access.account).cloned(),
1386                                contract: contract_names.get(&storage_access.account).cloned(),
1387                                ..Default::default()
1388                            });
1389                        let layout = storage_layouts.get(&storage_access.account);
1390                        // Update state diff. Do not overwrite the initial value if already set.
1391                        match account_diff.state_diff.entry(storage_access.slot) {
1392                            Entry::Vacant(slot_state_diff) => {
1393                                // Get storage layout info for this slot
1394                                // Include mapping slots if available for the account
1395                                let mapping_slots = ccx
1396                                    .state
1397                                    .mapping_slots
1398                                    .as_ref()
1399                                    .and_then(|slots| slots.get(&storage_access.account));
1400
1401                                let mut slot_info = layout.and_then(|layout| {
1402                                    let decoder = SlotIdentifier::new(layout.clone());
1403                                    decoder.identify(&storage_access.slot, mapping_slots)
1404                                });
1405
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
1414                                slot_state_diff.insert(SlotStateDiff {
1415                                    previous_value: storage_access.previousValue,
1416                                    new_value: storage_access.newValue,
1417                                    slot_info,
1418                                });
1419                            }
1420                            Entry::Occupied(mut slot_state_diff) => {
1421                                let entry = slot_state_diff.get_mut();
1422                                entry.new_value = storage_access.newValue;
1423
1424                                // Update decoded values if we have slot info
1425                                if let Some(ref mut slot_info) = entry.slot_info {
1426                                    slot_info.decode_values(
1427                                        entry.previous_value,
1428                                        storage_access.newValue,
1429                                    );
1430                                }
1431                            }
1432                        }
1433                    }
1434                }
1435            });
1436    }
1437    state_diffs
1438}
1439
1440/// EIP-1967 implementation storage slot
1441const EIP1967_IMPL_SLOT: &str = "360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc";
1442
1443/// EIP-1822 UUPS implementation storage slot: keccak256("PROXIABLE")
1444const EIP1822_PROXIABLE_SLOT: &str =
1445    "c5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7";
1446
1447/// Helper function to get the contract data from the deployed code at an address.
1448fn get_contract_data<'a>(
1449    ccx: &'a mut CheatsCtxt,
1450    address: Address,
1451) -> Option<(&'a foundry_compilers::ArtifactId, &'a foundry_common::contracts::ContractData)> {
1452    // Check if we have available artifacts to match against
1453    let artifacts = ccx.state.config.available_artifacts.as_ref()?;
1454
1455    // Try to load the account and get its code
1456    let account = ccx.ecx.journaled_state.load_account(address).ok()?;
1457    let code = account.info.code.as_ref()?;
1458
1459    // Skip if code is empty
1460    if code.is_empty() {
1461        return None;
1462    }
1463
1464    // Try to find the artifact by deployed code
1465    let code_bytes = code.original_bytes();
1466    // First check for proxy patterns
1467    let hex_str = hex::encode(&code_bytes);
1468    let find_by_suffix =
1469        |suffix: &str| artifacts.iter().find(|(a, _)| a.identifier().ends_with(suffix));
1470    // Simple proxy detection based on storage slot patterns
1471    if hex_str.contains(EIP1967_IMPL_SLOT)
1472        && let Some(result) = find_by_suffix(":TransparentUpgradeableProxy")
1473    {
1474        return Some(result);
1475    } else if hex_str.contains(EIP1822_PROXIABLE_SLOT)
1476        && let Some(result) = find_by_suffix(":UUPSUpgradeable")
1477    {
1478        return Some(result);
1479    }
1480
1481    // Try exact match
1482    if let Some(result) = artifacts.find_by_deployed_code_exact(&code_bytes) {
1483        return Some(result);
1484    }
1485
1486    // Fallback to fuzzy matching if exact match fails
1487    artifacts.find_by_deployed_code(&code_bytes)
1488}
1489
1490/// Helper function to set / unset cold storage slot of the target address.
1491fn set_cold_slot(ccx: &mut CheatsCtxt, target: Address, slot: U256, cold: bool) {
1492    if let Some(account) = ccx.ecx.journaled_state.state.get_mut(&target)
1493        && let Some(storage_slot) = account.storage.get_mut(&slot)
1494    {
1495        storage_slot.is_cold = cold;
1496    }
1497}