feat(cheatcodes): include slot type and decode values in state diffs

crates/cheatcodes/src/evm.rs

@@ -6,11 +6,13 @@ use crate::{
     inspector::{Ecx, RecordDebugStepInfo},
 };
 use alloy_consensus::TxEnvelope;
+use alloy_dyn_abi::{DynSolType, DynSolValue};
 use alloy_genesis::{Genesis, GenesisAccount};
-use alloy_primitives::{Address, B256, U256, map::HashMap};
+use alloy_primitives::{Address, B256, U256, hex, map::HashMap};
 use alloy_rlp::Decodable;
 use alloy_sol_types::SolValue;
 use foundry_common::fs::{read_json_file, write_json_file};
+use foundry_compilers::artifacts::StorageLayout;
 use foundry_evm_core::{
     ContextExt,
     backend::{DatabaseExt, RevertStateSnapshotAction},
@@ -30,6 +32,7 @@ use std::{
     collections::{BTreeMap, HashSet, btree_map::Entry},
     fmt::Display,
     path::Path,
+    str::FromStr,
 };
 
 mod record_debug_step;
@@ -103,6 +106,58 @@ struct SlotStateDiff {
     previous_value: B256,
     /// Current storage value.
     new_value: B256,
+    /// Decoded values according to the Solidity type (e.g., uint256, address).
+    /// Only present when storage layout is available and decoding succeeds.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    decoded: Option<DecodedSlotValues>,
+
+    /// Storage layout metadata (variable name, type, offset).
+    /// Only present when contract has storage layout output.
+    #[serde(skip_serializing_if = "Option::is_none", flatten)]
+    slot_info: Option<SlotInfo>,
+}
+
+/// Storage slot metadata from the contract's storage layout.
+#[derive(Serialize, Debug)]
+struct SlotInfo {
+    /// Variable name (e.g., "owner", "values\[0\]", "config.maxSize").
+    label: String,
+    /// Solidity type, serialized as string (e.g., "uint256", "address").
+    #[serde(rename = "type", serialize_with = "serialize_dyn_sol_type")]
+    dyn_sol_type: DynSolType,
+    /// Byte offset within the 32-byte slot (0 for full slot, 0-31 for packed).
+    offset: i64,
+    /// Storage slot number as decimal string.
+    slot: String,
+}
+
+fn serialize_dyn_sol_type<S>(dyn_type: &DynSolType, serializer: S) -> Result<S::Ok, S::Error>
+where
+    S: serde::Serializer,
+{
+    serializer.serialize_str(&dyn_type.to_string())
+}
+
+/// Decoded storage values showing before and after states.
+#[derive(Debug)]
+struct DecodedSlotValues {
+    /// Decoded value before the state change.
+    previous_value: DynSolValue,
+    /// Decoded value after the state change.
+    new_value: DynSolValue,
+}
+
+impl Serialize for DecodedSlotValues {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        use serde::ser::SerializeStruct;
+        let mut state = serializer.serialize_struct("DecodedSlotValues", 2)?;
+        state.serialize_field("previousValue", &format_dyn_sol_value_raw(&self.previous_value))?;
+        state.serialize_field("newValue", &format_dyn_sol_value_raw(&self.new_value))?;
+        state.end()
+    }
 }
 
 /// Balance diff info.
@@ -170,11 +225,38 @@ impl Display for AccountStateDiffs {
         if !&self.state_diff.is_empty() {
             writeln!(f, "- state diff:")?;
             for (slot, slot_changes) in &self.state_diff {
-                writeln!(
-                    f,
-                    "@ {slot}: {} → {}",
-                    slot_changes.previous_value, slot_changes.new_value
-                )?;
+                match (&slot_changes.slot_info, &slot_changes.decoded) {
+                    (Some(slot_info), Some(decoded)) => {
+                        // Have both slot info and decoded values - only show decoded values
+                        writeln!(
+                            f,
+                            "@ {slot} ({}, {}): {} → {}",
+                            slot_info.label,
+                            slot_info.dyn_sol_type,
+                            format_dyn_sol_value_raw(&decoded.previous_value),
+                            format_dyn_sol_value_raw(&decoded.new_value)
+                        )?;
+                    }
+                    (Some(slot_info), None) => {
+                        // Have slot info but no decoded values - show raw hex values
+                        writeln!(
+                            f,
+                            "@ {slot} ({}, {}): {} → {}",
+                            slot_info.label,
+                            slot_info.dyn_sol_type,
+                            slot_changes.previous_value,
+                            slot_changes.new_value
+                        )?;
+                    }
+                    _ => {
+                        // No slot info - show raw hex values
+                        writeln!(
+                            f,
+                            "@ {slot}: {} → {}",
+                            slot_changes.previous_value, slot_changes.new_value
+                        )?;
+                    }
+                }
             }
         }
 
@@ -1257,12 +1339,20 @@ fn get_recorded_state_diffs(ccx: &mut CheatsCtxt) -> BTreeMap<Address, AccountSt
         }
     }
 
