feat(invariants): sample typed storage values (#11204)

* feat(`forge`): sample typed storage values

* arc it

* nit

* clippy

* nit

* strip file prefixes

* fmt

* don't add adjacent values to sample

* correctly match artifact identifier and name

* clippy

* nit

* handle simple mappings while inserting storage values

* cleanup + insert only DynSolType's found in mappings

* Nit: with_project_contracts

---------

Co-authored-by: grandizzy <[email protected]>

Author: yash-atreya

crates/common/src/contracts.rs

@@ -9,7 +9,7 @@ use foundry_compilers::{
     ArtifactId, Project, ProjectCompileOutput,
     artifacts::{
         BytecodeObject, CompactBytecode, CompactContractBytecode, CompactDeployedBytecode,
-        ConfigurableContractArtifact, ContractBytecodeSome, Offsets,
+        ConfigurableContractArtifact, ContractBytecodeSome, Offsets, StorageLayout,
     },
     utils::canonicalized,
 };
@@ -75,6 +75,8 @@ pub struct ContractData {
     pub bytecode: Option<BytecodeData>,
     /// Contract runtime code.
     pub deployed_bytecode: Option<BytecodeData>,
+    /// Contract storage layout, if available.
+    pub storage_layout: Option<Arc<StorageLayout>>,
 }
 
 impl ContractData {
@@ -120,6 +122,29 @@ impl ContractsByArtifact {
                         abi: abi?,
                         bytecode: bytecode.map(Into::into),
                         deployed_bytecode: deployed_bytecode.map(Into::into),
+                        storage_layout: None,
+                    },
+                ))
+            })
+            .collect();
+        Self(Arc::new(map))
+    }
+
+    /// Creates a new instance from project compile output, preserving storage layouts.
+    pub fn with_storage_layout(output: ProjectCompileOutput) -> Self {
+        let map = output
+            .into_artifacts()
+            .filter_map(|(id, artifact)| {
+                let name = id.name.clone();
+                let abi = artifact.abi?;
+                Some((
+                    id,
+                    ContractData {
+                        name,
+                        abi,
+                        bytecode: artifact.bytecode.map(Into::into),
+                        deployed_bytecode: artifact.deployed_bytecode.map(Into::into),
+                        storage_layout: artifact.storage_layout.map(Arc::new),
                     },
                 ))
             })

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

@@ -790,7 +790,11 @@ impl<'a> InvariantExecutor<'a> {
                     && self.artifact_filters.matches(identifier)
             })
             .map(|(addr, (identifier, abi))| {
-                (*addr, TargetedContract::new(identifier.clone(), abi.clone()))
+                (
+                    *addr,
+                    TargetedContract::new(identifier.clone(), abi.clone())
+                        .with_project_contracts(self.project_contracts),
+                )
             })
             .collect();
         let mut contracts = TargetedContracts { inner: contracts };
@@ -833,8 +837,12 @@ impl<'a> InvariantExecutor<'a> {
             // Identifiers are specified as an array, so we loop through them.
             for identifier in artifacts {
                 // Try to find the contract by name or identifier in the project's contracts.
-                if let Some(abi) = self.project_contracts.find_abi_by_name_or_identifier(identifier)
+                if let Some((_, contract_data)) =
+                    self.project_contracts.iter().find(|(artifact, _)| {
+                        &artifact.name == identifier || &artifact.identifier() == identifier
+                    })
                 {
+                    let abi = &contract_data.abi;
                     combined
                         // Check if there's an entry for the given key in the 'combined' map.
                         .entry(*addr)
@@ -844,7 +852,13 @@ impl<'a> InvariantExecutor<'a> {
                             entry.abi.functions.extend(abi.functions.clone());
                         })
                         // Otherwise insert it into the map.
-                        .or_insert_with(|| TargetedContract::new(identifier.to_string(), abi));
+                        .or_insert_with(|| {
+                            let mut contract =
+                                TargetedContract::new(identifier.to_string(), abi.clone());
+                            contract.storage_layout =
+                                contract_data.storage_layout.as_ref().map(Arc::clone);
+                            contract
+                        });
                 }
             }
         }
