foundry_cheatcodes/test/
expect.rs

Help
1use std::{
2    collections::VecDeque,
3    fmt::{self, Display},
4};
5
6use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Error, Result, Vm::*};
7use alloy_primitives::{
8    map::{hash_map::Entry, AddressHashMap, HashMap},
9    Address, Bytes, LogData as RawLog, U256,
10};
11use revm::{
12    context::JournalTr,
13    interpreter::{InstructionResult, Interpreter, InterpreterAction, InterpreterResult},
14};
15
16use super::revert_handlers::RevertParameters;
17/// Tracks the expected calls per address.
18///
19/// For each address, we track the expected calls per call data. We track it in such manner
20/// so that we don't mix together calldatas that only contain selectors and calldatas that contain
21/// selector and arguments (partial and full matches).
22///
23/// This then allows us to customize the matching behavior for each call data on the
24/// `ExpectedCallData` struct and track how many times we've actually seen the call on the second
25/// element of the tuple.
26pub type ExpectedCallTracker = HashMap<Address, HashMap<Bytes, (ExpectedCallData, u64)>>;
27
28#[derive(Clone, Debug)]
29pub struct ExpectedCallData {
30    /// The expected value sent in the call
31    pub value: Option<U256>,
32    /// The expected gas supplied to the call
33    pub gas: Option<u64>,
34    /// The expected *minimum* gas supplied to the call
35    pub min_gas: Option<u64>,
36    /// The number of times the call is expected to be made.
37    /// If the type of call is `NonCount`, this is the lower bound for the number of calls
38    /// that must be seen.
39    /// If the type of call is `Count`, this is the exact number of calls that must be seen.
40    pub count: u64,
41    /// The type of expected call.
42    pub call_type: ExpectedCallType,
43}
44
45/// The type of expected call.
46#[derive(Clone, Debug, PartialEq, Eq)]
47pub enum ExpectedCallType {
48    /// The call is expected to be made at least once.
49    NonCount,
50    /// The exact number of calls expected.
51    Count,
52}
53
54/// The type of expected revert.
55#[derive(Clone, Debug)]
56pub enum ExpectedRevertKind {
57    /// Expects revert from the next non-cheatcode call.
58    Default,
59    /// Expects revert from the next cheatcode call.
60    ///
61    /// The `pending_processing` flag is used to track whether we have exited
62    /// `expectCheatcodeRevert` context or not.
63    /// We have to track it to avoid expecting `expectCheatcodeRevert` call to revert itself.
64    Cheatcode { pending_processing: bool },
65}
66
67#[derive(Clone, Debug)]
68pub struct ExpectedRevert {
69    /// The expected data returned by the revert, None being any.
70    pub reason: Option<Vec<u8>>,
71    /// The depth at which the revert is expected.
72    pub depth: usize,
73    /// The type of expected revert.
74    pub kind: ExpectedRevertKind,
75    /// If true then only the first 4 bytes of expected data returned by the revert are checked.
76    pub partial_match: bool,
77    /// Contract expected to revert next call.
78    pub reverter: Option<Address>,
79    /// Address that reverted the call.
80    pub reverted_by: Option<Address>,
81    /// Max call depth reached during next call execution.
82    pub max_depth: usize,
83    /// Number of times this revert is expected.
84    pub count: u64,
85    /// Actual number of times this revert has been seen.
86    pub actual_count: u64,
87}
88
89#[derive(Clone, Debug)]
90pub struct ExpectedEmit {
91    /// The depth at which we expect this emit to have occurred
92    pub depth: usize,
93    /// The log we expect
94    pub log: Option<RawLog>,
95    /// The checks to perform:
96    /// ```text
97    /// ┌───────┬───────┬───────┬───────┬────┐
98    /// │topic 0│topic 1│topic 2│topic 3│data│
99    /// └───────┴───────┴───────┴───────┴────┘
100    /// ```
101    pub checks: [bool; 5],
102    /// If present, check originating address against this
103    pub address: Option<Address>,
104    /// If present, relax the requirement that topic 0 must be present. This allows anonymous
105    /// events with no indexed topics to be matched.
106    pub anonymous: bool,
107    /// Whether the log was actually found in the subcalls
108    pub found: bool,
109    /// Number of times the log is expected to be emitted
110    pub count: u64,
111}
112
113#[derive(Clone, Debug)]
114pub struct ExpectedCreate {
115    /// The address that deployed the contract
116    pub deployer: Address,
117    /// Runtime bytecode of the contract
118    pub bytecode: Bytes,
119    /// Whether deployed with CREATE or CREATE2
120    pub create_scheme: CreateScheme,
121}
122
123#[derive(Clone, Debug)]
124pub enum CreateScheme {
125    Create,
126    Create2,
127}
128
129impl Display for CreateScheme {
130    fn fmt(&self, f: &mut fmt::Formatter) -> std::fmt::Result {
131        match self {
132            Self::Create => write!(f, "CREATE"),
133            Self::Create2 => write!(f, "CREATE2"),
134        }
135    }
136}
137
138impl From<revm::context_interface::CreateScheme> for CreateScheme {
139    fn from(scheme: revm::context_interface::CreateScheme) -> Self {
140        match scheme {
141            revm::context_interface::CreateScheme::Create => Self::Create,
142            revm::context_interface::CreateScheme::Create2 { .. } => Self::Create2,
143            _ => unimplemented!("Unsupported create scheme"),
144        }
145    }
146}
147
148impl CreateScheme {
149    pub fn eq(&self, create_scheme: Self) -> bool {
150        matches!(
151            (self, create_scheme),
152            (Self::Create, Self::Create) | (Self::Create2, Self::Create2 { .. })
153        )
154    }
155}
156
157impl Cheatcode for expectCall_0Call {
158    fn apply(&self, state: &mut Cheatcodes) -> Result {
159        let Self { callee, data } = self;
160        expect_call(state, callee, data, None, None, None, 1, ExpectedCallType::NonCount)
161    }
162}
163
164impl Cheatcode for expectCall_1Call {
165    fn apply(&self, state: &mut Cheatcodes) -> Result {
166        let Self { callee, data, count } = self;
167        expect_call(state, callee, data, None, None, None, *count, ExpectedCallType::Count)
168    }
169}
170
171impl Cheatcode for expectCall_2Call {
172    fn apply(&self, state: &mut Cheatcodes) -> Result {
173        let Self { callee, msgValue, data } = self;
174        expect_call(state, callee, data, Some(msgValue), None, None, 1, ExpectedCallType::NonCount)
175    }
176}
177
178impl Cheatcode for expectCall_3Call {
179    fn apply(&self, state: &mut Cheatcodes) -> Result {
180        let Self { callee, msgValue, data, count } = self;
181        expect_call(
182            state,
183            callee,
184            data,
185            Some(msgValue),
186            None,
187            None,
188            *count,
189            ExpectedCallType::Count,
190        )
191    }
192}
193
194impl Cheatcode for expectCall_4Call {
195    fn apply(&self, state: &mut Cheatcodes) -> Result {
196        let Self { callee, msgValue, gas, data } = self;
197        expect_call(
198            state,
199            callee,
200            data,
201            Some(msgValue),
202            Some(*gas),
203            None,
204            1,
205            ExpectedCallType::NonCount,
206        )
207    }
208}
209
210impl Cheatcode for expectCall_5Call {
211    fn apply(&self, state: &mut Cheatcodes) -> Result {
212        let Self { callee, msgValue, gas, data, count } = self;
213        expect_call(
214            state,
215            callee,
216            data,
217            Some(msgValue),
218            Some(*gas),
219            None,
220            *count,
221            ExpectedCallType::Count,
222        )
223    }
224}
225
226impl Cheatcode for expectCallMinGas_0Call {
227    fn apply(&self, state: &mut Cheatcodes) -> Result {
228        let Self { callee, msgValue, minGas, data } = self;
229        expect_call(
230            state,
231            callee,
232            data,
233            Some(msgValue),
234            None,
235            Some(*minGas),
236            1,
237            ExpectedCallType::NonCount,
238        )
239    }
240}
241
242impl Cheatcode for expectCallMinGas_1Call {
243    fn apply(&self, state: &mut Cheatcodes) -> Result {
244        let Self { callee, msgValue, minGas, data, count } = self;
245        expect_call(
246            state,
247            callee,
248            data,
249            Some(msgValue),
250            None,
251            Some(*minGas),
252            *count,
253            ExpectedCallType::Count,
254        )
255    }
256}
257
258impl Cheatcode for expectEmit_0Call {
259    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
260        let Self { checkTopic1, checkTopic2, checkTopic3, checkData } = *self;
261        expect_emit(
262            ccx.state,
263            ccx.ecx.journaled_state.depth(),
264            [true, checkTopic1, checkTopic2, checkTopic3, checkData],
265            None,
266            false,
267            1,
268        )
269    }
270}
271
272impl Cheatcode for expectEmit_1Call {
273    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
274        let Self { checkTopic1, checkTopic2, checkTopic3, checkData, emitter } = *self;
275        expect_emit(
276            ccx.state,
277            ccx.ecx.journaled_state.depth(),
278            [true, checkTopic1, checkTopic2, checkTopic3, checkData],
279            Some(emitter),
280            false,
281            1,
282        )
283    }
284}
285
286impl Cheatcode for expectEmit_2Call {
287    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
288        let Self {} = self;
289        expect_emit(ccx.state, ccx.ecx.journaled_state.depth(), [true; 5], None, false, 1)
290    }
291}
292
293impl Cheatcode for expectEmit_3Call {
294    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
295        let Self { emitter } = *self;
296        expect_emit(ccx.state, ccx.ecx.journaled_state.depth(), [true; 5], Some(emitter), false, 1)
297    }
298}
299
300impl Cheatcode for expectEmit_4Call {
301    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
302        let Self { checkTopic1, checkTopic2, checkTopic3, checkData, count } = *self;
303        expect_emit(
304            ccx.state,
305            ccx.ecx.journaled_state.depth(),
306            [true, checkTopic1, checkTopic2, checkTopic3, checkData],
307            None,
308            false,
309            count,
310        )
311    }
312}
313
314impl Cheatcode for expectEmit_5Call {
315    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
316        let Self { checkTopic1, checkTopic2, checkTopic3, checkData, emitter, count } = *self;
317        expect_emit(
318            ccx.state,
319            ccx.ecx.journaled_state.depth(),
320            [true, checkTopic1, checkTopic2, checkTopic3, checkData],
321            Some(emitter),
322            false,
323            count,
324        )
325    }
326}
327
328impl Cheatcode for expectEmit_6Call {
329    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
330        let Self { count } = *self;
331        expect_emit(ccx.state, ccx.ecx.journaled_state.depth(), [true; 5], None, false, count)
332    }
333}
334
335impl Cheatcode for expectEmit_7Call {
336    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
337        let Self { emitter, count } = *self;
338        expect_emit(
339            ccx.state,
340            ccx.ecx.journaled_state.depth(),
341            [true; 5],
342            Some(emitter),
343            false,
344            count,
345        )
346    }
347}
348
349impl Cheatcode for expectEmitAnonymous_0Call {
350    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
351        let Self { checkTopic0, checkTopic1, checkTopic2, checkTopic3, checkData } = *self;
352        expect_emit(
353            ccx.state,
354            ccx.ecx.journaled_state.depth(),
355            [checkTopic0, checkTopic1, checkTopic2, checkTopic3, checkData],
356            None,
357            true,
358            1,
359        )
360    }
361}
362
363impl Cheatcode for expectEmitAnonymous_1Call {
364    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
365        let Self { checkTopic0, checkTopic1, checkTopic2, checkTopic3, checkData, emitter } = *self;
366        expect_emit(
367            ccx.state,
368            ccx.ecx.journaled_state.depth(),
369            [checkTopic0, checkTopic1, checkTopic2, checkTopic3, checkData],
370            Some(emitter),
371            true,
372            1,
373        )
374    }
375}
376
377impl Cheatcode for expectEmitAnonymous_2Call {
378    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
379        let Self {} = self;
380        expect_emit(ccx.state, ccx.ecx.journaled_state.depth(), [true; 5], None, true, 1)
381    }
382}
383
384impl Cheatcode for expectEmitAnonymous_3Call {
385    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
386        let Self { emitter } = *self;
387        expect_emit(ccx.state, ccx.ecx.journaled_state.depth(), [true; 5], Some(emitter), true, 1)
388    }
389}
390
391impl Cheatcode for expectCreateCall {
392    fn apply(&self, state: &mut Cheatcodes) -> Result {
393        let Self { bytecode, deployer } = self;
394        expect_create(state, bytecode.clone(), *deployer, CreateScheme::Create)
395    }
396}
397
398impl Cheatcode for expectCreate2Call {
399    fn apply(&self, state: &mut Cheatcodes) -> Result {
400        let Self { bytecode, deployer } = self;
401        expect_create(state, bytecode.clone(), *deployer, CreateScheme::Create2)
402    }
403}
404
405impl Cheatcode for expectRevert_0Call {
406    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
407        let Self {} = self;
408        expect_revert(ccx.state, None, ccx.ecx.journaled_state.depth(), false, false, None, 1)
409    }
410}
411
412impl Cheatcode for expectRevert_1Call {
413    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
414        let Self { revertData } = self;
415        expect_revert(
416            ccx.state,
417            Some(revertData.as_ref()),
418            ccx.ecx.journaled_state.depth(),
419            false,
420            false,
421            None,
422            1,
423        )
424    }
425}
426
427impl Cheatcode for expectRevert_2Call {
428    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
429        let Self { revertData } = self;
430        expect_revert(
431            ccx.state,
432            Some(revertData),
433            ccx.ecx.journaled_state.depth(),
434            false,
435            false,
436            None,
437            1,
438        )
439    }
440}
441
442impl Cheatcode for expectRevert_3Call {
443    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
444        let Self { reverter } = self;
445        expect_revert(
446            ccx.state,
447            None,
448            ccx.ecx.journaled_state.depth(),
449            false,
450            false,
451            Some(*reverter),
452            1,
453        )
454    }
455}
456
457impl Cheatcode for expectRevert_4Call {
458    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
459        let Self { revertData, reverter } = self;
460        expect_revert(
461            ccx.state,
462            Some(revertData.as_ref()),
463            ccx.ecx.journaled_state.depth(),
464            false,
465            false,
466            Some(*reverter),
467            1,
468        )
469    }
470}
471
472impl Cheatcode for expectRevert_5Call {
473    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
474        let Self { revertData, reverter } = self;
475        expect_revert(
476            ccx.state,
477            Some(revertData),
478            ccx.ecx.journaled_state.depth(),
479            false,
480            false,
481            Some(*reverter),
482            1,
483        )
484    }
485}
486
487impl Cheatcode for expectRevert_6Call {
488    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
489        let Self { count } = self;
490        expect_revert(ccx.state, None, ccx.ecx.journaled_state.depth(), false, false, None, *count)
491    }
492}
493
494impl Cheatcode for expectRevert_7Call {
495    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
496        let Self { revertData, count } = self;
497        expect_revert(
498            ccx.state,
499            Some(revertData.as_ref()),
500            ccx.ecx.journaled_state.depth(),
501            false,
502            false,
503            None,
504            *count,
505        )
506    }
507}
508
509impl Cheatcode for expectRevert_8Call {
510    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
511        let Self { revertData, count } = self;
512        expect_revert(
513            ccx.state,
514            Some(revertData),
515            ccx.ecx.journaled_state.depth(),
516            false,
517            false,
518            None,
519            *count,
520        )
521    }
522}
523
524impl Cheatcode for expectRevert_9Call {
525    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
526        let Self { reverter, count } = self;
527        expect_revert(
528            ccx.state,
529            None,
530            ccx.ecx.journaled_state.depth(),
531            false,
532            false,
533            Some(*reverter),
534            *count,
535        )
536    }
537}
538
539impl Cheatcode for expectRevert_10Call {
540    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
541        let Self { revertData, reverter, count } = self;
542        expect_revert(
543            ccx.state,
544            Some(revertData.as_ref()),
545            ccx.ecx.journaled_state.depth(),
546            false,
547            false,
548            Some(*reverter),
549            *count,
550        )
551    }
552}
553
554impl Cheatcode for expectRevert_11Call {
555    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
556        let Self { revertData, reverter, count } = self;
557        expect_revert(
558            ccx.state,
559            Some(revertData),
560            ccx.ecx.journaled_state.depth(),
561            false,
562            false,
563            Some(*reverter),
564            *count,
565        )
566    }
567}
568
569impl Cheatcode for expectPartialRevert_0Call {
570    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
571        let Self { revertData } = self;
572        expect_revert(
573            ccx.state,
574            Some(revertData.as_ref()),
575            ccx.ecx.journaled_state.depth(),
576            false,
577            true,
578            None,
579            1,
580        )
581    }
582}
583
584impl Cheatcode for expectPartialRevert_1Call {
585    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
586        let Self { revertData, reverter } = self;
587        expect_revert(
588            ccx.state,
589            Some(revertData.as_ref()),
590            ccx.ecx.journaled_state.depth(),
591            false,
592            true,
593            Some(*reverter),
594            1,
595        )
596    }
597}
598
599impl Cheatcode for _expectCheatcodeRevert_0Call {
600    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
601        expect_revert(ccx.state, None, ccx.ecx.journaled_state.depth(), true, false, None, 1)
602    }
603}
604
605impl Cheatcode for _expectCheatcodeRevert_1Call {
606    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
607        let Self { revertData } = self;
608        expect_revert(
609            ccx.state,
610            Some(revertData.as_ref()),
611            ccx.ecx.journaled_state.depth(),
612            true,
613            false,
614            None,
615            1,
616        )
617    }
618}
619
620impl Cheatcode for _expectCheatcodeRevert_2Call {
621    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
622        let Self { revertData } = self;
623        expect_revert(
624            ccx.state,
625            Some(revertData),
626            ccx.ecx.journaled_state.depth(),
627            true,
628            false,
629            None,
630            1,
631        )
632    }
633}
634
635impl Cheatcode for expectSafeMemoryCall {
636    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
637        let Self { min, max } = *self;
638        expect_safe_memory(ccx.state, min, max, ccx.ecx.journaled_state.depth().try_into()?)
639    }
640}
641
642impl Cheatcode for stopExpectSafeMemoryCall {
643    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
644        let Self {} = self;
645        ccx.state.allowed_mem_writes.remove(&ccx.ecx.journaled_state.depth().try_into()?);
646        Ok(Default::default())
647    }
648}
649
650impl Cheatcode for expectSafeMemoryCallCall {
651    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
652        let Self { min, max } = *self;
653        expect_safe_memory(ccx.state, min, max, (ccx.ecx.journaled_state.depth() + 1).try_into()?)
654    }
655}
656
657impl RevertParameters for ExpectedRevert {
658    fn reverter(&self) -> Option<Address> {
659        self.reverter
660    }
661
662    fn reason(&self) -> Option<&[u8]> {
663        self.reason.as_deref()
664    }
665
666    fn partial_match(&self) -> bool {
667        self.partial_match
668    }
669}
670
671/// Handles expected calls specified by the `expectCall` cheatcodes.
672///
673/// It can handle calls in two ways:
674/// - If the cheatcode was used with a `count` argument, it will expect the call to be made exactly
675///   `count` times. e.g. `vm.expectCall(address(0xc4f3), abi.encodeWithSelector(0xd34db33f), 4)`
676///   will expect the call to address(0xc4f3) with selector `0xd34db33f` to be made exactly 4 times.
677///   If the amount of calls is less or more than 4, the test will fail. Note that the `count`
678///   argument cannot be overwritten with another `vm.expectCall`. If this is attempted,
679///   `expectCall` will revert.
680/// - If the cheatcode was used without a `count` argument, it will expect the call to be made at
681///   least the amount of times the cheatcode was called. This means that `vm.expectCall` without a
682///   count argument can be called many times, but cannot be called with a `count` argument after it
683///   was called without one. If the latter happens, `expectCall` will revert. e.g
684///   `vm.expectCall(address(0xc4f3), abi.encodeWithSelector(0xd34db33f))` will expect the call to
685///   address(0xc4f3) and selector `0xd34db33f` to be made at least once. If the amount of calls is
686///   0, the test will fail. If the call is made more than once, the test will pass.
687#[expect(clippy::too_many_arguments)] // It is what it is
688fn expect_call(
689    state: &mut Cheatcodes,
690    target: &Address,
691    calldata: &Bytes,
692    value: Option<&U256>,
693    mut gas: Option<u64>,
694    mut min_gas: Option<u64>,
695    count: u64,
696    call_type: ExpectedCallType,
697) -> Result {
698    let expecteds = state.expected_calls.entry(*target).or_default();
699
700    if let Some(val) = value {
701        if *val > U256::ZERO {
702            // If the value of the transaction is non-zero, the EVM adds a call stipend of 2300 gas
703            // to ensure that the basic fallback function can be called.
704            let positive_value_cost_stipend = 2300;
705            if let Some(gas) = &mut gas {
706                *gas += positive_value_cost_stipend;
707            }
708            if let Some(min_gas) = &mut min_gas {
709                *min_gas += positive_value_cost_stipend;
710            }
711        }
712    }
713
714    match call_type {
715        ExpectedCallType::Count => {
716            // Get the expected calls for this target.
717            // In this case, as we're using counted expectCalls, we should not be able to set them
718            // more than once.
719            ensure!(
720                !expecteds.contains_key(calldata),
721                "counted expected calls can only bet set once"
722            );
723            expecteds.insert(
724                calldata.clone(),
725                (ExpectedCallData { value: value.copied(), gas, min_gas, count, call_type }, 0),
726            );
727        }
728        ExpectedCallType::NonCount => {
729            // Check if the expected calldata exists.
730            // If it does, increment the count by one as we expect to see it one more time.
731            match expecteds.entry(calldata.clone()) {
732                Entry::Occupied(mut entry) => {
733                    let (expected, _) = entry.get_mut();
734                    // Ensure we're not overwriting a counted expectCall.
735                    ensure!(
736                        expected.call_type == ExpectedCallType::NonCount,
737                        "cannot overwrite a counted expectCall with a non-counted expectCall"
738                    );
739                    expected.count += 1;
740                }
741                // If it does not exist, then create it.
742                Entry::Vacant(entry) => {
743                    entry.insert((
744                        ExpectedCallData { value: value.copied(), gas, min_gas, count, call_type },
745                        0,
746                    ));
747                }
748            }
749        }
750    }
751
752    Ok(Default::default())
753}
754
755fn expect_emit(
756    state: &mut Cheatcodes,
757    depth: usize,
758    checks: [bool; 5],
759    address: Option<Address>,
760    anonymous: bool,
761    count: u64,
762) -> Result {
763    let expected_emit =
764        ExpectedEmit { depth, checks, address, found: false, log: None, anonymous, count };
765    if let Some(found_emit_pos) = state.expected_emits.iter().position(|(emit, _)| emit.found) {
766        // The order of emits already found (back of queue) should not be modified, hence push any
767        // new emit before first found emit.
768        state.expected_emits.insert(found_emit_pos, (expected_emit, Default::default()));
769    } else {
770        // If no expected emits then push new one at the back of queue.
771        state.expected_emits.push_back((expected_emit, Default::default()));
772    }
773
774    Ok(Default::default())
775}
776
777pub(crate) fn handle_expect_emit(
778    state: &mut Cheatcodes,
779    log: &alloy_primitives::Log,
780    interpreter: &mut Interpreter,
781) {
782    // Fill or check the expected emits.
783    // We expect for emit checks to be filled as they're declared (from oldest to newest),
784    // so we fill them and push them to the back of the queue.
785    // If the user has properly filled all the emits, they'll end up in their original order.
786    // If not, the queue will not be in the order the events will be intended to be filled,
787    // and we'll be able to later detect this and bail.
788
789    // First, we can return early if all events have been matched.
790    // This allows a contract to arbitrarily emit more events than expected (additive behavior),
791    // as long as all the previous events were matched in the order they were expected to be.
792    if state.expected_emits.iter().all(|(expected, _)| expected.found) {
793        return
794    }
795
796    let should_fill_logs = state.expected_emits.iter().any(|(expected, _)| expected.log.is_none());
797    let index_to_fill_or_check = if should_fill_logs {
798        // If there's anything to fill, we start with the last event to match in the queue
799        // (without taking into account events already matched).
800        state
801            .expected_emits
802            .iter()
803            .position(|(emit, _)| emit.found)
804            .unwrap_or(state.expected_emits.len())
805            .saturating_sub(1)
806    } else {
807        // Otherwise, if all expected logs are filled, we start to check any unmatched event
808        // in the declared order, so we start from the front (like a queue).
809        0
810    };
811
812    let (mut event_to_fill_or_check, mut count_map) = state
813        .expected_emits
814        .remove(index_to_fill_or_check)
815        .expect("we should have an emit to fill or check");
816
817    let Some(expected) = &event_to_fill_or_check.log else {
818        // Unless the caller is trying to match an anonymous event, the first topic must be
819        // filled.
820        if event_to_fill_or_check.anonymous || !log.topics().is_empty() {
821            event_to_fill_or_check.log = Some(log.data.clone());
822            // If we only filled the expected log then we put it back at the same position.
823            state
824                .expected_emits
825                .insert(index_to_fill_or_check, (event_to_fill_or_check, count_map));
826        } else {
827            interpreter.control.instruction_result = InstructionResult::Revert;
828            interpreter.control.next_action = InterpreterAction::Return {
829                result: InterpreterResult {
830                    output: Error::encode("use vm.expectEmitAnonymous to match anonymous events"),
831                    gas: interpreter.control.gas,
832                    result: InstructionResult::Revert,
833                },
834            };
835        }
836        return
837    };
838
839    // Increment/set `count` for `log.address` and `log.data`
840    match count_map.entry(log.address) {
841        Entry::Occupied(mut entry) => {
842            // Checks and inserts the log into the map.
843            // If the log doesn't pass the checks, it is ignored and `count` is not incremented.
844            let log_count_map = entry.get_mut();
845            log_count_map.insert(&log.data);
846        }
847        Entry::Vacant(entry) => {
848            let mut log_count_map = LogCountMap::new(&event_to_fill_or_check);
849
850            if log_count_map.satisfies_checks(&log.data) {
851                log_count_map.insert(&log.data);
852
853                // Entry is only inserted if it satisfies the checks.
854                entry.insert(log_count_map);
855            }
856        }
857    }
858
859    event_to_fill_or_check.found = || -> bool {
860        if !checks_topics_and_data(event_to_fill_or_check.checks, expected, log) {
861            return false
862        }
863
864        // Maybe match source address.
865        if event_to_fill_or_check.address.is_some_and(|addr| addr != log.address) {
866            return false;
867        }
868
869        let expected_count = event_to_fill_or_check.count;
870
871        match event_to_fill_or_check.address {
872            Some(emitter) => count_map
873                .get(&emitter)
874                .is_some_and(|log_map| log_map.count(&log.data) >= expected_count),
875            None => count_map
876                .values()
877                .find(|log_map| log_map.satisfies_checks(&log.data))
878                .is_some_and(|map| map.count(&log.data) >= expected_count),
879        }
880    }();
881
882    // If we found the event, we can push it to the back of the queue
883    // and begin expecting the next event.
884    if event_to_fill_or_check.found {
885        state.expected_emits.push_back((event_to_fill_or_check, count_map));
886    } else {
887        // We did not match this event, so we need to keep waiting for the right one to
888        // appear.
889        state.expected_emits.push_front((event_to_fill_or_check, count_map));
890    }
891}
892
893/// Handles expected emits specified by the `expectEmit` cheatcodes.
894///
895/// The second element of the tuple counts the number of times the log has been emitted by a
896/// particular address
897pub type ExpectedEmitTracker = VecDeque<(ExpectedEmit, AddressHashMap<LogCountMap>)>;
898
899#[derive(Clone, Debug, Default)]
900pub struct LogCountMap {
901    checks: [bool; 5],
902    expected_log: RawLog,
903    map: HashMap<RawLog, u64>,
904}
905
906impl LogCountMap {
907    /// Instantiates `LogCountMap`.
908    fn new(expected_emit: &ExpectedEmit) -> Self {
909        Self {
910            checks: expected_emit.checks,
911            expected_log: expected_emit.log.clone().expect("log should be filled here"),
912            map: Default::default(),
913        }
914    }
915
916    /// Inserts a log into the map and increments the count.
917    ///
918    /// The log must pass all checks against the expected log for the count to increment.
919    ///
920    /// Returns true if the log was inserted and count was incremented.
921    fn insert(&mut self, log: &RawLog) -> bool {
922        // If its already in the map, increment the count without checking.
923        if self.map.contains_key(log) {
924            self.map.entry(log.clone()).and_modify(|c| *c += 1);
925
926            return true
927        }
928
929        if !self.satisfies_checks(log) {
930            return false
931        }
932
933        self.map.entry(log.clone()).and_modify(|c| *c += 1).or_insert(1);
934
935        true
936    }
937
938    /// Checks the incoming raw log against the expected logs topics and data.
939    fn satisfies_checks(&self, log: &RawLog) -> bool {
940        checks_topics_and_data(self.checks, &self.expected_log, log)
941    }
942
943    pub fn count(&self, log: &RawLog) -> u64 {
944        if !self.satisfies_checks(log) {
945            return 0
946        }
947
948        self.count_unchecked()
949    }
950
951    pub fn count_unchecked(&self) -> u64 {
952        self.map.values().sum()
953    }
954}
955
956fn expect_create(
957    state: &mut Cheatcodes,
958    bytecode: Bytes,
959    deployer: Address,
960    create_scheme: CreateScheme,
961) -> Result {
962    let expected_create = ExpectedCreate { bytecode, deployer, create_scheme };
963    state.expected_creates.push(expected_create);
964
965    Ok(Default::default())
966}
967
968fn expect_revert(
969    state: &mut Cheatcodes,
970    reason: Option<&[u8]>,
971    depth: usize,
972    cheatcode: bool,
973    partial_match: bool,
974    reverter: Option<Address>,
975    count: u64,
976) -> Result {
977    ensure!(
978        state.expected_revert.is_none(),
979        "you must call another function prior to expecting a second revert"
980    );
981    state.expected_revert = Some(ExpectedRevert {
982        reason: reason.map(<[_]>::to_vec),
983        depth,
984        kind: if cheatcode {
985            ExpectedRevertKind::Cheatcode { pending_processing: true }
986        } else {
987            ExpectedRevertKind::Default
988        },
989        partial_match,
990        reverter,
991        reverted_by: None,
992        max_depth: depth,
993        count,
994        actual_count: 0,
995    });
996    Ok(Default::default())
997}
998
999fn checks_topics_and_data(checks: [bool; 5], expected: &RawLog, log: &RawLog) -> bool {
1000    if log.topics().len() != expected.topics().len() {
1001        return false
1002    }
1003
1004    // Check topics.
1005    if !log
1006        .topics()
1007        .iter()
1008        .enumerate()
1009        .filter(|(i, _)| checks[*i])
1010        .all(|(i, topic)| topic == &expected.topics()[i])
1011    {
1012        return false
1013    }
1014
1015    // Check data
1016    if checks[4] && expected.data.as_ref() != log.data.as_ref() {
1017        return false
1018    }
1019
1020    true
1021}
1022
1023fn expect_safe_memory(state: &mut Cheatcodes, start: u64, end: u64, depth: u64) -> Result {
1024    ensure!(start < end, "memory range start ({start}) is greater than end ({end})");
1025    #[expect(clippy::single_range_in_vec_init)] // Wanted behaviour
1026    let offsets = state.allowed_mem_writes.entry(depth).or_insert_with(|| vec![0..0x60]);
1027    offsets.push(start..end);
1028    Ok(Default::default())
1029}
foundry_cheatcodes/test/expect.rs

foundry_cheatcodes/test/
expect.rs
