feat(forge): filter through artifacts on invariant testing (#2635)

* filter through contract names or identifiers

* use identifier method from ArtifactId

* filter newly generated contracts

* give a more helpful message

* improve docs

* fmt

* rename Abi to Artifact

* add in-run contract filtering test

* fix fmt

* add ContractsByArtifact and ContractsByArtifactExt

* add ContractsByAddress

* add ArtifactFilters

* replace ContractsByArtifactExt

* move some utils to common

* add missing dev dependency

* move override test from invariant tests

* fix ArtifactFilters.get_targeted_functions

Author: joshieDo

Cargo.lock

@@ -76,7 +76,7 @@ impl ScriptArgs {
         transactions: VecDeque<TypedTransaction>,
         script_config: &mut ScriptConfig,
         decoder: &mut CallTraceDecoder,
-        contracts: &BTreeMap<ArtifactId, (Abi, Vec<u8>)>,
+        contracts: &ContractsByArtifact,
     ) -> eyre::Result<VecDeque<TransactionWithMetadata>> {
         let mut runner = self
             .prepare_runner(script_config, script_config.evm_opts.sender, SimulationStage::OnChain)
@@ -93,8 +93,7 @@ impl ScriptArgs {
             .iter()
             .filter_map(|(addr, contract_id)| {
                 let contract_name = utils::get_contract_name(contract_id);
-                if let Some((_, (abi, _))) =
-                    contracts.iter().find(|(artifact, _)| artifact.name == contract_name)
+                if let Ok(Some((_, (abi, _)))) = contracts.find_by_name_or_identifier(contract_name)
                 {
                     return Some((*addr, (contract_name.to_string(), abi)))
                 }

cli/src/cmd/forge/script/executor.rs

@@ -30,7 +30,7 @@ use forge::{
         CallTraceArena, CallTraceDecoder, CallTraceDecoderBuilder, TraceKind,
     },
 };
-use foundry_common::{evm::EvmArgs, CONTRACT_MAX_SIZE};
+use foundry_common::{evm::EvmArgs, ContractsByArtifact, CONTRACT_MAX_SIZE};
 use foundry_config::Config;
 use foundry_utils::{encode_args, format_token, IntoFunction};
 use serde::{Deserialize, Serialize};
@@ -161,7 +161,7 @@ impl ScriptArgs {
         &self,
         script_config: &ScriptConfig,
         result: &mut ScriptResult,
-        known_contracts: &BTreeMap<ArtifactId, (Abi, Vec<u8>)>,
+        known_contracts: &ContractsByArtifact,
     ) -> eyre::Result<CallTraceDecoder> {
         let etherscan_identifier = EtherscanIdentifier::new(
             script_config.evm_opts.get_remote_chain_id(),
@@ -520,7 +520,7 @@ pub struct ScriptConfig {
 /// Data struct to help `ScriptSequence` verify contracts on `etherscan`.
 pub struct VerifyBundle {
     pub num_of_optimizations: Option<usize>,
-    pub known_contracts: BTreeMap<ArtifactId, (Abi, Vec<u8>)>,
+    pub known_contracts: ContractsByArtifact,
     pub etherscan_key: Option<String>,
     pub project_paths: ProjectPathsArgs,
     pub retry: RetryArgs,
@@ -530,7 +530,7 @@ impl VerifyBundle {
     pub fn new(
         project: &Project,
         config: &Config,
-        known_contracts: BTreeMap<ArtifactId, (Abi, Vec<u8>)>,
+        known_contracts: ContractsByArtifact,
         retry: RetryArgs,
     ) -> Self {
         let num_of_optimizations =

cli/src/cmd/forge/script/mod.rs

@@ -164,7 +164,7 @@ impl ScriptSequence {
                 if let (Some(contract_address), Some(data)) =
                     (receipt.contract_address, tx.typed_tx().data())
                 {
-                    for (artifact, (_contract, bytecode)) in &verify.known_contracts {
+                    for (artifact, (_contract, bytecode)) in verify.known_contracts.iter() {
                         // If it's a CREATE2, the tx.data comes with a 32-byte salt in the beginning
                         // of the transaction
                         if data.0.split_at(create2_offset).1.starts_with(bytecode) {

cli/src/cmd/forge/script/sequence.rs

@@ -14,7 +14,7 @@ use ethers::{
     },
 };
 use forge::executor::opts::EvmOpts;
-use foundry_common::TestFunctionExt;
+use foundry_common::{ContractsByArtifact, TestFunctionExt};
 use foundry_config::{figment::Figment, Chain as ConfigChain, Config};
 use foundry_utils::Retry;
 use std::{collections::BTreeMap, path::PathBuf};
@@ -169,22 +169,24 @@ pub fn needs_setup(abi: &Abi) -> bool {
 pub fn unwrap_contracts(
     contracts: &BTreeMap<ArtifactId, ContractBytecodeSome>,
     deployed_code: bool,
-) -> BTreeMap<ArtifactId, (Abi, Vec<u8>)> {
-    contracts
-        .iter()
-        .filter_map(|(id, c)| {
-            let bytecode = if deployed_code {
-                c.deployed_bytecode.clone().into_bytes()
-            } else {
-                c.bytecode.clone().object.into_bytes()
-            };
+) -> ContractsByArtifact {
+    ContractsByArtifact(
+        contracts
+            .iter()
+            .filter_map(|(id, c)| {
+                let bytecode = if deployed_code {
+                    c.deployed_bytecode.clone().into_bytes()
+                } else {
+                    c.bytecode.clone().object.into_bytes()
+                };
 
-            if let Some(bytecode) = bytecode {
-                return Some((id.clone(), (c.abi.clone(), bytecode.to_vec())))
-            }
-            None
-        })
-        .collect()
+                if let Some(bytecode) = bytecode {
+                    return Some((id.clone(), (c.abi.clone(), bytecode.to_vec())))
+                }
+                None
+            })
+            .collect(),
+    )
 }
 
 #[macro_export]

cli/src/cmd/utils.rs

@@ -15,6 +15,7 @@ foundry-config = { path = "../config" }
 
 # eth
 ethers-core = { git = "https://github.com/gakonst/ethers-rs", default-features = false }
+ethers-solc = { git = "https://github.com/gakonst/ethers-rs", default-features = false }
 ethers-providers = { git = "https://github.com/gakonst/ethers-rs", default-features = false }
 
 # io

common/Cargo.toml

@@ -0,0 +1,110 @@
+//! commonly used contract types and functions
+
+use ethers_core::{
+    abi::{Abi, Event, Function},
+    types::{Address, H256},
+};
+use ethers_solc::ArtifactId;
+use std::{
+    collections::BTreeMap,
+    ops::{Deref, DerefMut},
+};
+
+type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a (Abi, Vec<u8>));
+
+/// Wrapper type that maps an artifact to a contract ABI and bytecode.
+#[derive(Default)]
+pub struct ContractsByArtifact(pub BTreeMap<ArtifactId, (Abi, Vec<u8>)>);
+
+impl ContractsByArtifact {
+    /// Finds a contract which has a similar bytecode as `code`.
+    pub fn find_by_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef> {
+        self.iter().find(|(_, (_, known_code))| diff_score(known_code, code) < 0.1)
+    }
+    /// Finds a contract which has the same contract name or identifier as `id`. If more than one is
+    /// found, return error.
+    pub fn find_by_name_or_identifier(
+        &self,
+        id: &str,
+    ) -> eyre::Result<Option<ArtifactWithContractRef>> {
+        let contracts = self
+            .iter()
+            .filter(|(artifact, _)| artifact.name == id || artifact.identifier() == id)
+            .collect::<Vec<_>>();
+
+        if contracts.len() > 1 {
+            eyre::bail!("{id} has more than one implementation.");
+        }
+
+        Ok(contracts.first().cloned())
+    }
+
+    /// Flattens a group of contracts into maps of all events and functions
+    pub fn flatten(&self) -> (BTreeMap<[u8; 4], Function>, BTreeMap<H256, Event>, Abi) {
+        let flattened_funcs: BTreeMap<[u8; 4], Function> = self
+            .iter()
+            .flat_map(|(_name, (abi, _code))| {
+                abi.functions()
+                    .map(|func| (func.short_signature(), func.clone()))
+                    .collect::<BTreeMap<[u8; 4], Function>>()
+            })
+            .collect();
+
+        let flattened_events: BTreeMap<H256, Event> = self
+            .iter()
+            .flat_map(|(_name, (abi, _code))| {
+                abi.events()
+                    .map(|event| (event.signature(), event.clone()))
+                    .collect::<BTreeMap<H256, Event>>()
+            })
+            .collect();
+
+        // We need this for better revert decoding, and want it in abi form
+        let mut errors_abi = Abi::default();
+        self.iter().for_each(|(_name, (abi, _code))| {
+            abi.errors().for_each(|error| {
+                let entry =
+                    errors_abi.errors.entry(error.name.clone()).or_insert_with(Default::default);
+                entry.push(error.clone());
+            });
+        });
+        (flattened_funcs, flattened_events, errors_abi)
+    }
+}
+
+impl Deref for ContractsByArtifact {
+    type Target = BTreeMap<ArtifactId, (Abi, Vec<u8>)>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for ContractsByArtifact {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+/// Wrapper type that maps an address to a contract identifier and contract ABI.
+pub type ContractsByAddress = BTreeMap<Address, (String, Abi)>;
+
+/// Very simple fuzzy matching of contract bytecode.
+///
+/// Will fail for small contracts that are essentially all immutable variables.
+pub fn diff_score(a: &[u8], b: &[u8]) -> f64 {
+    let cutoff_len = usize::min(a.len(), b.len());
+    if cutoff_len == 0 {
+        return 1.0
+    }
+
+    let a = &a[..cutoff_len];
+    let b = &b[..cutoff_len];
+    let mut diff_chars = 0;
+    for i in 0..cutoff_len {
+        if a[i] != b[i] {
+            diff_chars += 1;
+        }
+    }
+    diff_chars as f64 / cutoff_len as f64
+}

common/src/contracts.rs

@@ -4,6 +4,7 @@
 
 pub mod calc;
 pub mod constants;
+pub mod contracts;
 pub mod errors;
 pub mod evm;
 pub mod fmt;
@@ -12,4 +13,5 @@ pub mod provider;
 pub use provider::*;
 pub mod traits;
 pub use constants::*;
+pub use contracts::*;
 pub use traits::*;