refactor: optimize trace identifiers

Author: DaniPopes

crates/evm/src/trace/decoder.rs

@@ -20,22 +20,27 @@ use std::collections::{BTreeMap, HashMap};
 
 /// Build a new [CallTraceDecoder].
 #[derive(Default)]
+#[must_use = "builders do nothing unless you call `build` on them"]
 pub struct CallTraceDecoderBuilder {
     decoder: CallTraceDecoder,
 }
 
 impl CallTraceDecoderBuilder {
+    /// Create a new builder.
+    #[inline]
     pub fn new() -> Self {
         Self { decoder: CallTraceDecoder::new().clone() }
     }
 
     /// Add known labels to the decoder.
+    #[inline]
     pub fn with_labels(mut self, labels: impl IntoIterator<Item = (Address, String)>) -> Self {
         self.decoder.labels.extend(labels);
         self
     }
 
     /// Add known events to the decoder.
+    #[inline]
     pub fn with_events(mut self, events: impl IntoIterator<Item = Event>) -> Self {
         for event in events {
             self.decoder
@@ -48,12 +53,14 @@ impl CallTraceDecoderBuilder {
     }
 
     /// Sets the verbosity level of the decoder.
+    #[inline]
     pub fn with_verbosity(mut self, level: u8) -> Self {
         self.decoder.verbosity = level;
         self
     }
 
     /// Build the decoder.
+    #[inline]
     pub fn build(self) -> CallTraceDecoder {
         self.decoder
     }
@@ -176,15 +183,10 @@ impl CallTraceDecoder {
     ///
     /// Unknown contracts are contracts that either lack a label or an ABI.
     pub fn identify(&mut self, trace: &CallTraceArena, identifier: &mut impl TraceIdentifier) {
-        let unidentified_addresses = trace
-            .addresses()
-            .into_iter()
-            .filter(|(address, _)| {
-                !self.labels.contains_key(address) || !self.contracts.contains_key(address)
-            })
-            .collect();
-
-        identifier.identify_addresses(unidentified_addresses).iter().for_each(|identity| {
+        let unidentified_addresses = trace.addresses().into_iter().filter(|&(address, _)| {
+            !self.labels.contains_key(address) || !self.contracts.contains_key(address)
+        });
+        for identity in identifier.identify_addresses(unidentified_addresses) {
             let address = identity.address;
 
             if let Some(contract) = &identity.contract {
@@ -216,11 +218,11 @@ impl CallTraceDecoder {
 
                 self.receive_contracts.entry(address).or_insert(abi.receive);
             }
-        });
+        }
     }
 
     pub async fn decode(&self, traces: &mut CallTraceArena) {
-        for node in traces.arena.iter_mut() {
+        for node in &mut traces.arena {
             // Set contract name
             if let Some(contract) = self.contracts.get(&node.trace.address).cloned() {
                 node.trace.contract = Some(contract);

crates/evm/src/trace/identifier/etherscan.rs

@@ -100,11 +100,11 @@ impl EtherscanIdentifier {
 }
 
 impl TraceIdentifier for EtherscanIdentifier {
-    fn identify_addresses(
-        &mut self,
-        addresses: Vec<(&Address, Option<&[u8]>)>,
-    ) -> Vec<AddressIdentity> {
-        trace!(target: "etherscanidentifier", "identify {} addresses", addresses.len());
+    fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity>
+    where
+        A: Iterator<Item = (&'a Address, Option<&'a [u8]>)>,
+    {
+        trace!(target: "etherscanidentifier", "identify {:?} addresses", addresses.size_hint().1);
 
         let Some(client) = self.client.clone() else {
             // no client was configured

crates/evm/src/trace/identifier/local.rs

@@ -1,56 +1,40 @@
 use super::{AddressIdentity, TraceIdentifier};
-use ethers::{
-    abi::{Abi, Address, Event},
-    prelude::ArtifactId,
-};
+use ethers::abi::{Address, Event};
 use foundry_common::contracts::{diff_score, ContractsByArtifact};
-use itertools::Itertools;
 use ordered_float::OrderedFloat;
-use std::{borrow::Cow, collections::BTreeMap};
+use std::borrow::Cow;
 
 /// A trace identifier that tries to identify addresses using local contracts.
-pub struct LocalTraceIdentifier {
-    local_contracts: BTreeMap<Vec<u8>, (ArtifactId, Abi)>,
+pub struct LocalTraceIdentifier<'a> {
+    known_contracts: &'a ContractsByArtifact,
 }
 
-impl LocalTraceIdentifier {
-    pub fn new(known_contracts: &ContractsByArtifact) -> Self {
-        Self {
-            local_contracts: known_contracts
-                .iter()
-                .map(|(id, (abi, runtime_code))| (runtime_code.clone(), (id.clone(), abi.clone())))
-                .collect(),
-        }
+impl<'a> LocalTraceIdentifier<'a> {
+    pub fn new(known_contracts: &'a ContractsByArtifact) -> Self {
+        Self { known_contracts }
     }
 
     /// Get all the events of the local contracts.
     pub fn events(&self) -> impl Iterator<Item = &Event> {
-        self.local_contracts.iter().flat_map(|(_, (_, abi))| abi.events())
+        self.known_contracts.iter().flat_map(|(_, (abi, _))| abi.events())
     }
 }
 
-impl TraceIdentifier for LocalTraceIdentifier {
-    fn identify_addresses(
-        &mut self,
-        addresses: Vec<(&Address, Option<&[u8]>)>,
-    ) -> Vec<AddressIdentity> {
+impl TraceIdentifier for LocalTraceIdentifier<'_> {
+    fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity>
+    where
+        A: Iterator<Item = (&'a Address, Option<&'a [u8]>)>,
+    {
         addresses
-            .into_iter()
             .filter_map(|(address, code)| {
                 let code = code?;
-                let (_, (_, (id, abi))) = self
-                    .local_contracts
+                let (_, (id, abi)) = self
+                    .known_contracts
                     .iter()
-                    .filter_map(|entry| {
-                        let score = diff_score(entry.0, code);
-                        if score < 0.1 {
-                            Some((OrderedFloat(score), entry))
-                        } else {
-                            None
-                        }
+                    .map(|(id, (abi, known_code))| {
+                        (OrderedFloat(diff_score(known_code, code)), (id, abi))
                     })
-                    .sorted_by_key(|(score, _)| *score)
-                    .next()?;
+                    .min_by_key(|(score, _)| *score)?;
 
                 Some(AddressIdentity {
                     address: *address,

crates/evm/src/trace/identifier/mod.rs

@@ -33,9 +33,7 @@ pub struct AddressIdentity<'a> {
 pub trait TraceIdentifier {
     // TODO: Update docs
     /// Attempts to identify an address in one or more call traces.
-    #[allow(clippy::type_complexity)]
-    fn identify_addresses(
-        &mut self,
-        addresses: Vec<(&Address, Option<&[u8]>)>,
-    ) -> Vec<AddressIdentity>;
+    fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity>
+    where
+        A: Iterator<Item = (&'a Address, Option<&'a [u8]>)>;
 }

crates/evm/src/trace/mod.rs

@@ -584,26 +584,23 @@ pub fn load_contracts(
     traces: Traces,
     known_contracts: Option<&ContractsByArtifact>,
 ) -> ContractsByAddress {
-    if let Some(contracts) = known_contracts {
-        let mut local_identifier = LocalTraceIdentifier::new(contracts);
-        let mut decoder = CallTraceDecoderBuilder::new().build();
-        for (_, trace) in &traces {
-            decoder.identify(trace, &mut local_identifier);
-        }
-
-        decoder
-            .contracts
-            .iter()
-            .filter_map(|(addr, name)| {
-                if let Ok(Some((_, (abi, _)))) = contracts.find_by_name_or_identifier(name) {
-                    return Some((*addr, (name.clone(), abi.clone())))
-                }
-                None
-            })
-            .collect()
-    } else {
-        BTreeMap::new()
+    let Some(contracts) = known_contracts else { return BTreeMap::new() };
+    let mut local_identifier = LocalTraceIdentifier::new(contracts);
+    let mut decoder = CallTraceDecoderBuilder::new().build();
+    for (_, trace) in &traces {
+        decoder.identify(trace, &mut local_identifier);
     }
+
+    decoder
+        .contracts
+        .iter()
+        .filter_map(|(addr, name)| {
+            if let Ok(Some((_, (abi, _)))) = contracts.find_by_name_or_identifier(name) {
+                return Some((*addr, (name.clone(), abi.clone())))
+            }
+            None
+        })
+        .collect()
 }
 
 /// creates the memory data in 32byte chunks

crates/forge/bin/cmd/script/mod.rs

@@ -237,7 +237,7 @@ impl ScriptArgs {
 
         let mut local_identifier = LocalTraceIdentifier::new(known_contracts);
         let mut decoder = CallTraceDecoderBuilder::new()
-            .with_labels(result.labeled_addresses.clone())
+            .with_labels(result.labeled_addresses.iter().map(|(a, s)| (*a, s.clone())))
             .with_verbosity(verbosity)
             .build();
 

crates/forge/bin/cmd/test/mod.rs

@@ -540,7 +540,8 @@ async fn test(
         Ok(TestOutcome::new(results, allow_failure))
     } else {
         // Set up identifiers
-        let mut local_identifier = LocalTraceIdentifier::new(&runner.known_contracts);
+        let known_contracts = runner.known_contracts.clone();
+        let mut local_identifier = LocalTraceIdentifier::new(&known_contracts);
         let remote_chain_id = runner.evm_opts.get_remote_chain_id();
         // Do not re-query etherscan for contracts that you've already queried today.
         let mut etherscan_identifier = EtherscanIdentifier::new(&config, remote_chain_id)?;