1use crate::{CheatsCtxt, Result};
4use alloy_primitives::{Address, U256};
5use alloy_sol_types::SolInterface;
6use foundry_evm_core::{
7 constants::MONAD_CHEATCODE_ADDRESS,
8 evm::{FoundryEvmFactory, FoundryEvmNetwork},
9};
10use monad_revm::{
11 api::block::{
12 syscall_on_epoch_change_calldata, syscall_reward_calldata, syscall_snapshot_calldata,
13 },
14 staking::{
15 StorageReader,
16 constants::SYSTEM_ADDRESS,
17 storage::{STAKING_ADDRESS, global_slots, validator_key, validator_offsets},
18 write::{
19 StakingStorage, handle_syscall_on_epoch_change, handle_syscall_reward,
20 handle_syscall_snapshot,
21 },
22 },
23};
24use revm::{
25 context::{ContextTr, JournalTr},
26 precompile::PrecompileHalt,
27 primitives::Log,
28};
29
30alloy_sol_types::sol! {
31 interface MonadVm {
33 function setEpoch(uint64 epoch, bool inDelayPeriod) external;
35
36 function setProposer(uint64 valId) external;
38
39 function setAccumulator(uint64 valId, uint256 value) external;
41
42 function blockReward(address author, uint256 reward) external;
44
45 function epochSnapshot() external;
47
48 function epochChange(uint64 newEpoch) external;
50
51 function epochBoundary(uint64 newEpoch) external;
53 }
54}
55
56pub(crate) fn is_monad_cheatcode_call<FEN: FoundryEvmNetwork>(target: Address) -> bool {
57 target == MONAD_CHEATCODE_ADDRESS
58 && FEN::EvmFactory::EXTRA_CHEATCODE_ADDRESSES.contains(&target)
59}
60
61pub(crate) fn apply_monad_cheatcode<FEN: FoundryEvmNetwork>(
62 ccx: &mut CheatsCtxt<'_, '_, FEN>,
63 input: &[u8],
64) -> Result {
65 let decoded = MonadVm::MonadVmCalls::abi_decode(input).map_err(|e| {
66 if let alloy_sol_types::Error::UnknownSelector { selector, .. } = e {
67 let msg = format!(
68 "unknown monad cheatcode with selector {selector}; \
69 check that your MonadVm interface matches this forge version"
70 );
71 return alloy_sol_types::Error::Other(std::borrow::Cow::Owned(msg));
72 }
73 e
74 })?;
75
76 match decoded {
77 MonadVm::MonadVmCalls::setEpoch(call) => apply_set_epoch(ccx, call),
78 MonadVm::MonadVmCalls::setProposer(call) => apply_set_proposer(ccx, call),
79 MonadVm::MonadVmCalls::setAccumulator(call) => apply_set_accumulator(ccx, call),
80 MonadVm::MonadVmCalls::blockReward(call) => apply_block_reward(ccx, call),
81 MonadVm::MonadVmCalls::epochSnapshot(call) => apply_epoch_snapshot(ccx, call),
82 MonadVm::MonadVmCalls::epochChange(call) => apply_epoch_change(ccx, call),
83 MonadVm::MonadVmCalls::epochBoundary(call) => apply_epoch_boundary(ccx, call),
84 }
85}
86
87fn u64_left_aligned(v: u64) -> U256 {
88 let mut bytes = [0u8; 32];
89 bytes[0..8].copy_from_slice(&v.to_be_bytes());
90 U256::from_be_bytes(bytes)
91}
92
93fn sstore_staking<FEN: FoundryEvmNetwork>(
94 ccx: &mut CheatsCtxt<'_, '_, FEN>,
95 key: U256,
96 value: U256,
97) -> Result<()> {
98 ccx.ecx.journal_mut().load_account(STAKING_ADDRESS)?;
99 ccx.ecx
100 .journal_mut()
101 .sstore(STAKING_ADDRESS, key, value)
102 .map_err(|e| fmt_err!("staking sstore failed: {:?}", e))?;
103 Ok(())
104}
105
106struct CheatsCtxtStorage<'a, 'b, 'db, FEN: FoundryEvmNetwork> {
107 ccx: &'a mut CheatsCtxt<'b, 'db, FEN>,
108}
109
110impl<FEN: FoundryEvmNetwork> StorageReader for CheatsCtxtStorage<'_, '_, '_, FEN> {
111 fn sload(&mut self, key: U256) -> core::result::Result<U256, PrecompileHalt> {
112 self.ccx
113 .ecx
114 .journal_mut()
115 .sload(STAKING_ADDRESS, key)
116 .map(|r| r.data)
117 .map_err(|e| PrecompileHalt::Other(format!("sload failed: {e:?}").into()))
118 }
119}
120
121impl<FEN: FoundryEvmNetwork> StakingStorage for CheatsCtxtStorage<'_, '_, '_, FEN> {
122 fn sstore(&mut self, key: U256, value: U256) -> core::result::Result<(), PrecompileHalt> {
123 self.ccx
124 .ecx
125 .journal_mut()
126 .sstore(STAKING_ADDRESS, key, value)
127 .map(|_| ())
128 .map_err(|e| PrecompileHalt::Other(format!("sstore failed: {e:?}").into()))
129 }
130
131 fn transfer(
132 &mut self,
133 from: Address,
134 to: Address,
135 amount: U256,
136 ) -> core::result::Result<(), PrecompileHalt> {
137 if amount.is_zero() {
138 return Ok(());
139 }
140
141 match self.ccx.ecx.journal_mut().transfer(from, to, amount) {
142 Ok(None) => Ok(()),
143 Ok(Some(e)) => Err(PrecompileHalt::Other(format!("transfer failed: {e:?}").into())),
144 Err(e) => Err(PrecompileHalt::Other(format!("transfer error: {e:?}").into())),
145 }
146 }
147
148 fn emit_log(&mut self, log: Log) -> core::result::Result<(), PrecompileHalt> {
149 self.ccx.ecx.journal_mut().log(log);
150 Ok(())
151 }
152}
153
154fn apply_set_epoch<FEN: FoundryEvmNetwork>(
155 ccx: &mut CheatsCtxt<'_, '_, FEN>,
156 call: MonadVm::setEpochCall,
157) -> Result {
158 let MonadVm::setEpochCall { epoch, inDelayPeriod } = call;
159 sstore_staking(ccx, global_slots::EPOCH, u64_left_aligned(epoch))?;
160
161 let boundary_val = if inDelayPeriod {
162 let mut bytes = [0u8; 32];
163 bytes[0] = 1;
164 U256::from_be_bytes(bytes)
165 } else {
166 U256::ZERO
167 };
168 sstore_staking(ccx, global_slots::IN_BOUNDARY, boundary_val)?;
169 Ok(Default::default())
170}
171
172fn apply_set_proposer<FEN: FoundryEvmNetwork>(
173 ccx: &mut CheatsCtxt<'_, '_, FEN>,
174 call: MonadVm::setProposerCall,
175) -> Result {
176 sstore_staking(ccx, global_slots::PROPOSER_VAL_ID, u64_left_aligned(call.valId))?;
177 Ok(Default::default())
178}
179
180fn apply_set_accumulator<FEN: FoundryEvmNetwork>(
181 ccx: &mut CheatsCtxt<'_, '_, FEN>,
182 call: MonadVm::setAccumulatorCall,
183) -> Result {
184 sstore_staking(
185 ccx,
186 validator_key(call.valId, validator_offsets::ACCUMULATED_REWARD_PER_TOKEN),
187 call.value,
188 )?;
189 Ok(Default::default())
190}
191
192fn apply_block_reward<FEN: FoundryEvmNetwork>(
193 ccx: &mut CheatsCtxt<'_, '_, FEN>,
194 call: MonadVm::blockRewardCall,
195) -> Result {
196 let MonadVm::blockRewardCall { author, reward } = call;
197 let calldata = syscall_reward_calldata(author, reward);
198
199 let mut storage = CheatsCtxtStorage { ccx };
200 handle_syscall_reward(&mut storage, &calldata, u64::MAX, &SYSTEM_ADDRESS, U256::ZERO)
201 .map_err(|e| fmt_err!("blockReward failed: {e}"))?;
202
203 storage
204 .ccx
205 .ecx
206 .journal_mut()
207 .balance_incr(STAKING_ADDRESS, reward)
208 .map_err(|e| fmt_err!("blockReward balance increment failed: {e:?}"))?;
209
210 Ok(Default::default())
211}
212
213fn apply_epoch_snapshot<FEN: FoundryEvmNetwork>(
214 ccx: &mut CheatsCtxt<'_, '_, FEN>,
215 _call: MonadVm::epochSnapshotCall,
216) -> Result {
217 let calldata = syscall_snapshot_calldata();
218 let mut storage = CheatsCtxtStorage { ccx };
219 handle_syscall_snapshot(&mut storage, &calldata, u64::MAX, &SYSTEM_ADDRESS)
220 .map_err(|e| fmt_err!("epochSnapshot failed: {e}"))?;
221 Ok(Default::default())
222}
223
224fn apply_epoch_change<FEN: FoundryEvmNetwork>(
225 ccx: &mut CheatsCtxt<'_, '_, FEN>,
226 call: MonadVm::epochChangeCall,
227) -> Result {
228 let calldata = syscall_on_epoch_change_calldata(call.newEpoch);
229 let mut storage = CheatsCtxtStorage { ccx };
230 handle_syscall_on_epoch_change(&mut storage, &calldata, u64::MAX, &SYSTEM_ADDRESS)
231 .map_err(|e| fmt_err!("epochChange failed: {e}"))?;
232 Ok(Default::default())
233}
234
235fn apply_epoch_boundary<FEN: FoundryEvmNetwork>(
236 ccx: &mut CheatsCtxt<'_, '_, FEN>,
237 call: MonadVm::epochBoundaryCall,
238) -> Result {
239 apply_epoch_snapshot(ccx, MonadVm::epochSnapshotCall {})?;
240 apply_epoch_change(ccx, MonadVm::epochChangeCall { newEpoch: call.newEpoch })?;
241 Ok(Default::default())
242}