@@ -944,7 +958,10 @@ impl<'a> InvariantExecutor<'a> {
                         address
                     )
                 })?;
-                entry.insert(TargetedContract::new(identifier.clone(), abi.clone()))
+                entry.insert(
+                    TargetedContract::new(identifier.clone(), abi.clone())
+                        .with_project_contracts(self.project_contracts),
+                )
             }
         };
         contract.add_selectors(selectors.iter().copied(), should_exclude)?;

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

@@ -1,5 +1,6 @@
 use alloy_json_abi::{Function, JsonAbi};
-use alloy_primitives::{Address, Bytes, Selector};
+use alloy_primitives::{Address, Bytes, Selector, map::HashMap};
+use foundry_compilers::artifacts::StorageLayout;
 use itertools::Either;
 use parking_lot::Mutex;
 use serde::{Deserialize, Serialize};
@@ -75,6 +76,7 @@ impl FuzzRunIdentifiedContracts {
                 abi: contract.abi.clone(),
                 targeted_functions: functions,
                 excluded_functions: Vec::new(),
+                storage_layout: contract.storage_layout.as_ref().map(Arc::clone),
             };
             targets.insert(*address, contract);
         }
@@ -146,6 +148,16 @@ impl TargetedContracts {
                 .map(|function| format!("{}.{}", contract.identifier.clone(), function.name))
         })
     }
+
+    /// Returns a map of contract addresses to their storage layouts.
+    pub fn get_storage_layouts(&self) -> HashMap<Address, Arc<StorageLayout>> {
+        self.inner
+            .iter()
+            .filter_map(|(addr, c)| {
+                c.storage_layout.as_ref().map(|layout| (*addr, Arc::clone(layout)))
+            })
+            .collect()
+    }
 }
 
 impl std::ops::Deref for TargetedContracts {
@@ -173,12 +185,33 @@ pub struct TargetedContract {
     pub targeted_functions: Vec<Function>,
     /// The excluded functions of the contract.
     pub excluded_functions: Vec<Function>,
+    /// The contract's storage layout, if available.
+    pub storage_layout: Option<Arc<StorageLayout>>,
 }
 
 impl TargetedContract {
     /// Returns a new `TargetedContract` instance.
     pub fn new(identifier: String, abi: JsonAbi) -> Self {
-        Self { identifier, abi, targeted_functions: Vec::new(), excluded_functions: Vec::new() }
+        Self {
+            identifier,
+            abi,
+            targeted_functions: Vec::new(),
+            excluded_functions: Vec::new(),
+            storage_layout: None,
+        }
+    }
+
+    /// Determines contract storage layout from project contracts. Needs `storageLayout` to be
+    /// enabled as extra output in project configuration.
+    pub fn with_project_contracts(mut self, project_contracts: &ContractsByArtifact) -> Self {
+        if let Some((src, name)) = self.identifier.split_once(':')
+            && let Some((_, contract_data)) = project_contracts.iter().find(|(artifact, _)| {
+                artifact.name == name && artifact.source.as_path().ends_with(src)
+            })
+        {
+            self.storage_layout = contract_data.storage_layout.as_ref().map(Arc::clone);
+        }
+        self
     }
 
     /// Helper to retrieve functions to fuzz for specified abi.

crates/evm/fuzz/src/strategies/state.rs

@@ -6,6 +6,7 @@ use alloy_primitives::{
     map::{AddressIndexSet, B256IndexSet, HashMap},
 };
 use foundry_common::ignore_metadata_hash;
+use foundry_compilers::artifacts::StorageLayout;
 use foundry_config::FuzzDictionaryConfig;
 use foundry_evm_core::utils::StateChangeset;
 use parking_lot::{RawRwLock, RwLock, lock_api::RwLockReadGuard};
@@ -72,8 +73,10 @@ impl EvmFuzzState {
             let (target_abi, target_function) = targets.fuzzed_artifacts(tx);
             dict.insert_logs_values(target_abi, logs, run_depth);
             dict.insert_result_values(target_function, result, run_depth);
+            // Get storage layouts for contracts in the state changeset
+            let storage_layouts = targets.get_storage_layouts();
+            dict.insert_new_state_values(state_changeset, &storage_layouts);
         }
-        dict.insert_new_state_values(state_changeset);
     }
 
     /// Removes all newly added entries from the dictionary.
@@ -151,7 +154,7 @@ impl FuzzDictionary {
                 // Sort storage values before inserting to ensure deterministic dictionary.
                 let values = account.storage.iter().collect::<BTreeMap<_, _>>();
                 for (slot, value) in values {
-                    self.insert_storage_value(slot, value);
+                    self.insert_storage_value(slot, value, None);
                 }
             }
         }
