feat(forge): add vm.setArbitraryStorage with overwrites (foundry-rs#10219)

* feat(forge): add vm.setArbitraryStorage with overwrites

* Changes after review: better naming, import

Author: grandizzy

crates/cheatcodes/assets/cheatcodes.json

@@ -2826,6 +2826,11 @@ interface Vm {
     #[cheatcode(group = Utilities)]
     function setArbitraryStorage(address target) external;
 
+    /// Utility cheatcode to set arbitrary storage for given target address and overwrite
+    /// any storage slots that have been previously set.
+    #[cheatcode(group = Utilities)]
+    function setArbitraryStorage(address target, bool overwrite) external;
+
     /// Sorts an array in ascending order.
     #[cheatcode(group = Utilities)]
     function sort(uint256[] calldata array) external returns (uint256[] memory);

crates/cheatcodes/spec/src/vm.rs

@@ -23,7 +23,7 @@ use crate::{
 };
 use alloy_primitives::{
     hex,
-    map::{AddressHashMap, HashMap},
+    map::{AddressHashMap, HashMap, HashSet},
     Address, Bytes, Log, TxKind, B256, U256,
 };
 use alloy_rpc_types::{
@@ -301,12 +301,19 @@ pub struct ArbitraryStorage {
     pub values: HashMap<Address, HashMap<U256, U256>>,
     /// Mapping of address with storage copied to arbitrary storage address source.
     pub copies: HashMap<Address, Address>,
+    /// Address with storage slots that should be overwritten even if previously set.
+    pub overwrites: HashSet<Address>,
 }
 
 impl ArbitraryStorage {
     /// Marks an address with arbitrary storage.
-    pub fn mark_arbitrary(&mut self, address: &Address) {
+    pub fn mark_arbitrary(&mut self, address: &Address, overwrite: bool) {
         self.values.insert(*address, HashMap::default());
+        if overwrite {
+            self.overwrites.insert(*address);
+        } else {
+            self.overwrites.remove(address);
+        }
     }
 
     /// Maps an address that copies storage with the arbitrary storage address.
@@ -1210,6 +1217,27 @@ where {
         }
     }
 
+    /// Whether the given slot of address with arbitrary storage should be overwritten.
+    /// True if address is marked as and overwrite and if no value was previously generated for
+    /// given slot.
+    pub fn should_overwrite_arbitrary_storage(
+        &self,
+        address: &Address,
+        storage_slot: U256,
+    ) -> bool {
+        match &self.arbitrary_storage {
+            Some(storage) => {
+                storage.overwrites.contains(address) &&
+                    storage
+                        .values
+                        .get(address)
+                        .and_then(|arbitrary_values| arbitrary_values.get(&storage_slot))
+                        .is_none()
+            }
+            None => false,
+        }
+    }
+
     /// Whether the given address is a copy of an address with arbitrary storage.
     pub fn is_arbitrary_storage_copy(&self, address: &Address) -> bool {
         match &self.arbitrary_storage {
@@ -1844,7 +1872,9 @@ impl Cheatcodes {
             return;
         };
 
-        if value.is_cold && value.data.is_zero() {
+        if (value.is_cold && value.data.is_zero()) ||
+            self.should_overwrite_arbitrary_storage(&target_address, key)
+        {
             if self.has_arbitrary_storage(&target_address) {
                 let arbitrary_value = self.rng().gen();
                 self.arbitrary_storage.as_mut().unwrap().save(

crates/cheatcodes/src/inspector.rs

@@ -193,10 +193,19 @@ impl Cheatcode for resumeTracingCall {
     }
 }
 
-impl Cheatcode for setArbitraryStorageCall {
+impl Cheatcode for setArbitraryStorage_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { target } = self;
-        ccx.state.arbitrary_storage().mark_arbitrary(target);
+        ccx.state.arbitrary_storage().mark_arbitrary(target, false);
+
+        Ok(Default::default())
+    }
+}
+
+impl Cheatcode for setArbitraryStorage_1Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { target, overwrite } = self;
+        ccx.state.arbitrary_storage().mark_arbitrary(target, *overwrite);
 
         Ok(Default::default())
     }

crates/cheatcodes/src/utils.rs

@@ -123,3 +123,33 @@ contract SymbolicStorageWithSeedTest is DSTest {
         assertEq(uint256(storage_value), 0);
     }
 }
+
+// <https://github.com/foundry-rs/foundry/issues/10084>
+contract ArbitraryStorageOverwriteWithSeedTest is DSTest {
+    Vm vm = Vm(HEVM_ADDRESS);
+    uint256 _value;
+
+    function testArbitraryStorageFalse(uint256 value) public {
+        _value = value;
+        vm.setArbitraryStorage(address(this), false);
+        assertEq(_value, value);
+    }
+
+    function testArbitraryStorageTrue(uint256 value) public {
+        _value = value;
+        vm.setArbitraryStorage(address(this), true);
+        assertTrue(_value != value);
+    }
+
+    function testArbitraryStorageFalse_setAfter(uint256 value) public {
+        vm.setArbitraryStorage(address(this), false);
+        _value = value;
+        assertEq(_value, value);
+    }
+
+    function testArbitraryStorageTrue_setAfter(uint256 value) public {
+        vm.setArbitraryStorage(address(this), true);
+        _value = value;
+        assertEq(_value, value);
+    }
+}