Skip to main content

foundry_cheatcodes/
monad.rs

1//! Monad staking cheatcodes.
2
3use 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    /// Monad-specific cheatcodes. Accessible via `MonadVm(MONAD_CHEATCODE_ADDRESS)`.
32    interface MonadVm {
33        /// Sets the current epoch and delay period for the staking precompile.
34        function setEpoch(uint64 epoch, bool inDelayPeriod) external;
35
36        /// Sets the current block proposer validator ID.
37        function setProposer(uint64 valId) external;
38
39        /// Directly sets a validator's accumulated reward per token.
40        function setAccumulator(uint64 valId, uint256 value) external;
41
42        /// Distribute block reward via the real syscallReward handler.
43        function blockReward(address author, uint256 reward) external;
44
45        /// Execute syscallSnapshot.
46        function epochSnapshot() external;
47
48        /// Execute syscallOnEpochChange.
49        function epochChange(uint64 newEpoch) external;
50
51        /// Convenience: epochSnapshot() then epochChange(newEpoch).
52        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}