feat(invariants): use SlotIdentifier for identifying complex types and sampling them

crates/cheatcodes/src/evm/mapping.rs

@@ -1,11 +1,7 @@
 use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
-use alloy_primitives::{Address, B256, U256, keccak256, map::AddressHashMap};
+use alloy_primitives::{Address, B256};
 use alloy_sol_types::SolValue;
 use foundry_common::mapping_slots::MappingSlots;
-use revm::{
-    bytecode::opcode,
-    interpreter::{Interpreter, interpreter_types::Jumps},
-};
 
 impl Cheatcode for startMappingRecordingCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
@@ -74,29 +70,3 @@ fn slot_child<'a>(
 ) -> Option<&'a Vec<B256>> {
     mapping_slot(state, target)?.children.get(slot)
 }
-
-#[cold]
-pub(crate) fn step(mapping_slots: &mut AddressHashMap<MappingSlots>, interpreter: &Interpreter) {
-    match interpreter.bytecode.opcode() {
-        opcode::KECCAK256 => {
-            if interpreter.stack.peek(1) == Ok(U256::from(0x40)) {
-                let address = interpreter.input.target_address;
-                let offset = interpreter.stack.peek(0).expect("stack size > 1").saturating_to();
-                let data = interpreter.memory.slice_len(offset, 0x40);
-                let low = B256::from_slice(&data[..0x20]);
-                let high = B256::from_slice(&data[0x20..]);
-                let result = keccak256(&*data);
-
-                mapping_slots.entry(address).or_default().seen_sha3.insert(result, (low, high));
-            }
-        }
-        opcode::SSTORE => {
-            if let Some(mapping_slots) = mapping_slots.get_mut(&interpreter.input.target_address)
-                && let Ok(slot) = interpreter.stack.peek(0)
-            {
-                mapping_slots.insert(slot.into());
-            }
-        }
-        _ => {}
-    }
-}

crates/cheatcodes/src/inspector.rs

