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