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

foundry_cheatcodes/evm/
mapping.rs
