feat(forge): Invariant Testing v2 (#1572)

* init

* invariant testing kinda working

* updates

* fmt

* wip

* wip

* wip

* check if there is a fuzzer for invariants

* less clones

* add support for targetContracts on invariant tests

* move load_contracts

* add TestOptions and invariant_depth as param

* pass TestOptions on fuzz tests

* fuzz senders as well

* light cleanup

* make counterexample list concise

* show reverts on invariants test reports

* add excludeContracts()

* refactor address fetching

* move invariant to fuzz module

* fuzz calldata from state changes

* move block into assert_invariances

* add union between selected senders and random

* fix sender on get_addresses

* wip

* add targetSelectors

* add fail_on_revert for invariant tests

* dont stop on the first invariant failure on each case

* create a new strategy tree if a new contract is created

* only collect contract addresses from NewlyCreated

* display contract and sig on displaying counter example

* add documentation

* generate the sequence lazily instead

* wip

* refactor invariants into multi file module

* refactor get_addresses to get_list

* add test cases

* add reentrancy_strat

* set reentrancy target as an union with random

* merge master

* make call_override a flag

* add inspector_config() and inspector_config_mut()

* always collect data, even without override set

* docs

* more docs

* more docs

* remove unnecessary changeset clone & docs

* refactor +prepare_fuzzing

* more explanations and better var names

* replace TestKindGas for a more generic TestKindReport

* add docs to strategies

* smol fixes

* format failure sequence

* pass TestOptions instead of fuzzer to multicontractrunner

* small fixes

* make counterexample an enum

* add InvariantFailures

* turn add_function into get_function

* improve error report on assert_invariants

* simplify refs

* only override_call_strat needs to be sboxed, revert others

* fix invariant test regression

* fix: set_replay after setting the last_sequence

* fix: test_contract address comparison on call gen

* check invariants before calling anything

* improve doc on invariant_call_override

* remove unused error map from testrunner

* reset executor instead of db

* add type alias InvariantPreparation

* move InvariantExecutor into the same file

* add return status

* small refactor

* const instead of static

* merge fixes: backend + testoptions

* use iterator for functions

* FuzzRunIdentifiedContracts now uses Mutex

* from_utf8_lossy instead of unsafe unchecked

* use Mutex for runner of RandomCallGenerator

* move RandomCallGenerator to its own module

* write to fmt

* small refactor: error.replay

* remove newlines

* add remaining is_invariant_test

Co-authored-by: Brock <brock.elmore@gmail.com>

Author: joshieDo

Cargo.lock

@@ -26,7 +26,7 @@ use forge::{
     executor::{inspector::CheatsConfig, opts::EvmOpts},
     result::SuiteResult,
     trace::identifier::LocalTraceIdentifier,
-    MultiContractRunnerBuilder,
+    MultiContractRunnerBuilder, TestOptions,
 };
 use foundry_common::{evm::EvmArgs, fs};
 use foundry_config::{figment::Figment, Config};
@@ -251,29 +251,26 @@ impl CoverageArgs {
         config: Config,
         evm_opts: EvmOpts,
     ) -> eyre::Result<()> {
-        // Setup the fuzzer
-        // TODO: Add CLI Options to modify the persistence
-        let cfg = proptest::test_runner::Config {
-            failure_persistence: None,
-            cases: config.fuzz_runs,
-            max_local_rejects: config.fuzz_max_local_rejects,
-            max_global_rejects: config.fuzz_max_global_rejects,
+        let test_options = TestOptions {
+            fuzz_runs: config.fuzz_runs,
+            fuzz_max_local_rejects: config.fuzz_max_local_rejects,
+            fuzz_max_global_rejects: config.fuzz_max_global_rejects,
             ..Default::default()
         };
-        let fuzzer = proptest::test_runner::TestRunner::new(cfg);
+
         let root = project.paths.root;
 
         let env = evm_opts.evm_env_blocking();
 
         // Build the contract runner
         let evm_spec = utils::evm_spec(&config.evm_version);
         let mut runner = MultiContractRunnerBuilder::default()
-            .fuzzer(fuzzer)
             .initial_balance(evm_opts.initial_balance)
             .evm_spec(evm_spec)
             .sender(evm_opts.sender)
             .with_fork(evm_opts.get_fork(&config, env.clone()))
             .with_cheats_config(CheatsConfig::new(&config, &evm_opts))
+            .with_test_options(test_options)
             .set_coverage(true)
             .build(root.clone(), output, env, evm_opts)?;
 
@@ -283,7 +280,8 @@ impl CoverageArgs {
         let local_identifier = LocalTraceIdentifier::new(&runner.known_contracts);
 
         // TODO: Coverage for fuzz tests
-        let handle = thread::spawn(move || runner.test(&self.filter, Some(tx)).unwrap());
+        let handle =
+            thread::spawn(move || runner.test(&self.filter, Some(tx), Default::default()).unwrap());
         for mut result in rx.into_iter().flat_map(|(_, suite)| suite.test_results.into_values()) {
             if let Some(hit_map) = result.coverage.take() {
                 for (_, trace) in &mut result.traces {

cli/src/cmd/forge/coverage.rs

@@ -10,7 +10,7 @@ use crate::cmd::{
 use clap::{Parser, ValueHint};
 use ethers::types::U256;
 use eyre::Context;
-use forge::result::TestKindGas;
+use forge::result::TestKindReport;
 use once_cell::sync::Lazy;
 use regex::Regex;
 use std::{
@@ -210,7 +210,7 @@ impl SnapshotConfig {
 pub struct SnapshotEntry {
     pub contract_name: String,
     pub signature: String,
-    pub gas_used: TestKindGas,
+    pub gas_used: TestKindReport,
 }
 
 impl FromStr for SnapshotEntry {
@@ -226,7 +226,9 @@ impl FromStr for SnapshotEntry {
                             Some(SnapshotEntry {
                                 contract_name: file.as_str().to_string(),
                                 signature: sig.as_str().to_string(),
-                                gas_used: TestKindGas::Standard(gas.as_str().parse().unwrap()),
+                                gas_used: TestKindReport::Standard {
+                                    gas: gas.as_str().parse().unwrap(),
+                                },
                             })
                         } else {
                             cap.name("runs")
@@ -237,10 +239,10 @@ impl FromStr for SnapshotEntry {
                                 .map(|(runs, avg, med)| SnapshotEntry {
                                     contract_name: file.as_str().to_string(),
                                     signature: sig.as_str().to_string(),
-                                    gas_used: TestKindGas::Fuzz {
+                                    gas_used: TestKindReport::Fuzz {
                                         runs: runs.as_str().parse().unwrap(),
-                                        median: med.as_str().parse().unwrap(),
-                                        mean: avg.as_str().parse().unwrap(),
+                                        median_gas: med.as_str().parse().unwrap(),
+                                        mean_gas: avg.as_str().parse().unwrap(),
                                     },
                                 })
                         }
@@ -274,13 +276,7 @@ fn write_to_snapshot_file(
 ) -> eyre::Result<()> {
     let mut out = String::new();
     for test in tests {
-        writeln!(
-            out,
-            "{}:{} {}",
-            test.contract_name(),
-            test.signature,
-            test.result.kind.gas_used()
-        )?;
+        writeln!(out, "{}:{} {}", test.contract_name(), test.signature, test.result.kind.report())?;
     }
     Ok(fs::write(path, out)?)
 }
@@ -289,8 +285,8 @@ fn write_to_snapshot_file(
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct SnapshotDiff {
     pub signature: String,
-    pub source_gas_used: TestKindGas,
-    pub target_gas_used: TestKindGas,
+    pub source_gas_used: TestKindReport,
+    pub target_gas_used: TestKindReport,
 }
 
 impl SnapshotDiff {
@@ -321,7 +317,7 @@ fn check(tests: Vec<Test>, snaps: Vec<SnapshotEntry>) -> bool {
         if let Some(target_gas) =
             snaps.get(&(test.contract_name().to_string(), test.signature.clone())).cloned()
         {
-            let source_gas = test.result.kind.gas_used();
+            let source_gas = test.result.kind.report();
             if source_gas.gas() != target_gas.gas() {
                 eprintln!(
                     "Diff in \"{}::{}\": consumed \"{}\" gas, expected \"{}\" gas ",
@@ -363,7 +359,7 @@ fn diff(tests: Vec<Test>, snaps: Vec<SnapshotEntry>) -> eyre::Result<()> {
             })?;
 
         diffs.push(SnapshotDiff {
-            source_gas_used: test.result.kind.gas_used(),
+            source_gas_used: test.result.kind.report(),
             signature: test.signature,
             target_gas_used,
         });
@@ -429,7 +425,7 @@ mod tests {
             SnapshotEntry {
                 contract_name: "Test".to_string(),
                 signature: "deposit()".to_string(),
-                gas_used: TestKindGas::Standard(7222)
+                gas_used: TestKindReport::Standard { gas: 7222 }
             }
         );
     }
@@ -443,7 +439,7 @@ mod tests {
             SnapshotEntry {
                 contract_name: "Test".to_string(),
                 signature: "deposit()".to_string(),
-                gas_used: TestKindGas::Fuzz { runs: 256, median: 200, mean: 100 }
+                gas_used: TestKindReport::Fuzz { runs: 256, median_gas: 200, mean_gas: 100 }
             }
         );
     }

cli/src/cmd/forge/snapshot.rs

@@ -8,6 +8,7 @@ use crate::{
     compile::ProjectCompiler,
     suggestions, utils,
 };
+use cast::fuzz::CounterExample;
 use clap::{AppSettings, Parser};
 use ethers::{solc::utils::RuntimeOrHandle, types::U256};
 use forge::{
@@ -19,11 +20,10 @@ use forge::{
         identifier::{EtherscanIdentifier, LocalTraceIdentifier},
         CallTraceDecoderBuilder, TraceKind,
     },
-    MultiContractRunner, MultiContractRunnerBuilder,
+    MultiContractRunner, MultiContractRunnerBuilder, TestOptions,
 };
 use foundry_common::evm::EvmArgs;
 use foundry_config::{figment, figment::Figment, Config};
-use proptest::test_runner::{RngAlgorithm, TestRng};
 use regex::Regex;
 use std::{collections::BTreeMap, path::PathBuf, sync::mpsc::channel, thread, time::Duration};
 use tracing::trace;
@@ -181,7 +181,7 @@ pub struct Test {
 
 impl Test {
     pub fn gas_used(&self) -> u64 {
-        self.result.kind.gas_used().gas()
+        self.result.kind.report().gas()
     }
 
     /// Returns the contract name of the artifact id
@@ -280,47 +280,47 @@ fn short_test_result(name: &str, result: &TestResult) {
     let status = if result.success {
         Paint::green("[PASS]".to_string())
     } else {
-        let txt = match (&result.reason, &result.counterexample) {
-            (Some(ref reason), Some(ref counterexample)) => {
-                format!("[FAIL. Reason: {reason}. Counterexample: {counterexample}]")
-            }
-            (None, Some(ref counterexample)) => {
-                format!("[FAIL. Counterexample: {counterexample}]")
-            }
-            (Some(ref reason), None) => {
-                format!("[FAIL. Reason: {reason}]")
-            }
-            (None, None) => "[FAIL]".to_string(),
-        };
+        let reason = result
+            .reason
+            .as_ref()
+            .map(|reason| format!("Reason: {reason}"))
+            .unwrap_or_else(|| "Reason: Undefined.".to_string());
+
+        let counterexample = result
+            .counterexample
+            .as_ref()
+            .map(|example| match example {
+                CounterExample::Single(eg) => format!(" Counterexample: {eg}]"),
+                CounterExample::Sequence(sequence) => {
+                    let mut inner_txt = String::new();
+
+                    for checkpoint in sequence {
+                        inner_txt += format!("\t\t{checkpoint}\n").as_str();
+                    }
+                    format!("]\n\t[Sequence]\n{inner_txt}\n")
+                }
+            })
+            .unwrap_or_else(|| "]".to_string());
 
-        Paint::red(txt)
+        Paint::red(format!("[FAIL. {reason}{counterexample}"))
     };
 
-    println!("{} {} {}", status, name, result.kind.gas_used());
+    println!("{} {} {}", status, name, result.kind.report());
 }
 
 pub fn custom_run(args: TestArgs) -> eyre::Result<TestOutcome> {
     // Merge all configs
     let (config, mut evm_opts) = args.config_and_evm_opts()?;
 
-    // Setup the fuzzer
-    // TODO: Add CLI Options to modify the persistence
-    let cfg = proptest::test_runner::Config {
-        failure_persistence: None,
-        cases: config.fuzz_runs,
-        max_local_rejects: config.fuzz_max_local_rejects,
-        max_global_rejects: config.fuzz_max_global_rejects,
-        ..Default::default()
-    };
-
-    let fuzzer = if let Some(ref fuzz_seed) = config.fuzz_seed {
-        let mut bytes: [u8; 32] = [0; 32];
-        fuzz_seed.to_big_endian(&mut bytes);
-        trace!(target: "forge::test", "executing test command");
-        let rng = TestRng::from_seed(RngAlgorithm::ChaCha, &bytes);
-        proptest::test_runner::TestRunner::new_with_rng(cfg, rng)
-    } else {
-        proptest::test_runner::TestRunner::new(cfg)
+    let test_options = TestOptions {
+        fuzz_runs: config.fuzz_runs,
+        fuzz_max_local_rejects: config.fuzz_max_local_rejects,
+        fuzz_max_global_rejects: config.fuzz_max_global_rejects,
+        fuzz_seed: config.fuzz_seed,
+        invariant_runs: config.invariant_runs,
+        invariant_depth: config.invariant_depth,
+        invariant_fail_on_revert: config.invariant_fail_on_revert,
+        invariant_call_override: config.invariant_call_override,
     };
 
     let mut filter = args.filter(&config);
@@ -350,20 +350,21 @@ pub fn custom_run(args: TestArgs) -> eyre::Result<TestOutcome> {
     let evm_spec = utils::evm_spec(&config.evm_version);
 
     let mut runner = MultiContractRunnerBuilder::default()
-        .fuzzer(fuzzer)
         .initial_balance(evm_opts.initial_balance)
         .evm_spec(evm_spec)
         .sender(evm_opts.sender)
         .with_fork(evm_opts.get_fork(&config, env.clone()))
         .with_cheats_config(CheatsConfig::new(&config, &evm_opts))
+        .with_test_options(test_options)
         .build(project.paths.root, output, env, evm_opts)?;
 
     if args.debug.is_some() {
         filter.test_pattern = args.debug;
+
         match runner.count_filtered_tests(&filter) {
                 1 => {
                     // Run the test
-                    let results = runner.test(&filter, None)?;
+                    let results = runner.test(&filter, None, test_options)?;
 
                     // Get the result of the single test
                     let (id, sig, test_kind, counterexample) = results.iter().map(|(id, SuiteResult{ test_results, .. })| {
@@ -375,7 +376,7 @@ pub fn custom_run(args: TestArgs) -> eyre::Result<TestOutcome> {
                     // Build debugger args if this is a fuzz test
                     let sig = match test_kind {
                         TestKind::Fuzz(cases) => {
-                            if let Some(counterexample) = counterexample {
+                            if let Some(CounterExample::Single(counterexample)) = counterexample {
                                 counterexample.calldata.to_string()
                             } else {
                                 cases.cases().first().expect("no fuzz cases run").calldata.to_string()
@@ -407,7 +408,16 @@ pub fn custom_run(args: TestArgs) -> eyre::Result<TestOutcome> {
     } else if args.list {
         list(runner, filter, args.json)
     } else {
-        test(config, runner, verbosity, filter, args.json, args.allow_failure, args.gas_report)
+        test(
+            config,
+            runner,
+            verbosity,
+            filter,
+            args.json,
+            args.allow_failure,
+            test_options,
+            args.gas_report,
+        )
     }
 }
 
@@ -438,6 +448,7 @@ fn test(
     filter: Filter,
     json: bool,
     allow_failure: bool,
+    test_options: TestOptions,
     gas_reporting: bool,
 ) -> eyre::Result<TestOutcome> {
     trace!(target: "forge::test", "running all tests");
@@ -462,7 +473,7 @@ fn test(
     }
 
     if json {
-        let results = runner.test(&filter, None)?;
+        let results = runner.test(&filter, None, test_options)?;
         println!("{}", serde_json::to_string(&results)?);
         Ok(TestOutcome::new(results, allow_failure))
     } else {
@@ -483,7 +494,7 @@ fn test(
         let (tx, rx) = channel::<(String, SuiteResult)>();
 
         // Run tests
-        let handle = thread::spawn(move || runner.test(&filter, Some(tx)).unwrap());
+        let handle = thread::spawn(move || runner.test(&filter, Some(tx), test_options).unwrap());
 
         let mut results: BTreeMap<String, SuiteResult> = BTreeMap::new();
         let mut gas_report = GasReport::new(config.gas_reports);

cli/src/cmd/forge/test/mod.rs

@@ -219,7 +219,7 @@ impl MultiWallet {
             }
         );
 
-        let mut error_msg = "".to_string();
+        let mut error_msg = String::new();
 
         // This is an actual used address
         if addresses.contains(&Config::DEFAULT_SENDER) {

cli/src/opts/multi_wallet.rs

@@ -60,6 +60,10 @@ forgetest!(can_extract_config_values, |prj: TestProject, mut cmd: TestCommand| {
         fuzz_max_local_rejects: 2000,
         fuzz_max_global_rejects: 100203,
         fuzz_seed: Some(1000.into()),
+        invariant_runs: 256,
+        invariant_depth: 15,
+        invariant_fail_on_revert: false,
+        invariant_call_override: false,
         ffi: true,
         sender: "00a329c0648769A73afAc7F9381D08FB43dBEA72".parse().unwrap(),
         tx_origin: "00a329c0648769A73afAc7F9F81E08FB43dBEA72".parse().unwrap(),