feat(invariant): assert all invariants

Author: grandizzy

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

@@ -220,6 +220,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,5 +1,6 @@
 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::{
@@ -8,6 +9,7 @@ use foundry_evm_core::{
 };
 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)]
@@ -17,16 +19,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(())
     }
 }
 
@@ -75,7 +101,8 @@ pub struct FailedInvariantCaseData {
 impl FailedInvariantCaseData {
     pub fn new<FEN: FoundryEvmNetwork>(
         invariant_contract: &InvariantContract<'_>,
-        invariant_config: &InvariantConfig,
+        shrink_run_limit: u32,
+        fail_on_revert: bool,
         targeted_contracts: &FuzzRunIdentifiedContracts,
         calldata: &[BasicTxDetails],
         call_result: RawCallResult<FEN>,
@@ -95,7 +122,7 @@ impl FailedInvariantCaseData {
                 revert_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 {
@@ -108,8 +135,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,
             assertion_failure: false,
         }
     }

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

@@ -5,7 +5,8 @@ use crate::{
     },
     inspectors::Fuzzer,
 };
-use alloy_primitives::{Address, Bytes, FixedBytes, I256, Selector, U256, map::AddressMap};
+use alloy_json_abi::Function;
+use alloy_primitives::{Address, Bytes, FixedBytes, I256, Selector, U256, map::{AddressMap, HashMap}};
 use alloy_sol_types::{SolCall, sol};
 use eyre::{ContextCompat, Result, eyre};
 use foundry_common::{
@@ -34,7 +35,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, did_fail_on_assert};
+use result::{assert_after_invariant, assert_invariants, can_continue, did_fail_on_assert, invariant_preflight_check};
 use revm::{context::Block, state::Account};
 use serde::{Deserialize, Serialize};
 use serde_json::json;
@@ -49,7 +50,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;
 pub use result::InvariantFuzzTestResult;
@@ -266,12 +267,8 @@ impl<FEN: FoundryEvmNetwork> InvariantTest<FEN> {
         last_call_results: Option<RawCallResult<FEN>>,
         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![],
@@ -291,13 +288,13 @@ impl<FEN: FoundryEvmNetwork> InvariantTest<FEN> {
     }
 
     /// Whether invariant test has errors or not.
-    const 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);
+    fn set_error(&mut self, invariant: &Function, error: InvariantFuzzError) {
+        self.test_data.failures.record_failure(invariant, error);
     }
 
     /// Set last invariant test call results.
@@ -455,7 +452,7 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
         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"));
         }
 
@@ -534,9 +531,10 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
                     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 {
@@ -602,7 +600,7 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
                                 || is_last_call
                         };
 
-                    let result = if should_check_invariant {
+                    let can_continue_result = if should_check_invariant {
                         can_continue(
                             &invariant_contract,
                             &mut invariant_test,
@@ -621,7 +619,8 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
                         {
                             let case_data = error::FailedInvariantCaseData::new(
                                 &invariant_contract,
-                                &self.config,
+                                self.config.shrink_run_limit,
+                                self.config.fail_on_revert,
                                 &invariant_test.targeted_contracts,
                                 &current_run.inputs,
                                 call_result,
@@ -630,12 +629,12 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
                             .with_assertion_failure(assertion_failure);
                             invariant_test.test_data.failures.revert_reason =
                                 Some(case_data.revert_reason.clone());
-                            invariant_test.test_data.failures.error = Some(if assertion_failure {
+                            invariant_test.set_error(invariant_contract.invariant_fn, if assertion_failure {
                                 InvariantFuzzError::BrokenInvariant(case_data)
                             } else {
                                 InvariantFuzzError::Revert(case_data)
                             });
-                            result::RichInvariantResults::new(false, None)
+                            false
                         } else if call_result.reverted
                             && !invariant_contract.is_optimization()
                             && !self.config.has_delay()
@@ -644,33 +643,32 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
                             // preserve their warp/roll contribution when building the final
                             // counterexample.
                             current_run.inputs.pop();
-                            result::RichInvariantResults::new(true, None)
+                            true
                         } else {
-                            result::RichInvariantResults::new(true, None)
+                            true
                         }
                     };
 
-                    if !result.can_continue || current_run.depth == self.config.depth - 1 {
+                    if !can_continue_result || 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_result {
                         let reason = invariant_test
                             .test_data
                             .failures
-                            .error
-                            .as_ref()
+                            .errors
+                            .values()
+                            .next()
                             .and_then(|e| e.revert_reason())
                             .unwrap_or_default();
                         failure_metrics.record_failure(
-                            &invariant_contract.invariant_function.name,
+                            &invariant_contract.invariant_fn.name,
                             invariant_contract.name,
                             &reason,
                         );
                         break 'stop;
                     }
-
-                    invariant_test.set_last_call_results(result.call_result);
                     current_run.depth += 1;
                 }
 
@@ -696,7 +694,9 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
             );
 
             // 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)
+            {
                 let success = assert_after_invariant(
                     &invariant_contract,
                     &mut invariant_test,
@@ -708,12 +708,13 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
                     let reason = invariant_test
                         .test_data
                         .failures
-                        .error
-                        .as_ref()
+                        .errors
+                        .values()
+                        .next()
                         .and_then(|e| e.revert_reason())
                         .unwrap_or_default();
                     failure_metrics.record_failure(
-                        &invariant_contract.invariant_function.name,
+                        &invariant_contract.invariant_fn.name,
                         invariant_contract.name,
                         &reason,
                     );
@@ -746,7 +747,7 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
                 // Display corpus metrics inline as JSON.
                 let metrics = build_invariant_progress_json(
                     SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(),
-                    &invariant_contract.invariant_function.name,
+                    &invariant_contract.invariant_fn.name,
                     &corpus_manager.metrics,
                     invariant_test.test_data.optimization_best_value,
                     throughput,
@@ -765,7 +766,7 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
 
         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,
@@ -825,15 +826,15 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
         // 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()));
         }
 
@@ -879,7 +880,7 @@ impl<'a, FEN: FoundryEvmNetwork> InvariantExecutor<'a, FEN> {
             fuzz_state,
             targeted_contracts,
             failures,
-            last_call_results,
+            None,
             self.runner.clone(),
         );
 

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

@@ -4,7 +4,7 @@ use crate::executors::{
     invariant::shrink::{shrink_sequence, shrink_sequence_value},
 };
 use alloy_dyn_abi::JsonAbiExt;
-use alloy_primitives::{I256, Log, map::HashMap};
+use alloy_primitives::{I256, Log, U256, map::HashMap};
 use eyre::Result;
 use foundry_common::{ContractsByAddress, ContractsByArtifact};
 use foundry_config::InvariantConfig;
@@ -69,7 +69,7 @@ pub fn replay_run<FEN: FoundryEvmNetwork>(
     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);
@@ -91,6 +91,41 @@ pub fn replay_run<FEN: FoundryEvmNetwork>(
     Ok(counterexample_sequence)
 }
 
+pub fn generate_counterexample<FEN: FoundryEvmNetwork>(
+    mut executor: Executor<FEN>,
+    known_contracts: &ContractsByArtifact,
+    mut ided_contracts: ContractsByAddress,
+    inputs: &[BasicTxDetails],
+    show_solidity: bool,
+) -> Result<Vec<BaseCounterExample>> {
+    if executor.inspector().tracer.is_none() {
+        executor.set_tracing(TraceMode::Call);
+    }
+
+    let mut counterexample_sequence = vec![];
+
+    for tx in inputs {
+        let call_result = executor.transact_raw(
+            tx.sender,
+            tx.call_details.target,
+            tx.call_details.calldata.clone(),
+            U256::ZERO,
+        )?;
+
+        ided_contracts
+            .extend(load_contracts(call_result.traces.iter().map(|a| &a.arena), known_contracts));
+
+        counterexample_sequence.push(BaseCounterExample::from_invariant_call(
+            tx,
+            &ided_contracts,
+            call_result.traces,
+            show_solidity,
+        ));
+    }
+
+    Ok(counterexample_sequence)
+}
+
 /// Replays and shrinks a call sequence, collecting logs and traces.
 ///
 /// For check mode (target_value=None): shrinks to find shortest failing sequence.