feat(cheatcodes): vm.getStorageSlots (#11537)

* feat(common): identify bytes and string slots using storageLayout

* fmt

* clippy

* aggregate and decode bytes and strings

* cleanup

* nits

* clippy

* fmt

* nit

* cleanup

* identify the slots using the length of data in base slot value

* nit

* fix

* fix mapping identification

* feat(cheatcodes): vm.getStorageSlots

* fix tests

* nit

* fmt

* nit

* nit

* fmt

* Update crates/common/src/slot_identifier.rs

* Update crates/common/src/slot_identifier.rs

* add mapping in test

* fmt

* error mapping and dyn arrays

* nit

* nit

---------

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

Author: yash-atreya

crates/cheatcodes/assets/cheatcodes.json

@@ -432,6 +432,10 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Safe)]
     function getStateDiffJson() external view returns (string memory diff);
 
+    /// Returns an array of storage slots occupied by the specified variable.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function getStorageSlots(address target, string calldata variableName) external view returns (uint256[] memory slots);
+
     /// Returns an array of `StorageAccess` from current `vm.stateStateDiffRecording` session
     #[cheatcode(group = Evm, safety = Safe)]
     function getStorageAccesses() external view returns (StorageAccess[] memory storageAccesses);

crates/cheatcodes/spec/src/vm.rs

@@ -9,14 +9,17 @@ use alloy_consensus::TxEnvelope;
 use alloy_genesis::{Genesis, GenesisAccount};
 use alloy_network::eip2718::EIP4844_TX_TYPE_ID;
 use alloy_primitives::{
-    Address, B256, U256, hex,
+    Address, B256, U256, hex, keccak256,
     map::{B256Map, HashMap},
 };
 use alloy_rlp::Decodable;
 use alloy_sol_types::SolValue;
 use foundry_common::{
     fs::{read_json_file, write_json_file},
-    slot_identifier::{SlotIdentifier, SlotInfo},
+    slot_identifier::{
+        ENCODING_BYTES, ENCODING_DYN_ARRAY, ENCODING_INPLACE, ENCODING_MAPPING, SlotIdentifier,
+        SlotInfo,
+    },
 };
 use foundry_evm_core::{
     ContextExt,
@@ -37,6 +40,7 @@ use std::{
     collections::{BTreeMap, HashSet, btree_map::Entry},
     fmt::Display,
     path::Path,
+    str::FromStr,
 };
 
 mod record_debug_step;
@@ -908,6 +912,83 @@ impl Cheatcode for getStateDiffJsonCall {
     }
 }
 
