feat(invariant): assert all invariants

Author: grandizzy

crates/evm/evm/src/executors/corpus.rs

@@ -100,6 +100,7 @@ pub(crate) struct CorpusMetrics {
 impl fmt::Display for CorpusMetrics {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         writeln!(f)?;
+        writeln!(f, "      Edge coverage metrics:")?;
         writeln!(f, "        - cumulative edges seen: {}", self.cumulative_edges_seen)?;
         writeln!(f, "        - cumulative features seen: {}", self.cumulative_features_seen)?;
         writeln!(f, "        - corpus count: {}", self.corpus_count)?;

crates/evm/evm/src/executors/invariant/error.rs

@@ -1,10 +1,11 @@
 use super::InvariantContract;
 use crate::executors::RawCallResult;
+use alloy_json_abi::Function;
 use alloy_primitives::{Address, Bytes};
-use foundry_config::InvariantConfig;
 use foundry_evm_core::decode::RevertDecoder;
 use foundry_evm_fuzz::{BasicTxDetails, Reason, invariant::FuzzRunIdentifiedContracts};
 use proptest::test_runner::TestError;
+use std::{collections::HashMap, fmt};
 
 /// Stores information about failures and reverts of the invariant tests.
 #[derive(Clone, Default)]
@@ -14,16 +15,40 @@ pub struct InvariantFailures {
     /// The latest revert reason of a run.
     pub revert_reason: Option<String>,
     /// Maps a broken invariant to its specific error.
-    pub error: Option<InvariantFuzzError>,
+    pub errors: HashMap<String, InvariantFuzzError>,
 }
 
 impl InvariantFailures {
     pub fn new() -> Self {
         Self::default()
     }
 
-    pub fn into_inner(self) -> (usize, Option<InvariantFuzzError>) {
-        (self.reverts, self.error)
+    pub fn into_inner(self) -> (usize, HashMap<String, InvariantFuzzError>) {
+        (self.reverts, self.errors)
+    }
+
+    pub fn record_failure(&mut self, invariant: &Function, failure: InvariantFuzzError) {
+        self.errors.insert(invariant.name.clone(), failure);
+    }
+
+    pub fn has_failure(&self, invariant: &Function) -> bool {
+        self.errors.contains_key(&invariant.name)
+    }
+
+    pub fn get_failure(&self, invariant: &Function) -> Option<&InvariantFuzzError> {
+        self.errors.get(&invariant.name)
+    }
+
+    pub fn can_continue(&self, invariants: usize) -> bool {
+        self.errors.len() < invariants
+    }
+}
+
+impl fmt::Display for InvariantFailures {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        writeln!(f)?;
+        writeln!(f, "      ❌ Failures: {}", self.errors.len())?;
+        Ok(())
     }
 }
 
@@ -70,10 +95,11 @@ pub struct FailedInvariantCaseData {
 impl FailedInvariantCaseData {
     pub fn new(
         invariant_contract: &InvariantContract<'_>,
-        invariant_config: &InvariantConfig,
+        shrink_run_limit: u32,
+        fail_on_revert: bool,
         targeted_contracts: &FuzzRunIdentifiedContracts,
         calldata: &[BasicTxDetails],
-        call_result: RawCallResult,
+        call_result: &RawCallResult,
         inner_sequence: &[Option<BasicTxDetails>],
     ) -> Self {
         // Collect abis of fuzzed and invariant contracts to decode custom error.
@@ -82,7 +108,7 @@ impl FailedInvariantCaseData {
             .with_abi(invariant_contract.abi)
             .decode(call_result.result.as_ref(), call_result.exit_reason);
 
-        let func = invariant_contract.invariant_function;
+        let func = invariant_contract.invariant_fn;
         debug_assert!(func.inputs.is_empty());
         let origin = func.name.as_str();
         Self {
@@ -95,8 +121,8 @@ impl FailedInvariantCaseData {
             addr: invariant_contract.address,
             calldata: func.selector().to_vec().into(),
             inner_sequence: inner_sequence.to_vec(),
-            shrink_run_limit: invariant_config.shrink_run_limit,
-            fail_on_revert: invariant_config.fail_on_revert,
+            shrink_run_limit,
+            fail_on_revert,
         }
     }
 }

crates/evm/evm/src/executors/invariant/mod.rs

@@ -2,6 +2,7 @@ use crate::{
     executors::{Executor, RawCallResult},
     inspectors::Fuzzer,
 };
+use alloy_json_abi::Function;
 use alloy_primitives::{
     Address, Bytes, FixedBytes, Selector, U256,
     map::{AddressMap, HashMap},
@@ -28,7 +29,7 @@ use foundry_evm_traces::{CallTraceArena, SparsedTraceArena};
 use indicatif::ProgressBar;
 use parking_lot::RwLock;
 use proptest::{strategy::Strategy, test_runner::TestRunner};
-use result::{assert_after_invariant, assert_invariants, can_continue};
+use result::{assert_after_invariant, can_continue};
 use revm::state::Account;
 use std::{
     collections::{HashMap as Map, btree_map::Entry},
@@ -41,7 +42,7 @@ pub use error::{InvariantFailures, InvariantFuzzError};
 use foundry_evm_coverage::HitMaps;
 
 mod replay;
-pub use replay::{replay_error, replay_run};
+pub use replay::{generate_counterexample, replay_error, replay_run};
 
 mod result;
 use foundry_common::{TestFunctionExt, sh_println};
@@ -52,6 +53,7 @@ use serde_json::json;
 mod shrink;
 use crate::executors::{
     DURATION_BETWEEN_METRICS_REPORT, EarlyExit, EvmError, FuzzTestTimer, corpus::CorpusManager,
+    invariant::result::invariant_preflight_check,
 };
 pub use shrink::check_sequence;
 
@@ -130,8 +132,6 @@ struct InvariantTestData {
     last_run_inputs: Vec<BasicTxDetails>,
     // Additional traces for gas report.
     gas_report_traces: Vec<Vec<CallTraceArena>>,
-    // Last call results of the invariant test.
-    last_call_results: Option<RawCallResult>,
     // Line coverage information collected from all fuzzed calls.
     line_coverage: Option<HitMaps>,
     // Metrics for each fuzzed selector.
@@ -160,19 +160,13 @@ impl InvariantTest {
         fuzz_state: EvmFuzzState,
         targeted_contracts: FuzzRunIdentifiedContracts,
         failures: InvariantFailures,
-        last_call_results: Option<RawCallResult>,
         branch_runner: TestRunner,
     ) -> Self {
-        let mut fuzz_cases = vec![];
-        if last_call_results.is_none() {
-            fuzz_cases.push(FuzzedCases::new(vec![]));
-        }
         let test_data = InvariantTestData {
-            fuzz_cases,
+            fuzz_cases: vec![],
             failures,
             last_run_inputs: vec![],
             gas_report_traces: vec![],
-            last_call_results,
             line_coverage: None,
             metrics: Map::default(),
             branch_runner,
@@ -186,18 +180,13 @@ impl InvariantTest {
     }
 
     /// Whether invariant test has errors or not.
-    fn has_errors(&self) -> bool {
-        self.test_data.failures.error.is_some()
+    fn has_errors(&self, invariant: &Function) -> bool {
+        self.test_data.failures.has_failure(invariant)
     }
 
     /// Set invariant test error.
-    fn set_error(&mut self, error: InvariantFuzzError) {
-        self.test_data.failures.error = Some(error);
-    }
-
-    /// Set last invariant test call results.
-    fn set_last_call_results(&mut self, call_result: Option<RawCallResult>) {
-        self.test_data.last_call_results = call_result;
+    fn set_error(&mut self, invariant: &Function, error: InvariantFuzzError) {
+        self.test_data.failures.record_failure(invariant, error);
     }
 
     /// Set last invariant run call sequence.
@@ -336,7 +325,7 @@ impl<'a> InvariantExecutor<'a> {
         early_exit: &EarlyExit,
     ) -> Result<InvariantFuzzTestResult> {
         // Throw an error to abort test run if the invariant function accepts input params
-        if !invariant_contract.invariant_function.inputs.is_empty() {
+        if !invariant_contract.invariant_fn.inputs.is_empty() {
             return Err(eyre!("Invariant test function should have no inputs"));
         }
 
@@ -421,9 +410,10 @@ impl<'a> InvariantExecutor<'a> {
                     current_run.inputs.pop();
                     current_run.rejects += 1;
                     if current_run.rejects > self.config.max_assume_rejects {
-                        invariant_test.set_error(InvariantFuzzError::MaxAssumeRejects(
-                            self.config.max_assume_rejects,
-                        ));
+                        invariant_test.set_error(
+                            invariant_contract.invariant_fn,
+                            InvariantFuzzError::MaxAssumeRejects(self.config.max_assume_rejects),
+                        );
                         break 'stop;
                     }
                 } else {
@@ -468,7 +458,7 @@ impl<'a> InvariantExecutor<'a> {
                     });
 
                     // Determine if test can continue or should exit.
-                    let result = can_continue(
+                    let can_continue = can_continue(
                         &invariant_contract,
                         &mut invariant_test,
                         &mut current_run,
@@ -477,15 +467,13 @@ impl<'a> InvariantExecutor<'a> {
                         &state_changeset,
                     )
                     .map_err(|e| eyre!(e.to_string()))?;
-                    if !result.can_continue || current_run.depth == self.config.depth - 1 {
+                    if !can_continue || current_run.depth == self.config.depth - 1 {
                         invariant_test.set_last_run_inputs(&current_run.inputs);
                     }
                     // If test cannot continue then stop current run and exit test suite.
-                    if !result.can_continue {
+                    if !can_continue {
                         break 'stop;
                     }
-
-                    invariant_test.set_last_call_results(result.call_result);
                     current_run.depth += 1;
                 }
 
@@ -501,7 +489,9 @@ impl<'a> InvariantExecutor<'a> {
             corpus_manager.process_inputs(&current_run.inputs, current_run.new_coverage);
 
             // Call `afterInvariant` only if it is declared and test didn't fail already.
-            if invariant_contract.call_after_invariant && !invariant_test.has_errors() {
+            if invariant_contract.call_after_invariant
+                && !invariant_test.has_errors(invariant_contract.invariant_fn)
+            {
                 assert_after_invariant(
                     &invariant_contract,
                     &mut invariant_test,
@@ -516,10 +506,18 @@ impl<'a> InvariantExecutor<'a> {
             if let Some(progress) = progress {
                 // If running with progress then increment completed runs.
                 progress.inc(1);
-                // Display metrics in progress bar.
+
+                let failures = &invariant_test.test_data.failures;
+                let mut parts = Vec::new();
+                // Add failures if present
+                if !failures.errors.is_empty() {
+                    parts.push(format!("{failures}"));
+                }
+                // Add edge coverage metrics if enabled
                 if edge_coverage_enabled {
-                    progress.set_message(format!("{}", &corpus_manager.metrics));
+                    parts.push(format!("{}", corpus_manager.metrics));
                 }
+                progress.set_message(parts.join(""));
             } else if edge_coverage_enabled
                 && last_metrics_report.elapsed() > DURATION_BETWEEN_METRICS_REPORT
             {
@@ -528,7 +526,7 @@ impl<'a> InvariantExecutor<'a> {
                     "timestamp": SystemTime::now()
                         .duration_since(UNIX_EPOCH)?
                         .as_secs(),
-                    "invariant": invariant_contract.invariant_function.name,
+                    "invariant": invariant_contract.invariant_fn.name,
                     "metrics": &corpus_manager.metrics,
                 });
                 let _ = sh_println!("{}", serde_json::to_string(&metrics)?);
@@ -543,7 +541,7 @@ impl<'a> InvariantExecutor<'a> {
 
         let result = invariant_test.test_data;
         Ok(InvariantFuzzTestResult {
-            error: result.failures.error,
+            errors: result.failures.errors,
             cases: result.fuzz_cases,
             reverts: result.failures.reverts,
             last_run_inputs: result.last_run_inputs,
@@ -627,15 +625,15 @@ impl<'a> InvariantExecutor<'a> {
         // already know if we can early exit the invariant run.
         // This does not count as a fuzz run. It will just register the revert.
         let mut failures = InvariantFailures::new();
-        let last_call_results = assert_invariants(
+        invariant_preflight_check(
             invariant_contract,
             &self.config,
             &targeted_contracts,
             &self.executor,
             &[],
             &mut failures,
         )?;
-        if let Some(error) = failures.error {
+        if let Some(error) = failures.get_failure(invariant_contract.invariant_fn) {
             return Err(eyre!(error.revert_reason().unwrap_or_default()));
         }
 
@@ -646,13 +644,8 @@ impl<'a> InvariantExecutor<'a> {
             None,
             Some(&targeted_contracts),
         )?;
-        let invariant_test = InvariantTest::new(
-            fuzz_state,
-            targeted_contracts,
-            failures,
-            last_call_results,
-            self.runner.clone(),
-        );
+        let invariant_test =
+            InvariantTest::new(fuzz_state, targeted_contracts, failures, self.runner.clone());
 
         Ok((invariant_test, corpus_manager))
     }

crates/evm/evm/src/executors/invariant/replay.rs

@@ -70,7 +70,7 @@ pub fn replay_run(
     let (invariant_result, invariant_success) = call_invariant_function(
         &executor,
         invariant_contract.address,
-        invariant_contract.invariant_function.abi_encode_input(&[])?.into(),
+        invariant_contract.invariant_fn.abi_encode_input(&[])?.into(),
     )?;
     traces.push((TraceKind::Execution, invariant_result.traces.clone().unwrap()));
     logs.extend(invariant_result.logs);
@@ -92,6 +92,47 @@ pub fn replay_run(
     Ok(counterexample_sequence)
 }
 
+pub fn generate_counterexample(
+    mut executor: Executor,
+    known_contracts: &ContractsByArtifact,
+    mut ided_contracts: ContractsByAddress,
+    inputs: &[BasicTxDetails],
+    show_solidity: bool,
+) -> Result<Vec<BaseCounterExample>> {
+    // We want traces for a failed case.
+    if executor.inspector().tracer.is_none() {
+        executor.set_tracing(TraceMode::Call);
+    }
+
+    let mut counterexample_sequence = vec![];
+
+    // Replay each call from the sequence, collect logs, traces and coverage.
+    for tx in inputs {
+        let call_result = executor.transact_raw(
+            tx.sender,
+            tx.call_details.target,
+            tx.call_details.calldata.clone(),
+            U256::ZERO,
+        )?;
+
+        // Identify newly generated contracts, if they exist.
+        ided_contracts
+            .extend(load_contracts(call_result.traces.iter().map(|a| &a.arena), known_contracts));
+
+        // Create counter example to be used in failed case.
+        counterexample_sequence.push(BaseCounterExample::from_invariant_call(
+            tx.sender,
+            tx.call_details.target,
+            &tx.call_details.calldata,
+            &ided_contracts,
+            call_result.traces,
+            show_solidity,
+        ));
+    }
+
+    Ok(counterexample_sequence)
+}
+
 /// Replays the error case, shrinks the failing sequence and collects all necessary traces.
 #[expect(clippy::too_many_arguments)]
 pub fn replay_error(