1use crate::{Cheatcode, CheatsCtxt, Result, Vm::*, evm::journaled_account};
2use alloy_primitives::Address;
3use revm::context::JournalTr;
4
5#[derive(Clone, Copy, Debug, Default)]
7pub struct Prank {
8 pub prank_caller: Address,
10 pub prank_origin: Address,
12 pub new_caller: Address,
14 pub new_origin: Option<Address>,
16 pub depth: usize,
18 pub single_call: bool,
20 pub delegate_call: bool,
22 pub used: bool,
24}
25
26impl Prank {
27 pub fn new(
29 prank_caller: Address,
30 prank_origin: Address,
31 new_caller: Address,
32 new_origin: Option<Address>,
33 depth: usize,
34 single_call: bool,
35 delegate_call: bool,
36 ) -> Self {
37 Self {
38 prank_caller,
39 prank_origin,
40 new_caller,
41 new_origin,
42 depth,
43 single_call,
44 delegate_call,
45 used: false,
46 }
47 }
48
49 pub fn first_time_applied(&self) -> Option<Self> {
52 if self.used { None } else { Some(Self { used: true, ..*self }) }
53 }
54}
55
56impl Cheatcode for prank_0Call {
57 fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
58 let Self { msgSender } = self;
59 prank(ccx, msgSender, None, true, false)
60 }
61}
62
63impl Cheatcode for startPrank_0Call {
64 fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
65 let Self { msgSender } = self;
66 prank(ccx, msgSender, None, false, false)
67 }
68}
69
70impl Cheatcode for prank_1Call {
71 fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
72 let Self { msgSender, txOrigin } = self;
73 prank(ccx, msgSender, Some(txOrigin), true, false)
74 }
75}
76
77impl Cheatcode for startPrank_1Call {
78 fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
79 let Self { msgSender, txOrigin } = self;
80 prank(ccx, msgSender, Some(txOrigin), false, false)
81 }
82}
83
84impl Cheatcode for prank_2Call {
85 fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
86 let Self { msgSender, delegateCall } = self;
87 prank(ccx, msgSender, None, true, *delegateCall)
88 }
89}
90
91impl Cheatcode for startPrank_2Call {
92 fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
93 let Self { msgSender, delegateCall } = self;
94 prank(ccx, msgSender, None, false, *delegateCall)
95 }
96}
97
98impl Cheatcode for prank_3Call {
99 fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
100 let Self { msgSender, txOrigin, delegateCall } = self;
101 prank(ccx, msgSender, Some(txOrigin), true, *delegateCall)
102 }
103}
104
105impl Cheatcode for startPrank_3Call {
106 fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
107 let Self { msgSender, txOrigin, delegateCall } = self;
108 prank(ccx, msgSender, Some(txOrigin), false, *delegateCall)
109 }
110}
111
112impl Cheatcode for stopPrankCall {
113 fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
114 let Self {} = self;
115 ccx.state.pranks.remove(&ccx.ecx.journaled_state.depth());
116 Ok(Default::default())
117 }
118}
119
120fn prank(
121 ccx: &mut CheatsCtxt,
122 new_caller: &Address,
123 new_origin: Option<&Address>,
124 single_call: bool,
125 delegate_call: bool,
126) -> Result {
127 let account = journaled_account(ccx.ecx, *new_caller)?;
131
132 if delegate_call {
134 ensure!(
135 account.info.clone().code.is_some_and(|code| !code.is_empty()),
136 "cannot `prank` delegate call from an EOA"
137 );
138 }
139
140 let depth = ccx.ecx.journaled_state.depth();
141 if let Some(Prank { used, single_call: current_single_call, .. }) = ccx.state.get_prank(depth) {
142 ensure!(used, "cannot overwrite a prank until it is applied at least once");
143 ensure!(
146 single_call == *current_single_call,
147 "cannot override an ongoing prank with a single vm.prank; \
148 use vm.startPrank to override the current prank"
149 );
150 }
151
152 let prank = Prank::new(
153 ccx.caller,
154 ccx.ecx.tx.caller,
155 *new_caller,
156 new_origin.copied(),
157 depth,
158 single_call,
159 delegate_call,
160 );
161
162 ensure!(
163 ccx.state.broadcast.is_none(),
164 "cannot `prank` for a broadcasted transaction; \
165 pass the desired `tx.origin` into the `broadcast` cheatcode call"
166 );
167
168 ccx.state.pranks.insert(prank.depth, prank);
169 Ok(Default::default())
170}