feat(cheatcodes): add deterministic random value generation with seed (…

…foundry-rs#8622) * feat(cheatcodes): add ability to set seed for `vm.randomUint()` * chore: move `vm.randomAddress` test to its own contract * feat(cheatcodes): add ability to set seed for `vm.randomAddress()` * feat: use global seed instead of introducing new cheatcodes * chore: clean up * chore: clean up tests * feat: add `fuzz.seed` as inline parameter in tests * chore: trim 0x prefix * chore: nit * test: update random tests * fix: inline parsing on fuzz seed * test: set seed and update random tests * chore: remove inline config for seed * chore: clean up * chore: clean up tests * test: remove deterministic tests from testdata * test: implement forgetest to test that forge test with a seed produces deterministic random values * test: fix tests * chore: clean up * test: remove seed * fix: clippy and forge-fmt * chore: clean up * chore: rename test contract * fix: lint * chore: move rng to state instead of creating a new one when calling `vm.random*` cheats * chore: nit * test: update tests * fix: clippy * chore: nit * chore: clean up * Update crates/cheatcodes/src/inspector.rs Co-authored-by: DaniPopes <[email protected]> * test: only check outputs are the same or different * chore: clean up * chore: nits --------- Co-authored-by: DaniPopes <[email protected]>
programskillforverification · Aug 8, 2024 · 1f33c6f · 1f33c6f
1 parent b574cdf
commit 1f33c6f
Show file tree

Hide file tree

Showing 7 changed files with 114 additions and 15 deletions.
diff --git a/crates/cheatcodes/src/config.rs b/crates/cheatcodes/src/config.rs
@@ -1,6 +1,6 @@
 use super::Result;
 use crate::{script::ScriptWallets, Vm::Rpc};
-use alloy_primitives::Address;
+use alloy_primitives::{Address, U256};
 use foundry_common::{fs::normalize_path, ContractsByArtifact};
 use foundry_compilers::{utils::canonicalize, ProjectPathsConfig};
 use foundry_config::{
@@ -54,6 +54,8 @@ pub struct CheatsConfig {
     pub running_version: Option<Version>,
     /// Whether to enable legacy (non-reverting) assertions.
     pub assertions_revert: bool,
+    /// Optional seed for the RNG algorithm.
+    pub seed: Option<U256>,
 }
 
 impl CheatsConfig {
@@ -93,6 +95,7 @@ impl CheatsConfig {
             available_artifacts,
             running_version,
             assertions_revert: config.assertions_revert,
+            seed: config.fuzz.seed,
         }
     }
 
@@ -221,6 +224,7 @@ impl Default for CheatsConfig {
             available_artifacts: Default::default(),
             running_version: Default::default(),
             assertions_revert: true,
+            seed: None,
         }
     }
 }

diff --git a/crates/cheatcodes/src/fs.rs b/crates/cheatcodes/src/fs.rs
@@ -618,7 +618,7 @@ mod tests {
             root: PathBuf::from(&env!("CARGO_MANIFEST_DIR")),
             ..Default::default()
         };
-        Cheatcodes { config: Arc::new(config), ..Default::default() }
+        Cheatcodes::new(Arc::new(config))
     }
 
     #[test]

diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs
@@ -29,6 +29,7 @@ use foundry_evm_core::{
     InspectorExt,
 };
 use itertools::Itertools;