-    // Look up contract names for all addresses
+    // Look up contract names and storage layouts for all addresses
     let mut contract_names = HashMap::new();
+    let mut storage_layouts = HashMap::new();
     for address in addresses_to_lookup {
         if let Some(name) = get_contract_name(ccx, address) {
             contract_names.insert(address, name);
         }
+
+        // Also get storage layout if available
+        if let Some((_artifact_id, contract_data)) = get_contract_data(ccx, address)
+            && let Some(storage_layout) = &contract_data.storage_layout
+        {
+            storage_layouts.insert(address, storage_layout.clone());
+        }
     }
 
     // Now process the records
@@ -1331,13 +1421,44 @@ fn get_recorded_state_diffs(ccx: &mut CheatsCtxt) -> BTreeMap<Address, AccountSt
                         // Update state diff. Do not overwrite the initial value if already set.
                         match account_diff.state_diff.entry(storage_access.slot) {
                             Entry::Vacant(slot_state_diff) => {
+                                // Get storage layout info for this slot
+                                let slot_info = storage_layouts
+                                    .get(&storage_access.account)
+                                    .and_then(|layout| get_slot_info(layout, &storage_access.slot));
+
+                                // Try to decode values if we have slot info
+                                let decoded = slot_info.as_ref().and_then(|info| {
+                                    let prev = decode_storage_value(
+                                        storage_access.previousValue,
+                                        &info.dyn_sol_type,
+                                    )?;
+                                    let new = decode_storage_value(
+                                        storage_access.newValue,
+                                        &info.dyn_sol_type,
+                                    )?;
+                                    Some(DecodedSlotValues { previous_value: prev, new_value: new })
+                                });
+
                                 slot_state_diff.insert(SlotStateDiff {
                                     previous_value: storage_access.previousValue,
                                     new_value: storage_access.newValue,
+                                    decoded,
+                                    slot_info,
                                 });
                             }
                             Entry::Occupied(mut slot_state_diff) => {
-                                slot_state_diff.get_mut().new_value = storage_access.newValue;
+                                let entry = slot_state_diff.get_mut();
+                                entry.new_value = storage_access.newValue;
+
+                                if let Some(slot_info) = &entry.slot_info
+                                    && let Some(ref mut decoded) = entry.decoded
+                                    && let Some(new_value) = decode_storage_value(
+                                        storage_access.newValue,
+                                        &slot_info.dyn_sol_type,
+                                    )
+                                {
+                                    decoded.new_value = new_value;
+                                }
                             }
                         }
                     }
@@ -1375,6 +1496,158 @@ fn get_contract_name(ccx: &mut CheatsCtxt, address: Address) -> Option<String> {
     None
 }
 