+impl Cheatcode for getStorageSlotsCall {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { target, variableName } = self;
+
+        let storage_layout = get_contract_data(ccx, *target)
+            .and_then(|(_, data)| data.storage_layout.as_ref().map(|layout| layout.clone()))
+            .ok_or_else(|| fmt_err!("Storage layout not available for contract at {target}. Try compiling contracts with `--extra-output storageLayout`"))?;
+
+        trace!(storage = ?storage_layout.storage, "fetched storage");
+
+        let storage = storage_layout
+            .storage
+            .iter()
+            .find(|s| s.label.to_lowercase() == *variableName.to_lowercase())
+            .ok_or_else(|| fmt_err!("variable '{variableName}' not found in storage layout"))?;
+
+        let storage_type = storage_layout
+            .types
+            .get(&storage.storage_type)
+            .ok_or_else(|| fmt_err!("storage type not found for variable {variableName}"))?;
+
+        if storage_type.encoding == ENCODING_MAPPING || storage_type.encoding == ENCODING_DYN_ARRAY
+        {
+            return Err(fmt_err!(
+                "cannot get storage slots for variables with mapping or dynamic array types"
+            ));
+        }
+
+        let slot = U256::from_str(&storage.slot).map_err(|_| {
+            fmt_err!("invalid slot {} format for variable {variableName}", storage.slot)
+        })?;
+
+        let mut slots = Vec::new();
+
+        // Always push the base slot
+        slots.push(slot);
+
+        if storage_type.encoding == ENCODING_INPLACE {
+            // For inplace encoding, calculate the number of slots needed
+            let num_bytes = U256::from_str(&storage_type.number_of_bytes).map_err(|_| {
+                fmt_err!(
+                    "invalid number_of_bytes {} for variable {variableName}",
+                    storage_type.number_of_bytes
+                )
+            })?;
+            let num_slots = num_bytes.div_ceil(U256::from(32));
+
+            // Start from 1 since base slot is already added
+            for i in 1..num_slots.to::<usize>() {
+                slots.push(slot + U256::from(i));
+            }
+        }
+
+        if storage_type.encoding == ENCODING_BYTES {
+            // Try to check if it's a long bytes/string by reading the current storage
+            // value
+            if let Ok(value) = ccx.ecx.journaled_state.sload(*target, slot) {
+                let value_bytes = value.data.to_be_bytes::<32>();
+                let length_byte = value_bytes[31];
+                // Check if it's a long bytes/string (LSB is 1)
+                if length_byte & 1 == 1 {
+                    // Calculate data slots for long bytes/string
+                    let length: U256 = value.data >> 1;
+                    let num_data_slots = length.to::<usize>().div_ceil(32);
+                    let data_start = U256::from_be_bytes(keccak256(B256::from(slot).0).0);
+
+                    for i in 0..num_data_slots {
+                        slots.push(data_start + U256::from(i));
+                    }
+                }
+            }
+        }
+
+        Ok(slots.abi_encode())
+    }
+}
+
 impl Cheatcode for getStorageAccessesCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let mut storage_accesses = Vec::new();

crates/cheatcodes/src/evm.rs

@@ -12,10 +12,15 @@ use serde::Serialize;
 use std::{collections::BTreeMap, str::FromStr, sync::Arc};
 use tracing::trace;
 
-// Constants for storage type encodings
-const ENCODING_INPLACE: &str = "inplace";
-const ENCODING_MAPPING: &str = "mapping";
-const ENCODING_BYTES: &str = "bytes";
+/// "inplace" encoding type for variables that fit in one storage slot i.e 32 bytes
+pub const ENCODING_INPLACE: &str = "inplace";
+/// "mapping" encoding type for Solidity mappings, which use keccak256 hash-based storage
+pub const ENCODING_MAPPING: &str = "mapping";
+/// "bytes" encoding type for bytes and string types, which use either inplace or keccak256
+/// hash-based storage depending on length
+pub const ENCODING_BYTES: &str = "bytes";
+/// "dynamic_array" encoding type for dynamic arrays, which uses keccak256 hash-based storage
+pub const ENCODING_DYN_ARRAY: &str = "dynamic_array";
 
 /// Information about a storage slot including its label, type, and decoded values.
 #[derive(Serialize, Debug)]

crates/common/src/slot_identifier.rs

