feat(invariant): fuzz with values from events and return values (#7666)

* feat(invariant): scrape return values and add to fuzz dictionary

* Perist mined values between runs

* Refactor, add persistent samples

* Apply weight to collected sample values

* Add Function to BasicTxDetails (if has outputs), to be used for decoding. Decode results and persist per types. Use typed samples when fuzzing from state.

* Fix clippy and fmt

* Use prop-perturb take 1

* Decode logs using target abi, populate type samples

* Fmt

* Fix clippy, add calldetails type

* Fix fmt test

* Insert call sample once

* Proper function naming

* Generate state values bias using strategy

* Add BasicTxDetails and CallTargetDetails struct, add Function always to call details and use it to generate counterexample

* Tests cleanup

* Code cleanup

* Move args in CallDetails

* Fallback to old impl if we are not able to decode logs

* Refactor collect values fn

* Get abi from FuzzedContracts

* Lookup function from identified target abi.

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

@@ -205,11 +205,16 @@ impl<'a> InvariantExecutor<'a> {
             let mut assume_rejects_counter = 0;
 
             while current_run < self.config.depth {
-                let (sender, (address, calldata)) = inputs.last().expect("no input generated");
+                let tx = inputs.last().expect("no input generated");
 
                 // Execute call from the randomly generated sequence and commit state changes.
                 let call_result = executor
-                    .call_raw_committing(*sender, *address, calldata.clone(), U256::ZERO)
+                    .call_raw_committing(
+                        tx.sender,
+                        tx.call_details.target,
+                        tx.call_details.calldata.clone(),
+                        U256::ZERO,
+                    )
                     .expect("could not make raw evm call");
 
                 if call_result.result.as_ref() == MAGIC_ASSUME {
@@ -226,7 +231,16 @@ impl<'a> InvariantExecutor<'a> {
                     let mut state_changeset =
                         call_result.state_changeset.to_owned().expect("no changesets");
 
-                    collect_data(&mut state_changeset, sender, &call_result, &fuzz_state);
+                    if !&call_result.reverted {
+                        collect_data(
+                            &mut state_changeset,
+                            &targeted_contracts,
+                            tx,
+                            &call_result,
+                            &fuzz_state,
+                            self.config.depth,
+                        );
+                    }
 
                     // Collect created contracts and add to fuzz targets only if targeted contracts
                     // are updatable.
@@ -244,7 +258,7 @@ impl<'a> InvariantExecutor<'a> {
                     }
 
                     fuzz_runs.push(FuzzCase {
-                        calldata: calldata.clone(),
+                        calldata: tx.call_details.calldata.clone(),
                         gas: call_result.gas_used,
                         stipend: call_result.stipend,
                     });
@@ -639,27 +653,39 @@ impl<'a> InvariantExecutor<'a> {
 /// randomly generated addresses.
 fn collect_data(
     state_changeset: &mut HashMap<Address, revm::primitives::Account>,
-    sender: &Address,
+    fuzzed_contracts: &FuzzRunIdentifiedContracts,
+    tx: &BasicTxDetails,
     call_result: &RawCallResult,
     fuzz_state: &EvmFuzzState,
+    run_depth: u32,
 ) {
     // Verify it has no code.
     let mut has_code = false;
-    if let Some(Some(code)) = state_changeset.get(sender).map(|account| account.info.code.as_ref())
+    if let Some(Some(code)) =
+        state_changeset.get(&tx.sender).map(|account| account.info.code.as_ref())
     {
         has_code = !code.is_empty();
     }
 
     // We keep the nonce changes to apply later.
     let mut sender_changeset = None;
     if !has_code {
-        sender_changeset = state_changeset.remove(sender);
+        sender_changeset = state_changeset.remove(&tx.sender);
     }
 
-    fuzz_state.collect_state_from_call(&call_result.logs, &*state_changeset);
+    // Collect values from fuzzed call result and add them to fuzz dictionary.
+    let (fuzzed_contract_abi, fuzzed_function) = fuzzed_contracts.fuzzed_artifacts(tx);
+    fuzz_state.collect_values_from_call(
+        fuzzed_contract_abi.as_ref(),
+        fuzzed_function.as_ref(),
+        &call_result.result,
+        &call_result.logs,
+        &*state_changeset,
+        run_depth,
+    );
 
     // Re-add changes
     if let Some(changed) = sender_changeset {
-        state_changeset.insert(*sender, changed);
+        state_changeset.insert(tx.sender, changed);
     }
 }

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

@@ -35,9 +35,13 @@ pub fn replay_run(
     let mut counterexample_sequence = vec![];
 
     // Replay each call from the sequence, collect logs, traces and coverage.
-    for (sender, (addr, bytes)) in inputs.iter() {
-        let call_result =
-            executor.call_raw_committing(*sender, *addr, bytes.clone(), U256::ZERO)?;
+    for tx in inputs.iter() {
+        let call_result = executor.call_raw_committing(
+            tx.sender,
+            tx.call_details.target,
+            tx.call_details.calldata.clone(),
+            U256::ZERO,
+        )?;
         logs.extend(call_result.logs);
         traces.push((TraceKind::Execution, call_result.traces.clone().unwrap()));
 
@@ -57,9 +61,9 @@ pub fn replay_run(
 
         // Create counter example to be used in failed case.
         counterexample_sequence.push(BaseCounterExample::create(
-            *sender,
-            *addr,
-            bytes,
+            tx.sender,
+            tx.call_details.target,
+            &tx.call_details.calldata,
             &ided_contracts,
             call_result.traces,
         ));

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

@@ -120,9 +120,13 @@ fn check_sequence(
     let mut sequence_failed = false;
     // Apply the shrinked candidate sequence.
     for call_index in sequence {
-        let (sender, (addr, bytes)) = &calls[call_index];
-        let call_result =
-            executor.call_raw_committing(*sender, *addr, bytes.clone(), U256::ZERO)?;
+        let tx = &calls[call_index];
+        let call_result = executor.call_raw_committing(
+            tx.sender,
+            tx.call_details.target,
+            tx.call_details.calldata.clone(),
+            U256::ZERO,
+        )?;
         if call_result.reverted && failed_case.fail_on_revert {
             // Candidate sequence fails test.
             // We don't have to apply remaining calls to check sequence.

crates/evm/fuzz/src/inspector.rs

@@ -82,16 +82,14 @@ impl Fuzzer {
                 !call_generator.used
             {
                 // There's only a 30% chance that an override happens.
-                if let Some((sender, (contract, input))) =
-                    call_generator.next(call.context.caller, call.contract)
-                {
-                    *call.input = input.0;
-                    call.context.caller = sender;
-                    call.contract = contract;
+                if let Some(tx) = call_generator.next(call.context.caller, call.contract) {
+                    *call.input = tx.call_details.calldata.0;
+                    call.context.caller = tx.sender;
+                    call.contract = tx.call_details.target;
 
                     // TODO: in what scenarios can the following be problematic
-                    call.context.code_address = contract;
-                    call.context.address = contract;
+                    call.context.code_address = tx.call_details.target;
+                    call.context.address = tx.call_details.target;
 
                     call_generator.used = true;
                 }

crates/evm/fuzz/src/invariant/call_override.rs

@@ -1,5 +1,5 @@
-use super::BasicTxDetails;
-use alloy_primitives::{Address, Bytes};
+use super::{BasicTxDetails, CallDetails};
+use alloy_primitives::Address;
 use parking_lot::{Mutex, RwLock};
 use proptest::{
     option::weighted,
@@ -17,7 +17,7 @@ pub struct RandomCallGenerator {
     /// Runner that will generate the call from the strategy.
     pub runner: Arc<Mutex<TestRunner>>,
     /// Strategy to be used to generate calls from `target_reference`.
-    pub strategy: SBoxedStrategy<Option<(Address, Bytes)>>,
+    pub strategy: SBoxedStrategy<Option<CallDetails>>,
     /// Reference to which contract we want a fuzzed calldata from.
     pub target_reference: Arc<RwLock<Address>>,
     /// Flag to know if a call has been overridden. Don't allow nesting for now.
@@ -33,7 +33,7 @@ impl RandomCallGenerator {
     pub fn new(
         test_address: Address,
         runner: TestRunner,
-        strategy: SBoxedStrategy<(Address, Bytes)>,
+        strategy: SBoxedStrategy<CallDetails>,
         target_reference: Arc<RwLock<Address>>,
     ) -> Self {
         let strategy = weighted(0.9, strategy).sboxed();
@@ -71,7 +71,7 @@ impl RandomCallGenerator {
             )
         } else {
             // TODO: Do we want it to be 80% chance only too ?
-            let new_caller = original_target;
+            let sender = original_target;
 
             // Set which contract we mostly (80% chance) want to generate calldata from.
             *self.target_reference.write() = original_caller;
@@ -82,7 +82,7 @@ impl RandomCallGenerator {
                 .new_tree(&mut self.runner.lock())
                 .unwrap()
                 .current()
-                .map(|(new_target, calldata)| (new_caller, (new_target, calldata)));
+                .map(|call_details| BasicTxDetails { sender, call_details });
 
             self.last_sequence.write().push(choice.clone());
             choice

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

@@ -26,10 +26,37 @@ impl FuzzRunIdentifiedContracts {
     pub fn new(targets: TargetedContracts, is_updatable: bool) -> Self {
         Self { targets: Arc::new(Mutex::new(targets)), is_updatable }
     }
+
+    /// Returns fuzzed contract abi and fuzzed function from address and provided calldata.
+    /// Used to decode return values and logs in order to add values into fuzz dictionary.
+    pub fn fuzzed_artifacts(&self, tx: &BasicTxDetails) -> (Option<JsonAbi>, Option<Function>) {
+        match self.targets.lock().get(&tx.call_details.target) {
+            Some((_, abi, _)) => (
+                Some(abi.to_owned()),
+                abi.functions().find(|f| f.selector() == tx.call_details.calldata[..4]).cloned(),
+            ),
+            None => (None, None),
+        }
+    }
 }
 
-/// (Sender, (TargetContract, Calldata))
-pub type BasicTxDetails = (Address, (Address, Bytes));
+/// Details of a transaction generated by invariant strategy for fuzzing a target.
+#[derive(Clone, Debug)]
+pub struct BasicTxDetails {
+    // Transaction sender address.
+    pub sender: Address,
+    // Transaction call details.
+    pub call_details: CallDetails,
+}
+
+/// Call details of a transaction generated to fuzz invariant target.
+#[derive(Clone, Debug)]
+pub struct CallDetails {
+    // Address of target contract.
+    pub target: Address,
+    // The data of the transaction.
+    pub calldata: Bytes,
+}
 
 /// Test contract which is testing its invariants.
 #[derive(Clone, Debug)]

crates/evm/fuzz/src/strategies/invariants.rs

@@ -1,11 +1,11 @@
 use super::{fuzz_calldata, fuzz_param_from_state};
 use crate::{
-    invariant::{BasicTxDetails, FuzzRunIdentifiedContracts, SenderFilters},
+    invariant::{BasicTxDetails, CallDetails, FuzzRunIdentifiedContracts, SenderFilters},
     strategies::{fuzz_calldata_from_state, fuzz_param, EvmFuzzState},
     FuzzFixtures,
 };
 use alloy_json_abi::{Function, JsonAbi};
-use alloy_primitives::{Address, Bytes};
+use alloy_primitives::Address;
 use parking_lot::RwLock;
 use proptest::prelude::*;
 use std::{rc::Rc, sync::Arc};
@@ -16,7 +16,7 @@ pub fn override_call_strat(
     contracts: FuzzRunIdentifiedContracts,
     target: Arc<RwLock<Address>>,
     fuzz_fixtures: FuzzFixtures,
-) -> SBoxedStrategy<(Address, Bytes)> {
+) -> SBoxedStrategy<CallDetails> {
     let contracts_ref = contracts.targets.clone();
     proptest::prop_oneof![
         80 => proptest::strategy::LazyJust::new(move || *target.read()),
@@ -101,6 +101,7 @@ fn generate_call(
                 (sender, contract)
             })
         })
+        .prop_map(|(sender, call_details)| BasicTxDetails { sender, call_details })
         .boxed()
 }
 
@@ -166,9 +167,9 @@ fn select_random_function(
 pub fn fuzz_contract_with_calldata(
     fuzz_state: &EvmFuzzState,
     fuzz_fixtures: &FuzzFixtures,
-    contract: Address,
+    target: Address,
     func: Function,
-) -> impl Strategy<Value = (Address, Bytes)> {
+) -> impl Strategy<Value = CallDetails> {
     // We need to compose all the strategies generated for each parameter in all possible
     // combinations.
     // `prop_oneof!` / `TupleUnion` `Arc`s for cheap cloning.
@@ -179,6 +180,6 @@ pub fn fuzz_contract_with_calldata(
     ]
     .prop_map(move |calldata| {
         trace!(input=?calldata);
-        (contract, calldata)
+        CallDetails { target, calldata }
     })
 }

crates/evm/fuzz/src/strategies/param.rs

@@ -124,13 +124,27 @@ pub fn fuzz_param_from_state(
     // Value strategy that uses the state.
     let value = || {
         let state = state.clone();
-        // Use `Index` instead of `Selector` to not iterate over the entire dictionary.
-        any::<prop::sample::Index>().prop_map(move |index| {
-            let state = state.dictionary_read();
-            let values = state.values();
-            let index = index.index(values.len());
-            *values.iter().nth(index).unwrap()
-        })
+        let param = param.clone();
+        // Generate a bias and use it to pick samples or non-persistent values (50 / 50).
+        // Use `Index` instead of `Selector` when selecting a value to avoid iterating over the
+        // entire dictionary.
+        ((0..100).prop_flat_map(Just), any::<prop::sample::Index>()).prop_map(
+            move |(bias, index)| {
+                let state = state.dictionary_read();
+                let values = match bias {
+                    x if x < 50 => {
+                        if let Some(sample_values) = state.samples(param.clone()) {
+                            sample_values
+                        } else {
+                            state.values()
+                        }
+                    }
+                    _ => state.values(),
+                };
+                let index = index.index(values.len());
+                *values.iter().nth(index).unwrap()
+            },
+        )
     };
 
     // Convert the value based on the parameter type