fix(invariant): weight invariant selectors by number of selectors (fo…

…undry-rs#8176)

* fix(invariant): weight invariant selectors by number of selectors

- Consolidate FuzzRunIdentifiedContracts logic
- add function to flatten contracts function in order to be used by strategy
- test

* Changes after review: cleanup

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

@@ -18,7 +18,7 @@ use foundry_evm_fuzz::{
         ArtifactFilters, BasicTxDetails, FuzzRunIdentifiedContracts, InvariantContract,
         RandomCallGenerator, SenderFilters, TargetedContracts,
     },
-    strategies::{collect_created_contracts, invariant_strat, override_call_strat, EvmFuzzState},
+    strategies::{invariant_strat, override_call_strat, EvmFuzzState},
     FuzzCase, FuzzFixtures, FuzzedCases,
 };
 use foundry_evm_traces::CallTraceArena;
@@ -251,17 +251,14 @@ impl<'a> InvariantExecutor<'a> {
 
                     // Collect created contracts and add to fuzz targets only if targeted contracts
                     // are updatable.
-                    if targeted_contracts.is_updatable {
-                        if let Err(error) = collect_created_contracts(
-                            &state_changeset,
-                            self.project_contracts,
-                            self.setup_contracts,
-                            &self.artifact_filters,
-                            &targeted_contracts,
-                            &mut created_contracts,
-                        ) {
-                            warn!(target: "forge::test", "{error}");
-                        }
+                    if let Err(error) = &targeted_contracts.collect_created_contracts(
+                        &state_changeset,
+                        self.project_contracts,
+                        self.setup_contracts,
+                        &self.artifact_filters,
+                        &mut created_contracts,
+                    ) {
+                        warn!(target: "forge::test", "{error}");
                     }
 
                     fuzz_runs.push(FuzzCase {
@@ -319,12 +316,7 @@ impl<'a> InvariantExecutor<'a> {
             }
 
             // We clear all the targeted contracts created during this run.
-            if !created_contracts.is_empty() {
-                let mut writable_targeted = targeted_contracts.targets.lock();
-                for addr in created_contracts.iter() {
-                    writable_targeted.remove(addr);
-                }
-            }
+            let _ = &targeted_contracts.clear_created_contracts(created_contracts);
 
             if gas_report_traces.borrow().len() < self.config.gas_report_samples as usize {
                 gas_report_traces.borrow_mut().push(run_traces);

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

@@ -25,8 +25,7 @@ impl ArtifactFilters {
     /// Gets all the targeted functions from `artifact`. Returns error, if selectors do not match
     /// the `artifact`.
     ///
-    /// An empty vector means that it targets any mutable function. See `select_random_function` for
-    /// more.
+    /// An empty vector means that it targets any mutable function.
     pub fn get_targeted_functions(
         &self,
         artifact: &ArtifactId,

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

@@ -8,6 +8,8 @@ pub use call_override::RandomCallGenerator;
 
 mod filters;
 pub use filters::{ArtifactFilters, SenderFilters};
+use foundry_common::{ContractsByAddress, ContractsByArtifact};
+use foundry_evm_core::utils::StateChangeset;
 
 pub type TargetedContracts = BTreeMap<Address, (String, JsonAbi, Vec<Function>)>;
 
@@ -44,6 +46,93 @@ impl FuzzRunIdentifiedContracts {
         };
         f(abi, abi_f);
     }
+
+    /// Returns flatten target contract address and functions to be fuzzed.
+    /// Includes contract targeted functions if specified, else all mutable contract functions.
+    pub fn fuzzed_functions(&self) -> Vec<(Address, Function)> {
+        let mut fuzzed_functions = vec![];
+        for (contract, (_, abi, functions)) in self.targets.lock().iter() {
+            if !abi.functions.is_empty() {
+                for function in abi_fuzzed_functions(abi, functions) {
+                    fuzzed_functions.push((*contract, function.clone()));
+                }
+            }
+        }
+        fuzzed_functions
+    }
+
+    /// If targets are updatable, collect all contracts created during an invariant run (which
+    /// haven't been discovered yet).
+    pub fn collect_created_contracts(
+        &self,
+        state_changeset: &StateChangeset,
+        project_contracts: &ContractsByArtifact,
+        setup_contracts: &ContractsByAddress,
+        artifact_filters: &ArtifactFilters,
+        created_contracts: &mut Vec<Address>,
+    ) -> eyre::Result<()> {
+        if self.is_updatable {
+            let mut targets = self.targets.lock();
+            for (address, account) in state_changeset {
+                if setup_contracts.contains_key(address) {
+                    continue;
+                }
+                if !account.is_touched() {
+                    continue;
+                }
+                let Some(code) = &account.info.code else {
+                    continue;
+                };
+                if code.is_empty() {
+                    continue;
+                }
+                let Some((artifact, contract)) =
+                    project_contracts.find_by_deployed_code(code.original_byte_slice())
+                else {
+                    continue;
+                };
+                let Some(functions) =
+                    artifact_filters.get_targeted_functions(artifact, &contract.abi)?
+                else {
+                    continue;
+                };
+                created_contracts.push(*address);
+                targets.insert(*address, (artifact.name.clone(), contract.abi.clone(), functions));
+            }
+        }
+        Ok(())
+    }
+
+    /// Clears targeted contracts created during an invariant run.
+    pub fn clear_created_contracts(&self, created_contracts: Vec<Address>) {
+        if !created_contracts.is_empty() {
+            let mut targets = self.targets.lock();
+            for addr in created_contracts.iter() {
+                targets.remove(addr);
+            }
+        }
+    }
+}
+
+/// Helper to retrieve functions to fuzz for specified abi.
+/// Returns specified targeted functions if any, else mutable abi functions.
+pub(crate) fn abi_fuzzed_functions(
+    abi: &JsonAbi,
+    targeted_functions: &[Function],
+) -> Vec<Function> {
+    if !targeted_functions.is_empty() {
+        targeted_functions.to_vec()
+    } else {
+        abi.functions()
+            .filter(|&func| {
+                !matches!(
+                    func.state_mutability,
+                    alloy_json_abi::StateMutability::Pure | alloy_json_abi::StateMutability::View
+                )
+            })
+            .cloned()
+            .collect()
+    }
 }
 
 /// Details of a transaction generated by invariant strategy for fuzzing a target.

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

@@ -1,10 +1,13 @@
 use super::{fuzz_calldata, fuzz_param_from_state};
 use crate::{
-    invariant::{BasicTxDetails, CallDetails, FuzzRunIdentifiedContracts, SenderFilters},
+    invariant::{
+        abi_fuzzed_functions, BasicTxDetails, CallDetails, FuzzRunIdentifiedContracts,
+        SenderFilters,
+    },
     strategies::{fuzz_calldata_from_state, fuzz_param, EvmFuzzState},
     FuzzFixtures,
 };
-use alloy_json_abi::{Function, JsonAbi};
+use alloy_json_abi::Function;
 use alloy_primitives::Address;
 use parking_lot::RwLock;
 use proptest::prelude::*;
@@ -37,7 +40,8 @@ pub fn override_call_strat(
                 let (_, contract_specs) = contracts.iter().nth(rand_index).unwrap();
                 contract_specs
             });
-            select_random_function(abi, functions)
+            let fuzzed_functions = abi_fuzzed_functions(abi, functions);
+            any::<prop::sample::Index>().prop_map(move |index| index.get(&fuzzed_functions).clone())
         };
 
         func.prop_flat_map(move |func| {
@@ -64,27 +68,18 @@ pub fn invariant_strat(
     fuzz_fixtures: FuzzFixtures,
 ) -> impl Strategy<Value = BasicTxDetails> {
     let senders = Rc::new(senders);
-    any::<prop::sample::Selector>()
-        .prop_flat_map(move |selector| {
-            let (contract, func) = {
-                let contracts = contracts.targets.lock();
-                let contracts =
-                    contracts.iter().filter(|(_, (_, abi, _))| !abi.functions.is_empty());
-                let (&contract, (_, abi, functions)) = selector.select(contracts);
-
-                let func = select_random_function(abi, functions);
-                (contract, func)
-            };
-
-            let senders = senders.clone();
-            let fuzz_state = fuzz_state.clone();
-            let fuzz_fixtures = fuzz_fixtures.clone();
-            func.prop_flat_map(move |func| {
-                let sender = select_random_sender(&fuzz_state, senders.clone(), dictionary_weight);
-                let contract =
-                    fuzz_contract_with_calldata(&fuzz_state, &fuzz_fixtures, contract, func);
-                (sender, contract)
-            })
+    any::<prop::sample::Index>()
+        .prop_flat_map(move |index| {
+            let (target_address, target_function) =
+                index.get(&contracts.fuzzed_functions()).clone();
+            let sender = select_random_sender(&fuzz_state, senders.clone(), dictionary_weight);
+            let call_details = fuzz_contract_with_calldata(
+                &fuzz_state,
+                &fuzz_fixtures,
+                target_address,
+                target_function,
+            );
+            (sender, call_details)
         })
         .prop_map(|(sender, call_details)| BasicTxDetails { sender, call_details })
 }
@@ -112,30 +107,6 @@ fn select_random_sender(
     }
 }
 
-/// Strategy to select a random mutable function from the abi.
-///
-/// If `targeted_functions` is not empty, select one from it. Otherwise, take any
-/// of the available abi functions.
-fn select_random_function(
-    abi: &JsonAbi,
-    targeted_functions: &[Function],
-) -> impl Strategy<Value = Function> {
-    let functions = if !targeted_functions.is_empty() {
-        targeted_functions.to_vec()
-    } else {
-        abi.functions()
-            .filter(|&func| {
-                !matches!(
-                    func.state_mutability,
-                    alloy_json_abi::StateMutability::Pure | alloy_json_abi::StateMutability::View
-                )
-            })
-            .cloned()
-            .collect()
-    };
-    any::<prop::sample::Index>().prop_map(move |index| index.get(&functions).clone())
-}
-
 /// Given a function, it returns a proptest strategy which generates valid abi-encoded calldata
 /// for that function's input types.
 pub fn fuzz_contract_with_calldata(

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

@@ -11,7 +11,7 @@ mod calldata;
 pub use calldata::{fuzz_calldata, fuzz_calldata_from_state};
 
 mod state;
-pub use state::{collect_created_contracts, EvmFuzzState};
+pub use state::EvmFuzzState;
 
 mod invariants;
 pub use invariants::{fuzz_contract_with_calldata, invariant_strat, override_call_strat};

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

@@ -1,8 +1,7 @@
-use crate::invariant::{ArtifactFilters, BasicTxDetails, FuzzRunIdentifiedContracts};
+use crate::invariant::{BasicTxDetails, FuzzRunIdentifiedContracts};
 use alloy_dyn_abi::{DynSolType, DynSolValue, EventExt, FunctionExt};
 use alloy_json_abi::{Function, JsonAbi};
 use alloy_primitives::{Address, Bytes, Log, B256, U256};
-use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
 use foundry_config::FuzzDictionaryConfig;
 use foundry_evm_core::utils::StateChangeset;
 use indexmap::IndexSet;
@@ -381,43 +380,3 @@ impl FuzzDictionary {
         );
     }
 }
-
-/// Collects all created contracts from a StateChangeset which haven't been discovered yet. Stores
-/// them at `targeted_contracts` and `created_contracts`.
-pub fn collect_created_contracts(
-    state_changeset: &StateChangeset,
-    project_contracts: &ContractsByArtifact,
-    setup_contracts: &ContractsByAddress,
-    artifact_filters: &ArtifactFilters,
-    targeted_contracts: &FuzzRunIdentifiedContracts,
-    created_contracts: &mut Vec<Address>,
-) -> eyre::Result<()> {
-    let mut writable_targeted = targeted_contracts.targets.lock();
-    for (address, account) in state_changeset {
-        if setup_contracts.contains_key(address) {
-            continue;
-        }
-        if !account.is_touched() {
-            continue;
-        }
-        let Some(code) = &account.info.code else {
-            continue;
-        };
-        if code.is_empty() {
-            continue;
-        }
-        let Some((artifact, contract)) =
-            project_contracts.find_by_deployed_code(code.original_byte_slice())
-        else {
-            continue;
-        };
-        let Some(functions) = artifact_filters.get_targeted_functions(artifact, &contract.abi)?
-        else {
-            continue;
-        };
-        created_contracts.push(*address);
-        writable_targeted
-            .insert(*address, (artifact.name.clone(), contract.abi.clone(), functions));
-    }
-    Ok(())
-}

crates/forge/tests/it/invariant.rs

@@ -264,6 +264,10 @@ async fn test_invariant() {
                     ),
                     ("invariant_success()", true, None, None, None),
                 ],
+            ),
+            (
+                "default/fuzz/invariant/common/InvariantSelectorsWeight.t.sol:InvariantSelectorsWeightTest",
+                vec![("invariant_selectors_weight()", true, None, None, None)],
             )
         ]),
     );
@@ -765,3 +769,26 @@ async fn test_invariant_after_invariant() {
         )]),
     );
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_invariant_selectors_weight() {
+    let mut opts = TEST_DATA_DEFAULT.test_opts.clone();
+    opts.fuzz.seed = Some(U256::from(119u32));
+
+    let filter = Filter::new(".*", ".*", ".*fuzz/invariant/common/InvariantSelectorsWeight.t.sol");
+    let mut runner = TEST_DATA_DEFAULT.runner();
+    runner.test_options = opts.clone();
+    runner.test_options.invariant.runs = 1;
+    runner.test_options.invariant.depth = 30;
+    runner.test_options.invariant.failure_persist_dir =
+        Some(tempfile::tempdir().unwrap().into_path());
+
+    let results = runner.test_collect(&filter);
+    assert_multiple(
+        &results,
+        BTreeMap::from([(
+            "default/fuzz/invariant/common/InvariantSelectorsWeight.t.sol:InvariantSelectorsWeightTest",
+            vec![("invariant_selectors_weight()", true, None, None, None)],
+        )]),
+    )
+}