From d141ffa1e1f3185f997ba23df987ca9504d8d0a8 Mon Sep 17 00:00:00 2001 From: DaniPopes <57450786+DaniPopes@users.noreply.github.com> Date: Tue, 18 Jun 2024 09:38:53 +0200 Subject: [PATCH] perf: new-type TargetedContracts (#8180) --- .../evm/evm/src/executors/invariant/error.rs | 2 +- crates/evm/evm/src/executors/invariant/mod.rs | 60 +++-- crates/evm/fuzz/src/invariant/mod.rs | 212 +++++++++++------- crates/evm/fuzz/src/strategies/invariants.rs | 27 +-- crates/evm/fuzz/src/strategies/state.rs | 6 +- crates/forge/tests/it/invariant.rs | 2 +- .../common/InvariantShrinkWithAssert.t.sol | 2 +- 7 files changed, 179 insertions(+), 132 deletions(-) diff --git a/crates/evm/evm/src/executors/invariant/error.rs b/crates/evm/evm/src/executors/invariant/error.rs index 22d751ed5b02..7c47ac0a47c6 100644 --- a/crates/evm/evm/src/executors/invariant/error.rs +++ b/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)); diff --git a/crates/evm/evm/src/executors/invariant/mod.rs b/crates/evm/evm/src/executors/invariant/mod.rs index 1eee0cafc6fd..0918b3ebac17 100644 --- a/crates/evm/evm/src/executors/invariant/mod.rs +++ b/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::, _>>()?; - - targeted_contracts.insert(address, (name.to_string(), abi.clone(), functions)); - } + }; + contract.add_selectors(selectors.iter().copied())?; Ok(()) } diff --git a/crates/evm/fuzz/src/invariant/mod.rs b/crates/evm/fuzz/src/invariant/mod.rs index fdc0b9d56c31..9138443be221 100644 --- a/crates/evm/fuzz/src/invariant/mod.rs +++ b/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)>; +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
, ) -> 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 { - 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, +} + +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 { + 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; + + 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, +} + +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 { + 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, + ) -> eyre::Result<()> { + for selector in selectors { + self.targeted_functions.push(self.get_function(selector)?.clone()); + } + Ok(()) } } diff --git a/crates/evm/fuzz/src/strategies/invariants.rs b/crates/evm/fuzz/src/strategies/invariants.rs index 55af574c5989..c7d04dd1ab02 100644 --- a/crates/evm/fuzz/src/strategies/invariants.rs +++ b/crates/evm/fuzz/src/strategies/invariants.rs @@ -1,9 +1,6 @@ use super::{fuzz_calldata, fuzz_param_from_state}; use crate::{ - invariant::{ - abi_fuzzed_functions, BasicTxDetails, CallDetails, FuzzRunIdentifiedContracts, - SenderFilters, - }, + invariant::{BasicTxDetails, CallDetails, FuzzRunIdentifiedContracts, SenderFilters}, strategies::{fuzz_calldata_from_state, fuzz_param, EvmFuzzState}, FuzzFixtures, }; @@ -11,6 +8,7 @@ use alloy_json_abi::Function; use alloy_primitives::Address; use parking_lot::RwLock; use proptest::prelude::*; +use rand::seq::IteratorRandom; use std::{rc::Rc, sync::Arc}; /// Given a target address, we generate random calldata. @@ -32,15 +30,13 @@ pub fn override_call_strat( let func = { let contracts = contracts.targets.lock(); - let (_, abi, functions) = contracts.get(&target_address).unwrap_or_else(|| { + let contract = contracts.get(&target_address).unwrap_or_else(|| { // Choose a random contract if target selected by lazy strategy is not in fuzz run // identified contracts. This can happen when contract is created in `setUp` call // but is not included in targetContracts. - let rand_index = rand::thread_rng().gen_range(0..contracts.len()); - let (_, contract_specs) = contracts.iter().nth(rand_index).unwrap(); - contract_specs + contracts.values().choose(&mut rand::thread_rng()).unwrap() }); - let fuzzed_functions: Vec<_> = abi_fuzzed_functions(abi, functions).cloned().collect(); + let fuzzed_functions: Vec<_> = contract.abi_fuzzed_functions().cloned().collect(); any::().prop_map(move |index| index.get(&fuzzed_functions).clone()) }; @@ -68,16 +64,17 @@ pub fn invariant_strat( fuzz_fixtures: FuzzFixtures, ) -> impl Strategy { let senders = Rc::new(senders); - any::() - .prop_flat_map(move |index| { - let (target_address, target_function) = - index.get(&contracts.fuzzed_functions()).clone(); + any::() + .prop_flat_map(move |selector| { + let contracts = contracts.targets.lock(); + let functions = contracts.fuzzed_functions(); + let (target_address, target_function) = selector.select(functions); 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, + *target_address, + target_function.clone(), ); (sender, call_details) }) diff --git a/crates/evm/fuzz/src/strategies/state.rs b/crates/evm/fuzz/src/strategies/state.rs index cf29c94209a5..f4f1dde92482 100644 --- a/crates/evm/fuzz/src/strategies/state.rs +++ b/crates/evm/fuzz/src/strategies/state.rs @@ -64,10 +64,12 @@ impl EvmFuzzState { run_depth: u32, ) { let mut dict = self.inner.write(); - fuzzed_contracts.with_fuzzed_artifacts(tx, |target_abi, target_function| { + { + let targets = fuzzed_contracts.targets.lock(); + let (target_abi, target_function) = targets.fuzzed_artifacts(tx); dict.insert_logs_values(target_abi, logs, run_depth); dict.insert_result_values(target_function, result, run_depth); - }); + } dict.insert_new_state_values(state_changeset); } diff --git a/crates/forge/tests/it/invariant.rs b/crates/forge/tests/it/invariant.rs index 19500ab63df5..b6e564d49f76 100644 --- a/crates/forge/tests/it/invariant.rs +++ b/crates/forge/tests/it/invariant.rs @@ -773,7 +773,7 @@ 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)); + opts.fuzz.seed = Some(U256::from(100u32)); let filter = Filter::new(".*", ".*", ".*fuzz/invariant/common/InvariantSelectorsWeight.t.sol"); let mut runner = TEST_DATA_DEFAULT.runner(); diff --git a/testdata/default/fuzz/invariant/common/InvariantShrinkWithAssert.t.sol b/testdata/default/fuzz/invariant/common/InvariantShrinkWithAssert.t.sol index c189e2507629..d5dcfe6740a3 100644 --- a/testdata/default/fuzz/invariant/common/InvariantShrinkWithAssert.t.sol +++ b/testdata/default/fuzz/invariant/common/InvariantShrinkWithAssert.t.sol @@ -86,7 +86,7 @@ contract InvariantShrinkWithAssert is DSTest { } function invariant_with_assert() public { - assertTrue(counter.number() != 3); + assertTrue(counter.number() != 3, "wrong counter"); } }