@@ -17,7 +17,7 @@ use foundry_test_utils::{Filter, init_tracing};
 async fn test_cheats_local(test_data: &ForgeTestData) {
     let mut filter = Filter::new(".*", ".*", &format!(".*cheats{RE_PATH_SEPARATOR}*"))
         .exclude_paths("Fork")
-        .exclude_contracts("(Isolated|WithSeed|StateDiff)");
+        .exclude_contracts("(Isolated|WithSeed|StateDiff|GetStorageSlotsTest)");
 
     // Exclude FFI tests on Windows because no `echo`, and file tests that expect certain file paths
     if cfg!(windows) {
@@ -40,7 +40,7 @@ async fn test_cheats_local(test_data: &ForgeTestData) {
 /// Executes subset of all cheat code tests in isolation mode.
 async fn test_cheats_local_isolated(test_data: &ForgeTestData) {
     let filter = Filter::new(".*", ".*(Isolated)", &format!(".*cheats{RE_PATH_SEPARATOR}*"))
-        .exclude_contracts("StateDiff");
+        .exclude_contracts("(StateDiff|GetStorageSlotsTest)");
 
     let runner = test_data.runner_with(|config| {
         config.isolate = true;
@@ -78,7 +78,11 @@ async fn test_state_diff_storage_layout() {
         let output = get_compiled(&mut project);
         ForgeTestData { project, output, config: config.into(), profile }
     };
-    let filter = Filter::new(".*", "StateDiff", &format!(".*cheats{RE_PATH_SEPARATOR}*"));
+    let filter = Filter::new(
+        ".*",
+        "(StateDiff|GetStorageSlotsTest)",
+        &format!(".*cheats{RE_PATH_SEPARATOR}*"),
+    );
 
     let runner = test_data.runner_with(|config| {
         config.fs_permissions = FsPermissions::new(vec![PathPermission::read_write("./")]);

crates/forge/tests/it/cheats.rs

@@ -0,0 +1,93 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+
+contract StorageContract {
+    // Simple variables - 1 slot each
+    uint256 public value; // Slot 0
+    address public owner; // Slot 1
+
+    // Fixed array - 3 consecutive slots
+    uint256[3] public numbers; // Slots 2, 3, 4
+
+    // Bytes variables
+    bytes public shortBytes; // Slot 5 (less than 32 bytes)
+    mapping(address => uint256) public balances; // Slot 6 Inserted in between to make sure we can still properly identify the bytes
+    bytes public longBytes; // Slot 7 (32+ bytes, will use multiple slots)
+
+    // String variables
+    string public shortString; // Slot 8 (less than 32 bytes)
+    string public longString; // Slot 9 (32+ bytes, will use multiple slots)
+
+    function setShortBytes(bytes memory _data) public {
+        shortBytes = _data;
+    }
+
+    function setLongBytes(bytes memory _data) public {
+        longBytes = _data;
+    }
+
+    function setShortString(string memory _str) public {
+        shortString = _str;
+    }
+
+    function setLongString(string memory _str) public {
+        longString = _str;
+    }
+
+    function setNumbers(uint256 a, uint256 b, uint256 c) public {
+        numbers[0] = a;
+        numbers[1] = b;
+        numbers[2] = c;
+    }
+}
+
+contract GetStorageSlotsTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    StorageContract storageContract;
+
+    function setUp() public {
+        storageContract = new StorageContract();
+    }
+
+    function testGetStorageSlots() public {
+        // Test 1: Simple variable
+        uint256[] memory slots = vm.getStorageSlots(address(storageContract), "value");
+        assertEq(slots.length, 1);
+        assertEq(slots[0], 0);
+
+        // Test 2: Fixed array (should return 3 consecutive slots)
+        slots = vm.getStorageSlots(address(storageContract), "numbers");
+        assertEq(slots.length, 3);
+        assertEq(slots[0], 2);
+        assertEq(slots[1], 3);
+        assertEq(slots[2], 4);
+
+        // Test 3: Short bytes (less than 32 bytes)
+        storageContract.setShortBytes(hex"deadbeef");
+        slots = vm.getStorageSlots(address(storageContract), "shortBytes");
+        assertEq(slots.length, 1);
+        assertEq(slots[0], 5);
+
+        // Test 4: Long bytes (100 bytes = 4 slots needed)
+        bytes memory longData = new bytes(100);
+        for (uint256 i = 0; i < 100; i++) {
+            longData[i] = bytes1(uint8(i));
+        }
+        storageContract.setLongBytes(longData);
+
+        slots = vm.getStorageSlots(address(storageContract), "longBytes");
+        // Should return 5 slots: 1 base slot + 4 data slots
+        assertEq(slots.length, 5);
+        assertEq(slots[0], 7); // Base slot
+
+        // Data slots start at keccak256(base_slot)
+        uint256 dataStart = uint256(keccak256(abi.encode(uint256(7))));
+        assertEq(slots[1], dataStart);
+        assertEq(slots[2], dataStart + 1);
+        assertEq(slots[3], dataStart + 2);
+        assertEq(slots[4], dataStart + 3);
+    }
+}