Fix + More freshStorage tests + Kontrol tests ported

crates/cheatcodes/src/evm.rs

@@ -13,8 +13,9 @@ use foundry_evm_core::{
     backend::{DatabaseExt, RevertSnapshotAction},
     constants::{CALLER, CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, TEST_CONTRACT_ADDRESS},
 };
+use rand::Rng;
 use revm::{
-    primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY},
+    primitives::{Account, Bytecode, EvmStorageSlot, SpecId, KECCAK_EMPTY},
     InnerEvmContext,
 };
 use std::{
@@ -89,7 +90,12 @@ impl Cheatcode for loadCall {
         let Self { target, slot } = *self;
         ensure_not_precompile!(&target, ccx);
         ccx.ecx.load_account(target)?;
-        let val = ccx.ecx.sload(target, slot.into())?;
+        let mut val = ccx.ecx.sload(target, slot.into())?;
+        if ccx.state.fresh_storage.contains(&target) && val.is_cold && val.data == U256::ZERO {
+            val.data = ccx.state.rng().gen();
+            let mut account = ccx.ecx.load_account(target)?;
+            account.storage.insert(slot.into(), EvmStorageSlot::new(val.data));
+        }
         Ok(val.abi_encode())
     }
 }

crates/cheatcodes/src/inspector.rs

