feat: add useSeed cheatcode to set RNG seed (#10698)

* feat: mock deterministic shuffle implementation

* feat: implement useSeed()

* fix(test): typo

* fix(test): correct implementation of randomUint()
chore: format

* fix: refactor seed handling logic and remove redundant assertion message

* chore: remove unused imports and redundant newline

* fix(test): add seed validation and tests for shuffle consistency

* chore: add Shuffle test

* Adapt to updated proptest

---------

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

Author: Another-DevX

crates/cheatcodes/assets/cheatcodes.json

@@ -2879,6 +2879,10 @@ interface Vm {
     #[cheatcode(group = Utilities)]
     function shuffle(uint256[] calldata array) external returns (uint256[] memory);
 
+    /// Set RNG seed.
+    #[cheatcode(group = Utilities)]
+    function setSeed(uint256 seed) external;
+
     /// Causes the next contract creation (via new) to fail and return its initcode in the returndata buffer.
     /// This allows type-safe access to the initcode payload that would be used for contract creation.
     /// Example usage:

crates/cheatcodes/spec/src/vm.rs

@@ -1264,6 +1264,13 @@ impl Cheatcodes {
         })
     }
 
+    pub fn set_seed(&mut self, seed: U256) {
+        self.test_runner = Some(TestRunner::new_with_rng(
+            proptest::test_runner::Config::default(),
+            TestRng::from_seed(RngAlgorithm::ChaCha, &seed.to_be_bytes::<32>()),
+        ));
+    }
+
     /// Returns existing or set a default `ArbitraryStorage` option.
     /// Used by `setArbitraryStorage` cheatcode to track addresses with arbitrary storage.
     pub fn arbitrary_storage(&mut self) -> &mut ArbitraryStorage {

crates/cheatcodes/src/inspector.rs

@@ -273,6 +273,14 @@ impl Cheatcode for shuffleCall {
     }
 }
 
+impl Cheatcode for setSeedCall {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { seed } = self;
+        ccx.state.set_seed(*seed);
+        Ok(Default::default())
+    }
+}
+
 /// Helper to generate a random `uint` value (with given bits or bounded if specified)
 /// from type strategy.
 fn random_uint(state: &mut Cheatcodes, bits: Option<U256>, bounds: Option<(U256, U256)>) -> Result {

crates/cheatcodes/src/utils.rs

@@ -0,0 +1,92 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+
+contract SeedTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    function testSeedAffectsRandom() public {
+        // Use a known seed
+        uint256 seed = 123456789;
+        vm.setSeed(seed);
+
+        // Call a foundry cheatcode to get a random value (this depends on the integration)
+        uint256 rand1 = uint256(vm.randomUint());
+
+        // Reset the seed and verify the result is the same
+        vm.setSeed(seed);
+        uint256 rand2 = uint256(vm.randomUint());
+
+        uint256 rand3 = uint256(vm.randomUint());
+        // If the seed is the same, the random value must be equal
+        assertEq(rand1, rand2);
+        assertTrue(rand1 != rand3);
+    }
+
+    function testSeedChangesRandom() public {
+        // Use one seed
+        vm.setSeed(1);
+        uint256 randA = uint256(vm.randomUint());
+
+        // Use a different seed
+        vm.setSeed(2);
+        uint256 randB = uint256(vm.randomUint());
+
+        // Values must be different
+        assertTrue(randA != randB, "Random value must be different if seed is different");
+    }
+
+    function testSeedAffectsShuffle() public {
+        // Use a known seed
+        uint256 seed = 123456789;
+        vm.setSeed(seed);
+
+        // Create two identical arrays
+        uint256[] memory array1 = new uint256[](5);
+        uint256[] memory array2 = new uint256[](5);
+        for (uint256 i = 0; i < 5; i++) {
+            array1[i] = i;
+            array2[i] = i;
+        }
+
+        // Shuffle both arrays with the same seed
+        array1 = vm.shuffle(array1);
+        vm.setSeed(seed); // Reset the seed to get the same shuffle pattern
+        array2 = vm.shuffle(array2);
+
+        // Compare elements - they should be identical after shuffle
+        for (uint256 i = 0; i < array1.length; i++) {
+            assertEq(array1[i], array2[i], "Arrays should be identical with same seed");
+        }
+    }
+
+    function testDifferentSeedsProduceDifferentShuffles() public {
+        // Create the initial array
+        uint256[] memory array1 = new uint256[](5);
+        uint256[] memory array2 = new uint256[](5);
+        for (uint256 i = 0; i < 5; i++) {
+            array1[i] = i;
+            array2[i] = i;
+        }
+
+        // Use first seed
+        vm.setSeed(1);
+        array1 = vm.shuffle(array1);
+
+        // Use second seed
+        vm.setSeed(2);
+        array2 = vm.shuffle(array2);
+
+        // Arrays should be different (we'll check at least one difference exists)
+        bool foundDifference = false;
+        for (uint256 i = 0; i < array1.length; i++) {
+            if (array1[i] != array2[i]) {
+                foundDifference = true;
+                break;
+            }
+        }
+        assertTrue(foundDifference, "Arrays should be different with different seeds");
+    }
+}

testdata/cheats/Vm.sol

@@ -0,0 +1,61 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+
+contract ShuffleTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    function testDeterministicShuffle() public {
+        // Use a known seed
+        uint256 seed = 123456789;
+        vm.setSeed(seed);
+
+        // Create two identical arrays
+        uint256[] memory array1 = new uint256[](5);
+        uint256[] memory array2 = new uint256[](5);
+        for (uint256 i = 0; i < 5; i++) {
+            array1[i] = i;
+            array2[i] = i;
+        }
+
+        // Shuffle both arrays with the same seed
+        array1 = vm.shuffle(array1);
+        vm.setSeed(seed); // Reset the seed to get the same shuffle pattern
+        array2 = vm.shuffle(array2);
+
+        // Compare elements - they should be identical after shuffle
+        for (uint256 i = 0; i < array1.length; i++) {
+            assertEq(array1[i], array2[i], "Arrays should be identical with same seed");
+        }
+    }
+
+    function testDifferentSeedsProduceDifferentShuffles() public {
+        // Create the initial array
+        uint256[] memory array1 = new uint256[](5);
+        uint256[] memory array2 = new uint256[](5);
+        for (uint256 i = 0; i < 5; i++) {
+            array1[i] = i;
+            array2[i] = i;
+        }
+
+        // Use first seed
+        vm.setSeed(1);
+        array1 = vm.shuffle(array1);
+
+        // Use second seed
+        vm.setSeed(2);
+        array2 = vm.shuffle(array2);
+
+        // Arrays should be different (we'll check at least one difference exists)
+        bool foundDifference = false;
+        for (uint256 i = 0; i < array1.length; i++) {
+            if (array1[i] != array2[i]) {
+                foundDifference = true;
+                break;
+            }
+        }
+        assertTrue(foundDifference, "Arrays should be different with different seeds");
+    }
+}