+/// Helper function to get the contract data from the deployed code at an address.
+fn get_contract_data<'a>(
+    ccx: &'a mut CheatsCtxt,
+    address: Address,
+) -> Option<(&'a foundry_compilers::ArtifactId, &'a foundry_common::contracts::ContractData)> {
+    // Check if we have available artifacts to match against
+    let artifacts = ccx.state.config.available_artifacts.as_ref()?;
+
+    // Try to load the account and get its code
+    let account = ccx.ecx.journaled_state.load_account(address).ok()?;
+    let code = account.info.code.as_ref()?;
+
+    // Skip if code is empty
+    if code.is_empty() {
+        return None;
+    }
+
+    // Try to find the artifact by deployed code
+    let code_bytes = code.original_bytes();
+    if let Some(result) = artifacts.find_by_deployed_code_exact(&code_bytes) {
+        return Some(result);
+    }
+
+    // Fallback to fuzzy matching if exact match fails
+    artifacts.find_by_deployed_code(&code_bytes)
+}
+
+/// Gets storage layout info for a specific slot.
+fn get_slot_info(storage_layout: &StorageLayout, slot: &B256) -> Option<SlotInfo> {
+    let slot = U256::from_be_bytes(slot.0);
+    let slot_str = slot.to_string();
+
+    for storage in &storage_layout.storage {
+        let base_slot = U256::from_str(&storage.slot).ok()?;
+        let storage_type = storage_layout.types.get(&storage.storage_type)?;
+        let dyn_type = DynSolType::parse(&storage_type.label).ok()?;
+
+        // Check for exact slot match
+        if storage.slot == slot_str {
+            let label = match &dyn_type {
+                DynSolType::FixedArray(_, _) => {
+                    // For arrays, label the base slot with indices
+                    format!("{}{}", storage.label, get_array_base_indices(&dyn_type))
+                }
+                _ => storage.label.clone(),
+            };
+
+            return Some(SlotInfo {
+                label,
+                dyn_sol_type: dyn_type,
+                offset: storage.offset,
+                slot: storage.slot.clone(),
+            });
+        }
+
+        // Check if slot is part of a static array
+        if let DynSolType::FixedArray(_, _) = &dyn_type
+            && let Ok(total_bytes) = storage_type.number_of_bytes.parse::<u64>()
+        {
+            let total_slots = total_bytes.div_ceil(32);
+
+            // Check if slot is within array range
+            if slot > base_slot && slot < base_slot + U256::from(total_slots) {
+                let index = (slot - base_slot).to::<u64>();
+                let label = format_array_element_label(&storage.label, &dyn_type, index);
+
+                return Some(SlotInfo {
+                    label,
+                    dyn_sol_type: dyn_type,
+                    offset: 0,
+                    slot: slot.to_string(),
+                });
+            }
+        }
+    }
+
+    None
+}
+
+/// Returns the base index [\0\] or [\0\][\0\] for a fixed array type depending on the dimensions.
+fn get_array_base_indices(dyn_type: &DynSolType) -> String {
+    match dyn_type {
+        DynSolType::FixedArray(inner, _) => {
+            if let DynSolType::FixedArray(_, _) = inner.as_ref() {
+                // Nested array (2D or higher)
+                format!("[0]{}", get_array_base_indices(inner))
+            } else {
+                // Simple 1D array
+                "[0]".to_string()
+            }
+        }
+        _ => String::new(),
+    }
+}
+
+/// Helper function to format an array element label given its index
+fn format_array_element_label(base_label: &str, dyn_type: &DynSolType, index: u64) -> String {
+    match dyn_type {
+        DynSolType::FixedArray(inner, _) => {
+            if let DynSolType::FixedArray(_, inner_size) = inner.as_ref() {
+                // 2D array: calculate row and column
+                let row = index / (*inner_size as u64);
+                let col = index % (*inner_size as u64);
+                format!("{base_label}[{row}][{col}]")
+            } else {
+                // 1D array
+                format!("{base_label}[{index}]")
+            }
+        }
+        _ => base_label.to_string(),
+    }
+}
+
+/// Helper function to decode a single storage value using its DynSolType
+fn decode_storage_value(value: B256, dyn_type: &DynSolType) -> Option<DynSolValue> {
+    // Storage values are always 32 bytes, stored as a single word
+    // For arrays, we need to unwrap to the base element type
+    let mut actual_type = dyn_type;
+    // Unwrap nested arrays to get to the base element type.
+    while let DynSolType::FixedArray(elem_type, _) = actual_type {
+        actual_type = elem_type.as_ref();
+    }
+
+    // Use abi_decode to decode the value
+    actual_type.abi_decode(&value.0).ok()
+}
+
+/// Helper function to format DynSolValue as raw string without type information
+fn format_dyn_sol_value_raw(value: &DynSolValue) -> String {
+    match value {
+        DynSolValue::Bool(b) => b.to_string(),
+        DynSolValue::Int(i, _) => i.to_string(),
+        DynSolValue::Uint(u, _) => u.to_string(),
+        DynSolValue::FixedBytes(bytes, size) => hex::encode_prefixed(&bytes.0[..*size]),
+        DynSolValue::Address(addr) => addr.to_string(),
+        DynSolValue::Function(func) => func.as_address_and_selector().1.to_string(),
+        DynSolValue::Bytes(bytes) => hex::encode_prefixed(bytes),
+        DynSolValue::String(s) => s.clone(),
+        DynSolValue::Array(values) | DynSolValue::FixedArray(values) => {
+            let formatted: Vec<String> = values.iter().map(format_dyn_sol_value_raw).collect();
+            format!("[{}]", formatted.join(", "))
+        }
+        DynSolValue::Tuple(values) => {
+            let formatted: Vec<String> = values.iter().map(format_dyn_sol_value_raw).collect();
+            format!("({})", formatted.join(", "))
+        }
+        DynSolValue::CustomStruct { name: _, prop_names: _, tuple } => {
+            format_dyn_sol_value_raw(&DynSolValue::Tuple(tuple.clone()))
+        }
+    }
+}
+
 /// Helper function to set / unset cold storage slot of the target address.
 fn set_cold_slot(ccx: &mut CheatsCtxt, target: Address, slot: U256, cold: bool) {
     if let Some(account) = ccx.ecx.journaled_state.state.get_mut(&target)