foundry_cheatcodes/
state.rs

1//! Cheatcode state shared by both execution engines.
2
3use crate::{
4    CheatcodeAnalysis, CheatsConfig, ExpectedCallTracker, Wallets,
5    evm::mock::{MockCallDataContext, MockCallReturnData},
6    inspector::ArbitraryStorage,
7};
8use alloy_primitives::{
9    Address, Bytes, U256,
10    map::{AddressHashMap, HashMap},
11};
12use alloy_rpc_types::AccessList;
13use alloy_signer_local::PrivateKeySigner;
14use foundry_common::mapping_slots::MappingSlots;
15use foundry_wallets::wallet_multi::MultiWallet;
16use proptest::test_runner::{RngAlgorithm, TestRng, TestRunner};
17use rand::Rng;
18use serde_json::Value;
19use std::{
20    collections::{BTreeMap, VecDeque},
21    fs::File,
22    io::BufReader,
23    path::PathBuf,
24    sync::Arc,
25};
26
27/// Cheatcode state that does not depend on the execution engine.
28#[derive(Clone, Debug)]
29pub struct CheatcodesCore {
30    /// Additional, user configurable context this Inspector has access to when inspecting a call.
31    pub config: Arc<CheatsConfig>,
32    /// Solar compiler instance, to grant syntactic and semantic analysis capabilities
33    pub analysis: Option<CheatcodeAnalysis>,
34    /// Address labels
35    pub labels: AddressHashMap<String>,
36    /// Test-scoped context holding data that needs to be reset every test run
37    pub test_context: TestContext,
38    /// Whether to commit FS changes such as file creations, writes and deletes.
39    /// Used to prevent duplicate changes file executing non-committing calls.
40    pub fs_commit: bool,
41    /// Serialized JSON values.
42    // **Note**: both must a BTreeMap to ensure the order of the keys is deterministic.
43    pub serialized_jsons: BTreeMap<String, BTreeMap<String, Value>>,
44    /// Optional cheatcodes `TestRunner`. Used for generating random values from uint and int
45    /// strategies.
46    test_runner: Option<TestRunner>,
47    /// Unlocked wallets used in scripts and testing of scripts.
48    pub wallets: Option<Wallets>,
49    /// Parsed secp256k1 private-key signers for repeated `vm.addr` / `vm.sign` calls.
50    pub private_key_signers: HashMap<U256, PrivateKeySigner>,
51    /// Mocked calls
52    // **Note**: inner must a BTreeMap because of special `Ord` impl for `MockCallDataContext`
53    pub mocked_calls: HashMap<Address, BTreeMap<MockCallDataContext, VecDeque<MockCallReturnData>>>,
54    /// Mocked functions. Maps target address to be mocked to pair of (calldata, mock address).
55    pub mocked_functions: HashMap<Address, HashMap<Bytes, Address>>,
56    /// Expected calls
57    pub expected_calls: ExpectedCallTracker,
58    /// Mapping slots.
59    pub mapping_slots: Option<AddressHashMap<MappingSlots>>,
60    /// Recorded logs
61    pub recorded_logs: Option<Vec<crate::Vm::Log>>,
62    /// Current EIP-2930 access lists.
63    pub access_list: Option<AccessList>,
64    /// Whether the next contract creation should be intercepted to return its initcode.
65    pub intercept_next_create_call: bool,
66    /// Addresses with arbitrary storage.
67    pub arbitrary_storage: Option<ArbitraryStorage>,
68}
69
70impl Default for CheatcodesCore {
71    fn default() -> Self {
72        Self::new(Default::default())
73    }
74}
75
76impl CheatcodesCore {
77    /// Creates a new `CheatcodesCore` with the given settings.
78    pub fn new(config: Arc<CheatsConfig>) -> Self {
79        Self {
80            analysis: None,
81            labels: config.labels.clone(),
82            config,
83            test_context: Default::default(),
84            fs_commit: true,
85            serialized_jsons: Default::default(),
86            test_runner: Default::default(),
87            wallets: Default::default(),
88            private_key_signers: Default::default(),
89            mocked_calls: Default::default(),
90            mocked_functions: Default::default(),
91            expected_calls: Default::default(),
92            mapping_slots: Default::default(),
93            recorded_logs: Default::default(),
94            access_list: Default::default(),
95            intercept_next_create_call: Default::default(),
96            arbitrary_storage: Default::default(),
97        }
98    }
99
100    /// Enables cheatcode analysis capabilities by providing a solar compiler instance.
101    pub fn set_analysis(&mut self, analysis: CheatcodeAnalysis) {
102        self.analysis = Some(analysis);
103    }
104
105    /// Returns the configured wallets if available, else creates a new instance.
106    pub fn wallets(&mut self) -> &Wallets {
107        self.wallets.get_or_insert_with(|| Wallets::new(MultiWallet::default(), None))
108    }
109
110    /// Sets the unlocked wallets.
111    pub fn set_wallets(&mut self, wallets: Wallets) {
112        self.wallets = Some(wallets);
113    }
114
115    pub fn rng(&mut self) -> &mut impl Rng {
116        self.test_runner().rng()
117    }
118
119    pub fn test_runner(&mut self) -> &mut TestRunner {
120        self.test_runner.get_or_insert_with(|| match self.config.seed {
121            Some(seed) => TestRunner::new_with_rng(
122                proptest::test_runner::Config::default(),
123                TestRng::from_seed(RngAlgorithm::ChaCha, &seed.to_be_bytes::<32>()),
124            ),
125            None => TestRunner::new(proptest::test_runner::Config::default()),
126        })
127    }
128
129    pub fn set_seed(&mut self, seed: U256) {
130        self.test_runner = Some(TestRunner::new_with_rng(
131            proptest::test_runner::Config::default(),
132            TestRng::from_seed(RngAlgorithm::ChaCha, &seed.to_be_bytes::<32>()),
133        ));
134    }
135
136    /// Returns struct definitions from the analysis, if available.
137    pub fn struct_defs(&self) -> Option<&foundry_common::fmt::StructDefinitions> {
138        self.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok())
139    }
140
141    /// Returns existing or set a default `ArbitraryStorage` option.
142    /// Used by `setArbitraryStorage` cheatcode to track addresses with arbitrary storage.
143    pub fn arbitrary_storage(&mut self) -> &mut ArbitraryStorage {
144        self.arbitrary_storage.get_or_insert_with(ArbitraryStorage::default)
145    }
146}
147
148/// Contains additional, test specific resources that should be kept for the duration of the test
149#[derive(Debug, Default)]
150pub struct TestContext {
151    /// Buffered readers for files opened for reading (path => BufReader mapping)
152    pub opened_read_files: HashMap<PathBuf, BufReader<File>>,
153}
154
155/// Every time we clone `Context`, we want it to be empty
156impl Clone for TestContext {
157    fn clone(&self) -> Self {
158        Default::default()
159    }
160}
161
162impl TestContext {
163    /// Clears the context.
164    pub fn clear(&mut self) {
165        self.opened_read_files.clear();
166    }
167}
foundry_cheatcodes/state.rs

foundry_cheatcodes/
state.rs