@@ -1478,12 +1478,14 @@ impl Cheatcodes {
         interpreter: &mut Interpreter,
         ecx: &mut EvmContext<DB>,
     ) {
-        let key = try_or_return!(interpreter.stack().peek(0));
-        let target_address = interpreter.contract().target_address;
-        if let Ok(value) = ecx.sload(target_address, key) {
-            if value.is_cold && value.data == U256::ZERO {
-                if let Ok(mut target_account) = ecx.load_account(target_address) {
-                    target_account.storage.insert(key, EvmStorageSlot::new(self.rng().gen()));
+        if interpreter.current_opcode() == op::SLOAD {
+            let key = try_or_return!(interpreter.stack().peek(0));
+            let target_address = interpreter.contract().target_address;
+            if let Ok(value) = ecx.sload(target_address, key) {
+                if value.is_cold && value.data == U256::ZERO {
+                    if let Ok(mut target_account) = ecx.load_account(target_address) {
+                        target_account.storage.insert(key, EvmStorageSlot::new(self.rng().gen()));
+                    }
                 }
             }
         }

crates/cheatcodes/src/utils.rs

@@ -153,13 +153,7 @@ impl Cheatcode for resumeTracingCall {
 impl Cheatcode for freshStorageCall {
     fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { target } = self;
-        if let Ok(target_contract) = ccx.load_account(*target) {
-            ensure!(
-                target_contract.is_created(),
-                "contract must be newly created in this transaction"
-            );
-            ccx.state.fresh_storage.push(*target);
-        }
+        ccx.state.fresh_storage.push(*target);
 
         Ok(Default::default())
     }

crates/cheatcodes/tests/random.rs

@@ -1,5 +1,7 @@
 //! Contains various tests for `random*` and `fresh*` cheatcodes.
 
+use alloy_primitives::U256;
+use foundry_config::{Config, FuzzConfig};
 use foundry_test_utils::{forgetest_init, str, util::OutputExt};
 
 // tests that `forge test` with a seed produces deterministic random values for uint and addresses.
@@ -133,10 +135,6 @@ contract CounterFreshStorageTest is DSTest {
         // Change slot and make sure new value retrieved
         counter.setOwner(index, address(111));
         assertEq(counter.getOwner(index), address(111));
-
-        require(counter.a() != 0);
-        require(counter.b() != address(0));
-        require(counter.c() != 0);
     }
 
 
@@ -150,6 +148,21 @@ contract CounterFreshStorageTest is DSTest {
         counter.setA(11);
         assertEq(counter.a(), 11);
     }
+
+    function test_fresh_storage_multiple_read_writes() public {
+        Counter counter = new Counter();
+        vm.freshStorage(address(counter));
+        uint256 slot1 = vm.randomUint(0, 100);
+        uint256 slot2 = vm.randomUint(0, 100);
+        require(slot1 != slot2, "random positions should be different");
+        address alice = counter.owners(slot1);
+        address bob = counter.owners(slot2);
+        require(alice != bob, "random storage values should be different");
+        counter.setOwner(slot1, bob);
+        counter.setOwner(slot2, alice);
+        assertEq(alice, counter.owners(slot2));
+        assertEq(bob, counter.owners(slot1));
+    }
 }
      "#,
     )
@@ -158,7 +171,119 @@ contract CounterFreshStorageTest is DSTest {
     cmd.args(["test"]).assert_success().stdout_eq(str![[r#"
 ...
 [PASS] test_fresh_storage() ([GAS])
+[PASS] test_fresh_storage_multiple_read_writes() ([GAS])
 [PASS] test_fresh_storage_warm() ([GAS])
 ...
 "#]]);
 });
+
+// tests deterministic `freshStorage` cheatcode.
+forgetest_init!(test_fresh_storage_with_seed, |prj, cmd| {
+    prj.wipe_contracts();
+    prj.insert_ds_test();
+    prj.insert_vm();
+    prj.clear();
+
+    let config = Config {
+        fuzz: { FuzzConfig { seed: Some(U256::from(100)), ..Default::default() } },
+        ..Default::default()
+    };
+    prj.write_config(config);
+
+    prj.add_source(
+        "Counter.t.sol",
+        r#"pragma solidity 0.8.24;
+import {Vm} from "./Vm.sol";
+import {DSTest} from "./test.sol";
+contract Counter {
+    uint256[] public a;
+    address[] public b;
+    int8[] public c;
+    bytes32[] public d;
+}
+
+contract CounterFreshStorageTest is DSTest {
+    Vm vm = Vm(HEVM_ADDRESS);
+
+    function test_fresh_storage_with_seed() public {
+        Counter counter = new Counter();
+        vm.freshStorage(address(counter));
+        assertEq(counter.a(11), 85286582241781868037363115933978803127245343755841464083427462398552335014708);
+        assertEq(counter.b(22), 0x939180Daa938F9e18Ff0E76c112D25107D358B02);
+        assertEq(counter.c(33), -104);
+        assertEq(counter.d(44), 0x6c178fa9c434f142df61a5355cc2b8d07be691b98dabf5b1a924f2bce97a19c7);
+    }
+}
+     "#,
+    )
+    .unwrap();
+
+    cmd.args(["test"]).assert_success().stdout_eq(str![[r#"
+...
+[PASS] test_fresh_storage_with_seed() ([GAS])
+...
+"#]]);
+});
+
+// Kontrol `symbolicStorage` cheatcode tests ported to `freshStorage`.
+forgetest_init!(test_fresh_storage_kontrol, |prj, cmd| {
+    prj.wipe_contracts();
+    prj.insert_ds_test();
+    prj.insert_vm();
+    prj.clear();
+
+    let config = Config {
+        fuzz: { FuzzConfig { seed: Some(U256::from(100)), ..Default::default() } },
+        ..Default::default()
+    };
+    prj.write_config(config);
+
+    prj.add_source(
+        "Counter.t.sol",
+        r#"pragma solidity 0.8.24;
+import {Vm} from "./Vm.sol";
+import {DSTest} from "./test.sol";
+contract SymbolicStore {
+    uint256 public testNumber = 1337; // slot 0
+    constructor() {}
+}
+
+contract SymbolicStorageTest is DSTest {
+    Vm vm = Vm(HEVM_ADDRESS);
+
+    function test_SymbolicStorage() public {
+         uint256 slot = vm.randomUint(0, 100);
+         address addr = 0xEA674fdDe714fd979de3EdF0F56AA9716B898ec8;
+         vm.freshStorage(addr);
+         bytes32 value = vm.load(addr, bytes32(slot));
+         assertEq(uint256(value), 85286582241781868037363115933978803127245343755841464083427462398552335014708);
+         // Load slot again and make sure we get same value.
+         bytes32 value1 = vm.load(addr, bytes32(slot));
+         assertEq(uint256(value), uint256(value1));
+    }
+
+    function test_SymbolicStorage1() public {
+        uint256 slot = vm.randomUint(0, 100);
+        SymbolicStore myStore = new SymbolicStore();
+        vm.freshStorage(address(myStore));
+        bytes32 value = vm.load(address(myStore), bytes32(uint256(slot)));
+        assertEq(uint256(value), 85286582241781868037363115933978803127245343755841464083427462398552335014708);
+    }
+
+    function testEmptyInitialStorage(uint256 slot) public {
+        bytes32 storage_value = vm.load(address(vm), bytes32(slot));
+        assertEq(uint256(storage_value), 0);
+    }
+}
+     "#,
+    )
+    .unwrap();
+
+    cmd.args(["test"]).assert_success().stdout_eq(str![[r#"
+...
+[PASS] testEmptyInitialStorage(uint256) (runs: 256, [AVG_GAS])
+[PASS] test_SymbolicStorage() ([GAS])
+[PASS] test_SymbolicStorage1() ([GAS])
+...
+"#]]);
+});