foundry_cheatcodes/evm/
mapping.rs

1use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
2use alloy_primitives::{Address, B256, U256, keccak256, map::AddressHashMap};
3use alloy_sol_types::SolValue;
4use foundry_common::mapping_slots::MappingSlots;
5use revm::{
6    bytecode::opcode,
7    interpreter::{Interpreter, interpreter_types::Jumps},
8};
9
10impl Cheatcode for startMappingRecordingCall {
11    fn apply(&self, state: &mut Cheatcodes) -> Result {
12        let Self {} = self;
13        state.mapping_slots.get_or_insert_default();
14        Ok(Default::default())
15    }
16}
17
18impl Cheatcode for stopMappingRecordingCall {
19    fn apply(&self, state: &mut Cheatcodes) -> Result {
20        let Self {} = self;
21        state.mapping_slots = None;
22        Ok(Default::default())
23    }
24}
25
26impl Cheatcode for getMappingLengthCall {
27    fn apply(&self, state: &mut Cheatcodes) -> Result {
28        let Self { target, mappingSlot } = self;
29        let result = slot_child(state, target, mappingSlot).map(Vec::len).unwrap_or(0);
30        Ok((result as u64).abi_encode())
31    }
32}
33
34impl Cheatcode for getMappingSlotAtCall {
35    fn apply(&self, state: &mut Cheatcodes) -> Result {
36        let Self { target, mappingSlot, idx } = self;
37        let result = slot_child(state, target, mappingSlot)
38            .and_then(|set| set.get(idx.saturating_to::<usize>()))
39            .copied()
40            .unwrap_or_default();
41        Ok(result.abi_encode())
42    }
43}
44
45impl Cheatcode for getMappingKeyAndParentOfCall {
46    fn apply(&self, state: &mut Cheatcodes) -> Result {
47        let Self { target, elementSlot: slot } = self;
48        let mut found = false;
49        let mut key = &B256::ZERO;
50        let mut parent = &B256::ZERO;
51        if let Some(slots) = mapping_slot(state, target) {
52            if let Some(key2) = slots.keys.get(slot) {
53                found = true;
54                key = key2;
55                parent = &slots.parent_slots[slot];
56            } else if let Some((key2, parent2)) = slots.seen_sha3.get(slot) {
57                found = true;
58                key = key2;
59                parent = parent2;
60            }
61        }
62        Ok((found, key, parent).abi_encode_params())
63    }
64}
65
66fn mapping_slot<'a>(state: &'a Cheatcodes, target: &'a Address) -> Option<&'a MappingSlots> {
67    state.mapping_slots.as_ref()?.get(target)
68}
69
70fn slot_child<'a>(
71    state: &'a Cheatcodes,
72    target: &'a Address,
73    slot: &'a B256,
74) -> Option<&'a Vec<B256>> {
75    mapping_slot(state, target)?.children.get(slot)
76}
77
78#[cold]
79pub(crate) fn step(mapping_slots: &mut AddressHashMap<MappingSlots>, interpreter: &Interpreter) {
80    match interpreter.bytecode.opcode() {
81        opcode::KECCAK256 => {
82            if interpreter.stack.peek(1) == Ok(U256::from(0x40)) {
83                let address = interpreter.input.target_address;
84                let offset = interpreter.stack.peek(0).expect("stack size > 1").saturating_to();
85                let data = interpreter.memory.slice_len(offset, 0x40);
86                let low = B256::from_slice(&data[..0x20]);
87                let high = B256::from_slice(&data[0x20..]);
88                let result = keccak256(&*data);
89
90                mapping_slots.entry(address).or_default().seen_sha3.insert(result, (low, high));
91            }
92        }
93        opcode::SSTORE => {
94            if let Some(mapping_slots) = mapping_slots.get_mut(&interpreter.input.target_address)
95                && let Ok(slot) = interpreter.stack.peek(0)
96            {
97                mapping_slots.insert(slot.into());
98            }
99        }
100        _ => {}
101    }
102}
foundry_cheatcodes/evm/mapping.rs

foundry_cheatcodes/evm/
mapping.rs