+use rand::{rngs::StdRng, Rng, SeedableRng};
 use revm::{
     interpreter::{
         opcode, CallInputs, CallOutcome, CallScheme, CreateInputs, CreateOutcome, EOFCreateInputs,
@@ -315,6 +316,9 @@ pub struct Cheatcodes {
     /// Breakpoints supplied by the `breakpoint` cheatcode.
     /// `char -> (address, pc)`
     pub breakpoints: Breakpoints,
+
+    /// Optional RNG algorithm.
+    rng: Option<StdRng>,
 }
 
 // This is not derived because calling this in `fn new` with `..Default::default()` creates a second
@@ -356,6 +360,7 @@ impl Cheatcodes {
             mapping_slots: Default::default(),
             pc: Default::default(),
             breakpoints: Default::default(),
+            rng: Default::default(),
         }
     }
 
@@ -926,6 +931,13 @@ impl Cheatcodes {
 
         None
     }
+
+    pub fn rng(&mut self) -> &mut impl Rng {
+        self.rng.get_or_insert_with(|| match self.config.seed {
+            Some(seed) => StdRng::from_seed(seed.to_be_bytes::<32>()),
+            None => StdRng::from_entropy(),
+        })
+    }
 }
 
 impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {

diff --git a/crates/cheatcodes/src/utils.rs b/crates/cheatcodes/src/utils.rs
@@ -55,22 +55,21 @@ impl Cheatcode for ensNamehashCall {
 }
 
 impl Cheatcode for randomUint_0Call {
-    fn apply(&self, _state: &mut Cheatcodes) -> Result {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self {} = self;
-        // Use thread_rng to get a random number
-        let mut rng = rand::thread_rng();
+        let rng = state.rng();
         let random_number: U256 = rng.gen();
         Ok(random_number.abi_encode())
     }
 }
 
 impl Cheatcode for randomUint_1Call {
-    fn apply(&self, _state: &mut Cheatcodes) -> Result {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { min, max } = *self;
         ensure!(min <= max, "min must be less than or equal to max");
         // Generate random between range min..=max
-        let mut rng = rand::thread_rng();
         let exclusive_modulo = max - min;
+        let rng = state.rng();
         let mut random_number = rng.gen::<U256>();
         if exclusive_modulo != U256::MAX {
             let inclusive_modulo = exclusive_modulo + U256::from(1);
@@ -82,9 +81,10 @@ impl Cheatcode for randomUint_1Call {
 }
 
 impl Cheatcode for randomAddressCall {
-    fn apply(&self, _state: &mut Cheatcodes) -> Result {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self {} = self;
-        let addr = Address::random();
+        let rng = state.rng();
+        let addr = Address::random_with(rng);
         Ok(addr.abi_encode())
     }
 }
diff --git a/crates/forge/tests/cli/test_cmd.rs b/crates/forge/tests/cli/test_cmd.rs
@@ -6,6 +6,7 @@ use foundry_test_utils::{
     rpc, str,
     util::{OutputExt, OTHER_SOLC_VERSION, SOLC_VERSION},
 };
+use similar_asserts::assert_eq;
 use std::{path::PathBuf, str::FromStr};
 
 // tests that test filters are handled correctly
@@ -597,7 +598,7 @@ contract GasWaster {
 contract GasLimitTest is Test {
     function test() public {
         vm.createSelectFork("<rpc>");
-        
+
         GasWaster waster = new GasWaster();
         waster.waste();
     }
@@ -613,7 +614,7 @@ contract GasLimitTest is Test {
 forgetest!(test_match_path, |prj, cmd| {
     prj.add_source(
         "dummy",
-        r"  
+        r"
 contract Dummy {
     function testDummy() public {}
 }
@@ -1048,3 +1049,76 @@ Traces:
 "#
     ]]);
 });
+
+// tests that `forge test` with a seed produces deterministic random values for uint and addresses.
+forgetest_init!(deterministic_randomness_with_seed, |prj, cmd| {
+    prj.wipe_contracts();
+    prj.add_test(
+        "DeterministicRandomnessTest.t.sol",
+        r#"
+import {Test, console} from "forge-std/Test.sol";
+
+contract DeterministicRandomnessTest is Test {
+
+    function testDeterministicRandomUint() public {
+        console.log(vm.randomUint());
+        console.log(vm.randomUint());
+        console.log(vm.randomUint());
+    }
+
+    function testDeterministicRandomUintRange() public {
+        uint256 min = 0;
+        uint256 max = 1000000000;
+        console.log(vm.randomUint(min, max));
+        console.log(vm.randomUint(min, max));
+        console.log(vm.randomUint(min, max));
+    }
+
+    function testDeterministicRandomAddress() public {
+        console.log(vm.randomAddress());
+        console.log(vm.randomAddress());
+        console.log(vm.randomAddress());
+    }
+}
+"#,
+    )
+    .unwrap();
+
+    // Extracts the test result section from the DeterministicRandomnessTest contract output.
+    fn extract_test_result(out: &str) -> &str {
+        let start = out
+            .find("for test/DeterministicRandomnessTest.t.sol:DeterministicRandomnessTest")
+            .unwrap();
+        let end = out.find("Suite result: ok.").unwrap();
+        &out[start..end]
+    }
+
+    // Run the test twice with the same seed and verify the outputs are the same.
+    let seed1 = "0xa1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2";
+    cmd.args(["test", "--fuzz-seed", seed1, "-vv"]).assert_success();
+    let out1 = cmd.stdout_lossy();
+    let res1 = extract_test_result(&out1);
+
+    cmd.forge_fuse();
+    cmd.args(["test", "--fuzz-seed", seed1, "-vv"]).assert_success();
+    let out2 = cmd.stdout_lossy();
+    let res2 = extract_test_result(&out2);
+
+    assert_eq!(res1, res2);
+
+    // Run the test with another seed and verify the output differs.
+    let seed2 = "0xb1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3f4a5b6c7d8e9f0a1b2";
+    cmd.forge_fuse();
+    cmd.args(["test", "--fuzz-seed", seed2, "-vv"]).assert_success();
+    let out3 = cmd.stdout_lossy();
+    let res3 = extract_test_result(&out3);
+    assert_ne!(res3, res1);
+
+    // Run the test without a seed and verify the outputs differs once again.
+    cmd.forge_fuse();
+    cmd.args(["test", "-vv"]).assert_success();
+    let out4 = cmd.stdout_lossy();
+    let res4 = extract_test_result(&out4);
+    assert_ne!(res4, res1);
+    assert_ne!(res4, res3);
+});
diff --git a/testdata/default/cheats/RandomAddress.t.sol b/testdata/default/cheats/RandomAddress.t.sol
@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity 0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+
+contract RandomAddress is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    function testRandomAddress() public {
+        vm.randomAddress();
+    }
+}
diff --git a/testdata/default/cheats/RandomUint.t.sol b/testdata/default/cheats/RandomUint.t.sol
@@ -26,8 +26,4 @@ contract RandomUint is DSTest {
         assertTrue(rand >= min, "rand >= min");
         assertTrue(rand <= max, "rand <= max");
     }
-
-    function testRandomAddress() public {
-        vm.randomAddress();
-    }
 }