perf: new-type TargetedContracts

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

@@ -80,7 +80,7 @@ impl FailedInvariantCaseData {
     ) -> Self {
         // Collect abis of fuzzed and invariant contracts to decode custom error.
         let revert_reason = RevertDecoder::new()
-            .with_abis(targeted_contracts.targets.lock().iter().map(|(_, (_, abi, _))| abi))
+            .with_abis(targeted_contracts.targets.lock().iter().map(|(_, c)| &c.abi))
             .with_abi(invariant_contract.abi)
             .decode(call_result.result.as_ref(), Some(call_result.exit_reason));
 

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

@@ -7,16 +7,13 @@ use alloy_sol_types::{sol, SolCall};
 use eyre::{eyre, ContextCompat, Result};
 use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
 use foundry_config::InvariantConfig;
-use foundry_evm_core::{
-    constants::{
-        CALLER, CHEATCODE_ADDRESS, DEFAULT_CREATE2_DEPLOYER, HARDHAT_CONSOLE_ADDRESS, MAGIC_ASSUME,
-    },
-    utils::get_function,
+use foundry_evm_core::constants::{
+    CALLER, CHEATCODE_ADDRESS, DEFAULT_CREATE2_DEPLOYER, HARDHAT_CONSOLE_ADDRESS, MAGIC_ASSUME,
 };
 use foundry_evm_fuzz::{
     invariant::{
         ArtifactFilters, BasicTxDetails, FuzzRunIdentifiedContracts, InvariantContract,
-        RandomCallGenerator, SenderFilters, TargetedContracts,
+        RandomCallGenerator, SenderFilters, TargetedContract, TargetedContracts,
     },
     strategies::{invariant_strat, override_call_strat, EvmFuzzState},
     FuzzCase, FuzzFixtures, FuzzedCases,
@@ -31,7 +28,7 @@ use proptest::{
 use result::{assert_after_invariant, assert_invariants, can_continue};
 use revm::primitives::HashMap;
 use shrink::shrink_sequence;
-use std::{cell::RefCell, collections::BTreeMap, sync::Arc};
+use std::{cell::RefCell, collections::btree_map::Entry, sync::Arc};
 
 mod error;
 pub use error::{InvariantFailures, InvariantFuzzError};
@@ -517,7 +514,7 @@ impl<'a> InvariantExecutor<'a> {
         let excluded =
             self.call_sol_default(to, &IInvariantTest::excludeContractsCall {}).excludedContracts;
 
-        let mut contracts: TargetedContracts = self
+        let contracts = self
             .setup_contracts
             .iter()
             .filter(|&(addr, (identifier, _))| {
@@ -528,8 +525,11 @@ impl<'a> InvariantExecutor<'a> {
                     (excluded.is_empty() || !excluded.contains(addr)) &&
                     self.artifact_filters.matches(identifier)
             })
-            .map(|(addr, (identifier, abi))| (*addr, (identifier.clone(), abi.clone(), vec![])))
+            .map(|(addr, (identifier, abi))| {
+                (*addr, TargetedContract::new(identifier.clone(), abi.clone()))
+            })
             .collect();
+        let mut contracts = TargetedContracts { inner: contracts };
 
         self.target_interfaces(to, &mut contracts)?;
 
@@ -561,7 +561,7 @@ impl<'a> InvariantExecutor<'a> {
         // the specified interfaces for the same address. For example:
         // `[(addr1, ["IERC20", "IOwnable"])]` and `[(addr1, ["IERC20"]), (addr1, ("IOwnable"))]`
         // should be equivalent.
-        let mut combined: TargetedContracts = BTreeMap::new();
+        let mut combined = TargetedContracts::new();
 
         // Loop through each address and its associated artifact identifiers.
         // We're borrowing here to avoid taking full ownership.
@@ -577,18 +577,18 @@ impl<'a> InvariantExecutor<'a> {
                         .entry(*addr)
                         // If the entry exists, extends its ABI with the function list.
                         .and_modify(|entry| {
-                            let (_, contract_abi, _) = entry;
-
                             // Extend the ABI's function list with the new functions.
-                            contract_abi.functions.extend(contract.abi.functions.clone());
+                            entry.abi.functions.extend(contract.abi.functions.clone());
                         })
                         // Otherwise insert it into the map.
-                        .or_insert_with(|| (identifier.to_string(), contract.abi.clone(), vec![]));
+                        .or_insert_with(|| {
+                            TargetedContract::new(identifier.to_string(), contract.abi.clone())
+                        });
                 }
             }
         }
 
-        targeted_contracts.extend(combined);
+        targeted_contracts.extend(combined.inner);
 
         Ok(())
     }
@@ -623,26 +623,18 @@ impl<'a> InvariantExecutor<'a> {
         selectors: &[Selector],
         targeted_contracts: &mut TargetedContracts,
     ) -> eyre::Result<()> {
-        if let Some((name, abi, address_selectors)) = targeted_contracts.get_mut(&address) {
-            // The contract is already part of our filter, and all we do is specify that we're
-            // only looking at specific functions coming from `bytes4_array`.
-            for &selector in selectors {
-                address_selectors.push(get_function(name, selector, abi).cloned()?);
+        let contract = match targeted_contracts.entry(address) {
+            Entry::Occupied(entry) => entry.into_mut(),
+            Entry::Vacant(entry) => {
+                let (identifier, abi) = self.setup_contracts.get(&address).ok_or_else(|| {
+                    eyre::eyre!(
+                        "[targetSelectors] address does not have an associated contract: {address}"
+                    )
+                })?;
+                entry.insert(TargetedContract::new(identifier.clone(), abi.clone()))
             }
-        } else {
-            let (name, abi) = self.setup_contracts.get(&address).ok_or_else(|| {
-                eyre::eyre!(
-                    "[targetSelectors] address does not have an associated contract: {address}"
-                )
-            })?;
-
-            let functions = selectors
-                .iter()
-                .map(|&selector| get_function(name, selector, abi).cloned())
-                .collect::<Result<Vec<_>, _>>()?;
-
-            targeted_contracts.insert(address, (name.to_string(), abi.clone(), functions));
-        }
+        };
+        contract.add_selectors(selectors.iter().copied())?;
         Ok(())
     }
 

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

@@ -1,5 +1,5 @@
 use alloy_json_abi::{Function, JsonAbi};
-use alloy_primitives::{Address, Bytes};
+use alloy_primitives::{Address, Bytes, Selector};
 use itertools::Either;
 use parking_lot::Mutex;
 use std::{collections::BTreeMap, sync::Arc};
@@ -10,11 +10,10 @@ 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>)>;
+use foundry_evm_core::utils::{get_function, StateChangeset};
 
 /// Contracts identified as targets during a fuzz run.
+///
 /// During execution, any newly created contract is added as target and used through the rest of
 /// the fuzz run if the collection is updatable (no `targetContract` specified in `setUp`).
 #[derive(Clone, Debug)]
@@ -26,42 +25,11 @@ pub struct FuzzRunIdentifiedContracts {
 }
 
 impl FuzzRunIdentifiedContracts {
+    /// Creates a new `FuzzRunIdentifiedContracts` instance.
     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 with_fuzzed_artifacts(
-        &self,
-        tx: &BasicTxDetails,
-        f: impl FnOnce(Option<&JsonAbi>, Option<&Function>),
-    ) {
-        let targets = self.targets.lock();
-        let (abi, abi_f) = match targets.get(&tx.call_details.target) {
-            Some((_, abi, _)) => {
-                (Some(abi), abi.functions().find(|f| f.selector() == tx.call_details.calldata[..4]))
-            }
-            None => (None, None),
-        };
-        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(
@@ -72,34 +40,41 @@ impl FuzzRunIdentifiedContracts {
         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));
+        if !self.is_updatable {
+            return Ok(());
+        }
+
+        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);
+            let contract = TargetedContract {
+                identifier: artifact.name.clone(),
+                abi: contract.abi.clone(),
+                targeted_functions: functions,
+            };
+            targets.insert(*address, contract);
         }
         Ok(())
     }
@@ -115,21 +90,102 @@ impl FuzzRunIdentifiedContracts {
     }
 }
 
-/// 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<'a>(
-    abi: &'a JsonAbi,
-    targeted_functions: &'a [Function],
-) -> impl Iterator<Item = &'a Function> {
-    if !targeted_functions.is_empty() {
-        Either::Left(targeted_functions.iter())
-    } else {
-        Either::Right(abi.functions().filter(|&func| {
-            !matches!(
-                func.state_mutability,
-                alloy_json_abi::StateMutability::Pure | alloy_json_abi::StateMutability::View
-            )
-        }))
+/// A collection of contracts identified as targets for invariant testing.
+#[derive(Clone, Debug, Default)]
+pub struct TargetedContracts {
+    /// The inner map of targeted contracts.
+    pub inner: BTreeMap<Address, TargetedContract>,
+}
+
+impl TargetedContracts {
+    /// Returns a new `TargetedContracts` instance.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// 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.inner.get(&tx.call_details.target) {
+            Some(c) => (
+                Some(&c.abi),
+                c.abi.functions().find(|f| f.selector() == tx.call_details.calldata[..4]),
+            ),
+            None => (None, None),
+        }
+    }
+
+    /// 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) -> impl Iterator<Item = (&Address, &Function)> {
+        self.inner
+            .iter()
+            .filter(|(_, c)| !c.abi.functions.is_empty())
+            .flat_map(|(contract, c)| c.abi_fuzzed_functions().map(move |f| (contract, f)))
+    }
+}
+
+impl std::ops::Deref for TargetedContracts {
+    type Target = BTreeMap<Address, TargetedContract>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl std::ops::DerefMut for TargetedContracts {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}
+
+/// A contract identified as targets for invariant testing.
+#[derive(Clone, Debug)]
+pub struct TargetedContract {
+    /// The contract identifier. This is only used in error messages.
+    pub identifier: String,
+    /// The contract's ABI.
+    pub abi: JsonAbi,
+    /// The targeted functions of the contract.
+    pub targeted_functions: Vec<Function>,
+}
+
+impl TargetedContract {
+    /// Returns a new `TargetedContract` instance.
+    pub fn new(identifier: String, abi: JsonAbi) -> Self {
+        Self { identifier, abi, targeted_functions: Vec::new() }
+    }
+
+    /// Helper to retrieve functions to fuzz for specified abi.
+    /// Returns specified targeted functions if any, else mutable abi functions.
+    pub fn abi_fuzzed_functions(&self) -> impl Iterator<Item = &Function> {
+        if !self.targeted_functions.is_empty() {
+            Either::Left(self.targeted_functions.iter())
+        } else {
+            Either::Right(self.abi.functions().filter(|&func| {
+                !matches!(
+                    func.state_mutability,
+                    alloy_json_abi::StateMutability::Pure | alloy_json_abi::StateMutability::View
+                )
+            }))
+        }
+    }
+
+    /// Returns the function for the given selector.
+    pub fn get_function(&self, selector: Selector) -> eyre::Result<&Function> {
+        get_function(&self.identifier, selector, &self.abi)
+    }
+
+    /// Adds the specified selectors to the targeted functions.
+    pub fn add_selectors(
+        &mut self,
+        selectors: impl IntoIterator<Item = Selector>,
+    ) -> eyre::Result<()> {
+        for selector in selectors {
+            self.targeted_functions.push(self.get_function(selector)?.clone());
+        }
+        Ok(())
     }
 }