foundry_cheatcodes/
utils.rs

Help
1//! Implementations of [`Utilities`](spec::Group::Utilities) cheatcodes.
2
3use crate::{Cheatcode, Cheatcodes, CheatcodesExecutor, CheatsCtxt, Result, Vm::*};
4use alloy_dyn_abi::{DynSolType, DynSolValue};
5use alloy_ens::namehash;
6use alloy_primitives::{aliases::B32, map::HashMap, B64, U256};
7use alloy_sol_types::SolValue;
8use foundry_evm_core::constants::DEFAULT_CREATE2_DEPLOYER;
9use proptest::prelude::Strategy;
10use rand::{seq::SliceRandom, Rng, RngCore};
11use revm::context::JournalTr;
12
13/// Contains locations of traces ignored via cheatcodes.
14///
15/// The way we identify location in traces is by (node_idx, item_idx) tuple where node_idx is an
16/// index of a call trace node, and item_idx is a value between 0 and `node.ordering.len()` where i
17/// represents point after ith item, and 0 represents the beginning of the node trace.
18#[derive(Debug, Default, Clone)]
19pub struct IgnoredTraces {
20    /// Mapping from (start_node_idx, start_item_idx) to (end_node_idx, end_item_idx) representing
21    /// ranges of trace nodes to ignore.
22    pub ignored: HashMap<(usize, usize), (usize, usize)>,
23    /// Keeps track of (start_node_idx, start_item_idx) of the last `vm.pauseTracing` call.
24    pub last_pause_call: Option<(usize, usize)>,
25}
26
27impl Cheatcode for labelCall {
28    fn apply(&self, state: &mut Cheatcodes) -> Result {
29        let Self { account, newLabel } = self;
30        state.labels.insert(*account, newLabel.clone());
31        Ok(Default::default())
32    }
33}
34
35impl Cheatcode for getLabelCall {
36    fn apply(&self, state: &mut Cheatcodes) -> Result {
37        let Self { account } = self;
38        Ok(match state.labels.get(account) {
39            Some(label) => label.abi_encode(),
40            None => format!("unlabeled:{account}").abi_encode(),
41        })
42    }
43}
44
45impl Cheatcode for computeCreateAddressCall {
46    fn apply(&self, _state: &mut Cheatcodes) -> Result {
47        let Self { nonce, deployer } = self;
48        ensure!(*nonce <= U256::from(u64::MAX), "nonce must be less than 2^64 - 1");
49        Ok(deployer.create(nonce.to()).abi_encode())
50    }
51}
52
53impl Cheatcode for computeCreate2Address_0Call {
54    fn apply(&self, _state: &mut Cheatcodes) -> Result {
55        let Self { salt, initCodeHash, deployer } = self;
56        Ok(deployer.create2(salt, initCodeHash).abi_encode())
57    }
58}
59
60impl Cheatcode for computeCreate2Address_1Call {
61    fn apply(&self, _state: &mut Cheatcodes) -> Result {
62        let Self { salt, initCodeHash } = self;
63        Ok(DEFAULT_CREATE2_DEPLOYER.create2(salt, initCodeHash).abi_encode())
64    }
65}
66
67impl Cheatcode for ensNamehashCall {
68    fn apply(&self, _state: &mut Cheatcodes) -> Result {
69        let Self { name } = self;
70        Ok(namehash(name).abi_encode())
71    }
72}
73
74impl Cheatcode for randomUint_0Call {
75    fn apply(&self, state: &mut Cheatcodes) -> Result {
76        random_uint(state, None, None)
77    }
78}
79
80impl Cheatcode for randomUint_1Call {
81    fn apply(&self, state: &mut Cheatcodes) -> Result {
82        let Self { min, max } = *self;
83        random_uint(state, None, Some((min, max)))
84    }
85}
86
87impl Cheatcode for randomUint_2Call {
88    fn apply(&self, state: &mut Cheatcodes) -> Result {
89        let Self { bits } = *self;
90        random_uint(state, Some(bits), None)
91    }
92}
93
94impl Cheatcode for randomAddressCall {
95    fn apply(&self, state: &mut Cheatcodes) -> Result {
96        Ok(DynSolValue::type_strategy(&DynSolType::Address)
97            .new_tree(state.test_runner())
98            .unwrap()
99            .current()
100            .abi_encode())
101    }
102}
103
104impl Cheatcode for randomInt_0Call {
105    fn apply(&self, state: &mut Cheatcodes) -> Result {
106        random_int(state, None)
107    }
108}
109
110impl Cheatcode for randomInt_1Call {
111    fn apply(&self, state: &mut Cheatcodes) -> Result {
112        let Self { bits } = *self;
113        random_int(state, Some(bits))
114    }
115}
116
117impl Cheatcode for randomBoolCall {
118    fn apply(&self, state: &mut Cheatcodes) -> Result {
119        let rand_bool: bool = state.rng().random();
120        Ok(rand_bool.abi_encode())
121    }
122}
123
124impl Cheatcode for randomBytesCall {
125    fn apply(&self, state: &mut Cheatcodes) -> Result {
126        let Self { len } = *self;
127        ensure!(
128            len <= U256::from(usize::MAX),
129            format!("bytes length cannot exceed {}", usize::MAX)
130        );
131        let mut bytes = vec![0u8; len.to::<usize>()];
132        state.rng().fill_bytes(&mut bytes);
133        Ok(bytes.abi_encode())
134    }
135}
136
137impl Cheatcode for randomBytes4Call {
138    fn apply(&self, state: &mut Cheatcodes) -> Result {
139        let rand_u32 = state.rng().next_u32();
140        Ok(B32::from(rand_u32).abi_encode())
141    }
142}
143
144impl Cheatcode for randomBytes8Call {
145    fn apply(&self, state: &mut Cheatcodes) -> Result {
146        let rand_u64 = state.rng().next_u64();
147        Ok(B64::from(rand_u64).abi_encode())
148    }
149}
150
151impl Cheatcode for pauseTracingCall {
152    fn apply_full(
153        &self,
154        ccx: &mut crate::CheatsCtxt,
155        executor: &mut dyn CheatcodesExecutor,
156    ) -> Result {
157        let Some(tracer) = executor.tracing_inspector().and_then(|t| t.as_ref()) else {
158            // No tracer -> nothing to pause
159            return Ok(Default::default())
160        };
161
162        // If paused earlier, ignore the call
163        if ccx.state.ignored_traces.last_pause_call.is_some() {
164            return Ok(Default::default())
165        }
166
167        let cur_node = &tracer.traces().nodes().last().expect("no trace nodes");
168        ccx.state.ignored_traces.last_pause_call = Some((cur_node.idx, cur_node.ordering.len()));
169
170        Ok(Default::default())
171    }
172}
173
174impl Cheatcode for resumeTracingCall {
175    fn apply_full(
176        &self,
177        ccx: &mut crate::CheatsCtxt,
178        executor: &mut dyn CheatcodesExecutor,
179    ) -> Result {
180        let Some(tracer) = executor.tracing_inspector().and_then(|t| t.as_ref()) else {
181            // No tracer -> nothing to unpause
182            return Ok(Default::default())
183        };
184
185        let Some(start) = ccx.state.ignored_traces.last_pause_call.take() else {
186            // Nothing to unpause
187            return Ok(Default::default())
188        };
189
190        let node = &tracer.traces().nodes().last().expect("no trace nodes");
191        ccx.state.ignored_traces.ignored.insert(start, (node.idx, node.ordering.len()));
192
193        Ok(Default::default())
194    }
195}
196
197impl Cheatcode for interceptInitcodeCall {
198    fn apply(&self, state: &mut Cheatcodes) -> Result {
199        let Self {} = self;
200        if !state.intercept_next_create_call {
201            state.intercept_next_create_call = true;
202        } else {
203            bail!("vm.interceptInitcode() has already been called")
204        }
205        Ok(Default::default())
206    }
207}
208
209impl Cheatcode for setArbitraryStorage_0Call {
210    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
211        let Self { target } = self;
212        ccx.state.arbitrary_storage().mark_arbitrary(target, false);
213
214        Ok(Default::default())
215    }
216}
217
218impl Cheatcode for setArbitraryStorage_1Call {
219    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
220        let Self { target, overwrite } = self;
221        ccx.state.arbitrary_storage().mark_arbitrary(target, *overwrite);
222
223        Ok(Default::default())
224    }
225}
226
227impl Cheatcode for copyStorageCall {
228    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
229        let Self { from, to } = self;
230
231        ensure!(
232            !ccx.state.has_arbitrary_storage(to),
233            "target address cannot have arbitrary storage"
234        );
235
236        if let Ok(from_account) = ccx.ecx.journaled_state.load_account(*from) {
237            let from_storage = from_account.storage.clone();
238            if let Ok(mut to_account) = ccx.ecx.journaled_state.load_account(*to) {
239                to_account.storage = from_storage;
240                if let Some(ref mut arbitrary_storage) = &mut ccx.state.arbitrary_storage {
241                    arbitrary_storage.mark_copy(from, to);
242                }
243            }
244        }
245
246        Ok(Default::default())
247    }
248}
249
250impl Cheatcode for sortCall {
251    fn apply(&self, _state: &mut Cheatcodes) -> Result {
252        let Self { array } = self;
253
254        let mut sorted_values = array.clone();
255        sorted_values.sort();
256
257        Ok(sorted_values.abi_encode())
258    }
259}
260
261impl Cheatcode for shuffleCall {
262    fn apply(&self, state: &mut Cheatcodes) -> Result {
263        let Self { array } = self;
264
265        let mut shuffled_values = array.clone();
266        let rng = state.rng();
267        shuffled_values.shuffle(rng);
268
269        Ok(shuffled_values.abi_encode())
270    }
271}
272
273/// Helper to generate a random `uint` value (with given bits or bounded if specified)
274/// from type strategy.
275fn random_uint(state: &mut Cheatcodes, bits: Option<U256>, bounds: Option<(U256, U256)>) -> Result {
276    if let Some(bits) = bits {
277        // Generate random with specified bits.
278        ensure!(bits <= U256::from(256), "number of bits cannot exceed 256");
279        return Ok(DynSolValue::type_strategy(&DynSolType::Uint(bits.to::<usize>()))
280            .new_tree(state.test_runner())
281            .unwrap()
282            .current()
283            .abi_encode())
284    }
285
286    if let Some((min, max)) = bounds {
287        ensure!(min <= max, "min must be less than or equal to max");
288        // Generate random between range min..=max
289        let exclusive_modulo = max - min;
290        let mut random_number: U256 = state.rng().random();
291        if exclusive_modulo != U256::MAX {
292            let inclusive_modulo = exclusive_modulo + U256::from(1);
293            random_number %= inclusive_modulo;
294        }
295        random_number += min;
296        return Ok(random_number.abi_encode())
297    }
298
299    // Generate random `uint256` value.
300    Ok(DynSolValue::type_strategy(&DynSolType::Uint(256))
301        .new_tree(state.test_runner())
302        .unwrap()
303        .current()
304        .abi_encode())
305}
306
307/// Helper to generate a random `int` value (with given bits if specified) from type strategy.
308fn random_int(state: &mut Cheatcodes, bits: Option<U256>) -> Result {
309    let no_bits = bits.unwrap_or(U256::from(256));
310    ensure!(no_bits <= U256::from(256), "number of bits cannot exceed 256");
311    Ok(DynSolValue::type_strategy(&DynSolType::Int(no_bits.to::<usize>()))
312        .new_tree(state.test_runner())
313        .unwrap()
314        .current()
315        .abi_encode())
316}
foundry_cheatcodes/utils.rs

foundry_cheatcodes/
utils.rs