@@ -4,7 +4,7 @@ use crate::{
     CheatsConfig, CheatsCtxt, DynCheatcode, Error, Result,
     Vm::{self, AccountAccess},
     evm::{
-        DealRecord, GasRecord, RecordAccess, mapping,
+        DealRecord, GasRecord, RecordAccess,
         mock::{MockCallDataContext, MockCallReturnData},
         prank::Prank,
     },
@@ -33,7 +33,9 @@ use alloy_rpc_types::{
 };
 use alloy_sol_types::{SolCall, SolInterface, SolValue};
 use foundry_common::{
-    SELECTOR_LEN, TransactionMaybeSigned, evm::Breakpoints, mapping_slots::MappingSlots,
+    SELECTOR_LEN, TransactionMaybeSigned,
+    evm::Breakpoints,
+    mapping_slots::{MappingSlots, step as mapping_step},
 };
 use foundry_evm_core::{
     InspectorExt,
@@ -1103,7 +1105,7 @@ impl Inspector<EthEvmContext<&mut dyn DatabaseExt>> for Cheatcodes {
 
         // `startMappingRecording`: record SSTORE and KECCAK256.
         if let Some(mapping_slots) = &mut self.mapping_slots {
-            mapping::step(mapping_slots, interpreter);
+            mapping_step(mapping_slots, interpreter);
         }
 
         // `snapshotGas*`: take a snapshot of the current gas.

crates/common/Cargo.toml

@@ -44,6 +44,8 @@ alloy-transport.workspace = true
 alloy-consensus = { workspace = true, features = ["k256"] }
 alloy-network.workspace = true
 
+revm.workspace = true
+
 solar.workspace = true
 
 tower.workspace = true

crates/common/src/mapping_slots.rs

@@ -1,4 +1,11 @@
-use alloy_primitives::{B256, map::B256HashMap};
+use alloy_primitives::{
+    B256, U256, keccak256,
+    map::{AddressHashMap, B256HashMap},
+};
+use revm::{
+    bytecode::opcode,
+    interpreter::{Interpreter, interpreter_types::Jumps},
+};
 
 /// Recorded mapping slots.
 #[derive(Clone, Debug, Default)]
@@ -36,3 +43,30 @@ impl MappingSlots {
         }
     }
 }
+
+/// Function to be used in Inspector::step to record mapping slots and keys
+#[cold]
+pub fn step(mapping_slots: &mut AddressHashMap<MappingSlots>, interpreter: &Interpreter) {
+    match interpreter.bytecode.opcode() {
+        opcode::KECCAK256 => {
+            if interpreter.stack.peek(1) == Ok(U256::from(0x40)) {
+                let address = interpreter.input.target_address;
+                let offset = interpreter.stack.peek(0).expect("stack size > 1").saturating_to();
+                let data = interpreter.memory.slice_len(offset, 0x40);
+                let low = B256::from_slice(&data[..0x20]);
+                let high = B256::from_slice(&data[0x20..]);
+                let result = keccak256(&*data);
+
+                mapping_slots.entry(address).or_default().seen_sha3.insert(result, (low, high));
+            }
+        }
+        opcode::SSTORE => {
+            if let Some(mapping_slots) = mapping_slots.get_mut(&interpreter.input.target_address)
+                && let Ok(slot) = interpreter.stack.peek(0)
+            {
+                mapping_slots.insert(slot.into());
+            }
+        }
+        _ => {}
+    }
+}

crates/common/src/slot_identifier.rs

@@ -197,6 +197,7 @@ impl SlotIdentifier {
     ///
     /// It can also identify whether a slot belongs to a mapping if provided with [`MappingSlots`].
     pub fn identify(&self, slot: &B256, mapping_slots: Option<&MappingSlots>) -> Option<SlotInfo> {
+        trace!(?slot, "identifying slot");
         let slot_u256 = U256::from_be_bytes(slot.0);
         let slot_str = slot_u256.to_string();
 
@@ -629,7 +630,7 @@ impl SlotIdentifier {
 
         // Check if the value is another mapping (nested case)
         if let Some(value_storage_type) = self.storage_layout.types.get(value_type_ref) {
-            if value_storage_type.encoding == "mapping" {
+            if value_storage_type.encoding == ENCODING_MAPPING {
                 // Recursively resolve the nested mapping
                 let (nested_keys, final_value, _) = self.resolve_mapping_type(value_type_ref)?;
                 key_types.extend(nested_keys);

crates/evm/evm/src/executors/invariant/mod.rs

@@ -2,7 +2,10 @@ use crate::{
     executors::{Executor, RawCallResult},
     inspectors::Fuzzer,
 };
-use alloy_primitives::{Address, Bytes, FixedBytes, Selector, U256, map::HashMap};
+use alloy_primitives::{
+    Address, Bytes, FixedBytes, Selector, U256,
+    map::{AddressMap, HashMap},
+};
 use alloy_sol_types::{SolCall, sol};
 use eyre::{ContextCompat, Result, eyre};
 use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
@@ -598,6 +601,15 @@ impl<'a> InvariantExecutor<'a> {
             ));
         }
 
+        // If any of the targeted contracts have the storage layout enabled then we can sample
+        // mapping values. To accomplish, we need to record the mapping storage slots and keys.
+        let fuzz_state =
+            if targeted_contracts.targets.lock().iter().any(|(_, t)| t.storage_layout.is_some()) {
+                fuzz_state.with_mapping_slots(AddressMap::default())
+            } else {
+                fuzz_state
+            };
+
         self.executor.inspector_mut().set_fuzzer(Fuzzer {
             call_generator,
             fuzz_state: fuzz_state.clone(),

crates/evm/fuzz/src/inspector.rs

@@ -1,4 +1,5 @@
 use crate::{invariant::RandomCallGenerator, strategies::EvmFuzzState};
+use foundry_common::mapping_slots::step as mapping_step;
 use revm::{
     Inspector,
     context::{ContextTr, Transaction},
@@ -26,6 +27,9 @@ where
         // We only collect `stack` and `memory` data before and after calls.
         if self.collect {
             self.collect_data(interp);
+            if let Some(mapping_slots) = &mut self.fuzz_state.mapping_slots {
+                mapping_step(mapping_slots, interp);
+            }
         }
     }