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