@@ -226,16 +229,21 @@ impl FuzzDictionary {
 
     /// Insert values from call state changeset into fuzz dictionary.
     /// These values are removed at the end of current run.
-    fn insert_new_state_values(&mut self, state_changeset: &StateChangeset) {
+    fn insert_new_state_values(
+        &mut self,
+        state_changeset: &StateChangeset,
+        storage_layouts: &HashMap<Address, Arc<StorageLayout>>,
+    ) {
         for (address, account) in state_changeset {
             // Insert basic account information.
             self.insert_value(address.into_word());
             // Insert push bytes.
             self.insert_push_bytes_values(address, &account.info);
             // Insert storage values.
             if self.config.include_storage {
+                let storage_layout = storage_layouts.get(address).map(|arc| arc.as_ref());
                 for (slot, value) in &account.storage {
-                    self.insert_storage_value(slot, &value.present_value);
+                    self.insert_storage_value(slot, &value.present_value, storage_layout);
                 }
             }
         }
@@ -288,19 +296,132 @@ impl FuzzDictionary {
         }
     }
 
+    /// Resolves storage types from a storage layout for a given slot and all mapping types.
+    /// Returns a tuple of (slot_type, mapping_types) where slot_type is the specific type
+    /// for the storage slot and mapping_types are all mapping value types found in the layout.
+    fn resolve_storage_types(
+        &self,
+        storage_layout: Option<&StorageLayout>,
+        storage_slot: &U256,
+    ) -> (Option<DynSolType>, Vec<DynSolType>) {
+        let Some(layout) = storage_layout else {
+            return (None, Vec::new());
+        };
+
+        // Try to determine the type of this specific storage slot
+        let slot_type =
+            layout.storage.iter().find(|s| s.slot == storage_slot.to_string()).and_then(
+                |storage| {
+                    layout
+                        .types
+                        .get(&storage.storage_type)
+                        .and_then(|t| DynSolType::parse(&t.label).ok())
+                },
+            );
+
+        // Collect all mapping value types from the layout
+        let mapping_types = layout
+            .types
+            .values()
+            .filter_map(|type_info| {
+                if type_info.encoding == "mapping"
+                    && let Some(t_value) = type_info.value.as_ref()
+                    && let Some(mapping_value) = t_value.strip_prefix("t_")
+                {
+                    DynSolType::parse(mapping_value).ok()
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        (slot_type, mapping_types)
+    }
+
     /// Insert values from single storage slot and storage value into fuzz dictionary.
     /// If storage values are newly collected then they are removed at the end of current run.
-    fn insert_storage_value(&mut self, storage_slot: &U256, storage_value: &U256) {
+    fn insert_storage_value(
+        &mut self,
+        storage_slot: &U256,
+        storage_value: &U256,
+        storage_layout: Option<&StorageLayout>,
+    ) {
+        // Always insert the slot itself
         self.insert_value(B256::from(*storage_slot));
-        self.insert_value(B256::from(*storage_value));
-        // also add the value below and above the storage value to the dictionary.
-        if *storage_value != U256::ZERO {
-            let below_value = storage_value - U256::from(1);
-            self.insert_value(B256::from(below_value));
+
+        let (slot_type, mapping_types) = self.resolve_storage_types(storage_layout, storage_slot);
+
+        if let Some(sol_type) = slot_type {
+            self.insert_decoded_storage_value(sol_type, storage_value);
+        } else if !mapping_types.is_empty() {
+            self.insert_mapping_storage_values(mapping_types, storage_value);
+        } else {
+            // No type information available, insert as raw values (old behavior)
+            self.insert_value(B256::from(*storage_value));
+            // also add the value below and above the storage value to the dictionary.
+            if *storage_value != U256::ZERO {
+                let below_value = storage_value - U256::from(1);
+                self.insert_value(B256::from(below_value));
+            }
+            if *storage_value != U256::MAX {
+                let above_value = storage_value + U256::from(1);
+                self.insert_value(B256::from(above_value));
+            }
         }
-        if *storage_value != U256::MAX {
-            let above_value = storage_value + U256::from(1);
-            self.insert_value(B256::from(above_value));
+    }
+
+    /// Insert decoded storage values into the fuzz dictionary.
+    /// Only simple static type values are inserted as sample values.
+    /// Complex types (dynamic arrays, structs) are inserted as raw values
+    fn insert_decoded_storage_value(&mut self, sol_type: DynSolType, storage_value: &U256) {
+        // Only insert values for types that can be represented as a single word
+        match &sol_type {
+            DynSolType::Address
+            | DynSolType::Uint(_)
+            | DynSolType::Int(_)
+            | DynSolType::Bool
+            | DynSolType::FixedBytes(_)
+            | DynSolType::Bytes => {
+                // Insert as a typed sample value
+                self.sample_values.entry(sol_type).or_default().insert(B256::from(*storage_value));
+            }
+            _ => {
+                // For complex types (arrays, mappings, structs), insert as raw value
+                self.insert_value(B256::from(*storage_value));
+            }
+        }
+    }
+
+    /// Insert storage values of mapping value types as sample values in the fuzz dictionary.
+    ///
+    /// ```solidity
+    /// mapping(uint256 => address) public myMapping;
+    /// // `address` is the mapping value type here.
+    /// // `uint256` is the mapping key type.
+    /// ```
+    ///
+    /// A storage value is inserted if and only if it can be decoded into one of the mapping
+    /// [`DynSolType`] value types found in the [`StorageLayout`].
+    ///
+    /// If decoding fails, the value is inserted as a raw value.
+    fn insert_mapping_storage_values(
+        &mut self,
+        mapping_types: Vec<DynSolType>,
+        storage_value: &U256,
+    ) {
+        for sol_type in mapping_types {
+            match sol_type.abi_decode(storage_value.as_le_slice()) {
+                Ok(_) => {
+                    self.sample_values
+                        .entry(sol_type)
+                        .or_default()
+                        .insert(B256::from(*storage_value));
+                }
+                Err(_) => {
+                    // If decoding fails, insert as raw value
+                    self.insert_value(B256::from(*storage_value));
+                }
+            }
         }
     }
 

crates/forge/src/multi_runner.rs

@@ -519,7 +519,10 @@ impl MultiContractRunnerBuilder {
             }
         }
 
-        let known_contracts = ContractsByArtifact::new(linked_contracts);
+        // Create known contracts with storage layout information
+        let known_contracts = ContractsByArtifact::with_storage_layout(
+            output.clone().with_stripped_file_prefixes(root),
+        );
 
         Ok(MultiContractRunner {
             contracts: deployable_contracts,