Skip to main content

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