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

foundry_cheatcodes/test/
expect.rs
