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