foundry_cheatcodes/evm/
mapping.rs

1use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
2use alloy_primitives::{Address, B256};
3use alloy_sol_types::SolValue;
4use foundry_common::mapping_slots::MappingSlots;
5use foundry_evm_core::evm::FoundryEvmNetwork;
6
7impl Cheatcode for startMappingRecordingCall {
8    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
9        let Self {} = self;
10        state.mapping_slots.get_or_insert_default();
11        Ok(Default::default())
12    }
13}
14
15impl Cheatcode for stopMappingRecordingCall {
16    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
17        let Self {} = self;
18        state.mapping_slots = None;
19        Ok(Default::default())
20    }
21}
22
23impl Cheatcode for getMappingLengthCall {
24    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
25        let Self { target, mappingSlot } = self;
26        let result = slot_child(state, target, mappingSlot).map(Vec::len).unwrap_or(0);
27        Ok((result as u64).abi_encode())
28    }
29}
30
31impl Cheatcode for getMappingSlotAtCall {
32    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
33        let Self { target, mappingSlot, idx } = self;
34        let result = slot_child(state, target, mappingSlot)
35            .and_then(|set| set.get(idx.saturating_to::<usize>()))
36            .copied()
37            .unwrap_or_default();
38        Ok(result.abi_encode())
39    }
40}
41
42impl Cheatcode for getMappingKeyAndParentOfCall {
43    fn apply<FEN: FoundryEvmNetwork>(&self, state: &mut Cheatcodes<FEN>) -> Result {
44        let Self { target, elementSlot: slot } = self;
45        let mut found = false;
46        let mut key = &B256::ZERO;
47        let mut parent = &B256::ZERO;
48        if let Some(slots) = mapping_slot(state, target) {
49            if let Some(key2) = slots.keys.get(slot) {
50                found = true;
51                key = key2;
52                parent = &slots.parent_slots[slot];
53            } else if let Some((key2, parent2)) = slots.seen_sha3.get(slot) {
54                found = true;
55                key = key2;
56                parent = parent2;
57            }
58        }
59        Ok((found, key, parent).abi_encode_params())
60    }
61}
62
63fn mapping_slot<'a, FEN: FoundryEvmNetwork>(
64    state: &'a Cheatcodes<FEN>,
65    target: &'a Address,
66) -> Option<&'a MappingSlots> {
67    state.mapping_slots.as_ref()?.get(target)
68}
69
70fn slot_child<'a, FEN: FoundryEvmNetwork>(
71    state: &'a Cheatcodes<FEN>,
72    target: &'a Address,
73    slot: &'a B256,
74) -> Option<&'a Vec<B256>> {
75    mapping_slot(state, target)?.children.get(slot)
76}
foundry_cheatcodes/evm/mapping.rs

foundry_cheatcodes/evm/
mapping.rs
