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

foundry_cheatcodes/
utils.rs
