refactor: optimize trace identifiers (foundry-rs#5680)

Author: DaniPopes

crates/chisel/src/dispatcher.rs

@@ -958,13 +958,13 @@ impl ChiselDispatcher {
             session_config.evm_opts.get_remote_chain_id(),
         )?;
 
-        let mut decoder =
-            CallTraceDecoderBuilder::new().with_labels(result.labeled_addresses.clone()).build();
-
-        decoder.add_signature_identifier(SignaturesIdentifier::new(
-            Config::foundry_cache_dir(),
-            session_config.foundry_config.offline,
-        )?);
+        let mut decoder = CallTraceDecoderBuilder::new()
+            .with_labels(result.labeled_addresses.iter().map(|(a, s)| (*a, s.clone())))
+            .with_signature_identifier(SignaturesIdentifier::new(
+                Config::foundry_cache_dir(),
+                session_config.foundry_config.offline,
+            )?)
+            .build();
 
         for (_, trace) in &mut result.traces {
             // decoder.identify(trace, &mut local_identifier);

crates/cli/src/utils/cmd.rs

@@ -378,12 +378,13 @@ pub async fn handle_traces(
         None
     });
 
-    let mut decoder = CallTraceDecoderBuilder::new().with_labels(labeled_addresses).build();
-
-    decoder.add_signature_identifier(SignaturesIdentifier::new(
-        Config::foundry_cache_dir(),
-        config.offline,
-    )?);
+    let mut decoder = CallTraceDecoderBuilder::new()
+        .with_labels(labeled_addresses)
+        .with_signature_identifier(SignaturesIdentifier::new(
+            Config::foundry_cache_dir(),
+            config.offline,
+        )?)
+        .build();
 
     for (_, trace) in &mut result.traces {
         decoder.identify(trace, &mut etherscan_identifier);

crates/evm/src/trace/decoder.rs

@@ -1,5 +1,5 @@
 use super::{
-    identifier::{SingleSignaturesIdentifier, TraceIdentifier},
+    identifier::{AddressIdentity, SingleSignaturesIdentifier, TraceIdentifier},
     CallTraceArena, RawOrDecodedCall, RawOrDecodedLog, RawOrDecodedReturnData,
 };
 use crate::{
@@ -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,21 @@ impl CallTraceDecoderBuilder {
     }
 
     /// Sets the verbosity level of the decoder.
+    #[inline]
     pub fn with_verbosity(mut self, level: u8) -> Self {
         self.decoder.verbosity = level;
         self
     }
 
+    /// Sets the signature identifier for events and functions.
+    #[inline]
+    pub fn with_signature_identifier(mut self, identifier: SingleSignaturesIdentifier) -> Self {
+        self.decoder.signature_identifier = Some(identifier);
+        self
+    }
+
     /// Build the decoder.
+    #[inline]
     pub fn build(self) -> CallTraceDecoder {
         self.decoder
     }
@@ -168,23 +182,26 @@ impl CallTraceDecoder {
         }
     }
 
-    pub fn add_signature_identifier(&mut self, identifier: SingleSignaturesIdentifier) {
-        self.signature_identifier = Some(identifier);
-    }
-
     /// Identify unknown addresses in the specified call trace using the specified identifier.
     ///
     /// Unknown contracts are contracts that either lack a label or an ABI.
+    #[inline]
     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| {
+        self.collect_identities(identifier.identify_addresses(self.addresses(trace)));
+    }
+
+    #[inline(always)]
+    fn addresses<'a>(
+        &'a self,
+        trace: &'a CallTraceArena,
+    ) -> impl Iterator<Item = (&'a Address, Option<&'a [u8]>)> + 'a {
+        trace.addresses().into_iter().filter(|&(address, _)| {
+            !self.labels.contains_key(address) || !self.contracts.contains_key(address)
+        })
+    }
+
+    fn collect_identities(&mut self, identities: Vec<AddressIdentity>) {
+        for identity in identities {
             let address = identity.address;
 
             if let Some(contract) = &identity.contract {
@@ -197,30 +214,31 @@ impl CallTraceDecoder {
 
             if let Some(abi) = &identity.abi {
                 // Store known functions for the address
-                abi.functions()
-                    .map(|func| (func.short_signature(), func.clone()))
-                    .for_each(|(sig, func)| self.functions.entry(sig).or_default().push(func));
+                for function in abi.functions() {
+                    self.functions
+                        .entry(function.short_signature())
+                        .or_default()
+                        .push(function.clone())
+                }
 
                 // Flatten events from all ABIs
-                abi.events()
-                    .map(|event| ((event.signature(), indexed_inputs(event)), event.clone()))
-                    .for_each(|(sig, event)| {
-                        self.events.entry(sig).or_default().push(event);
-                    });
+                for event in abi.events() {
+                    let sig = (event.signature(), indexed_inputs(event));
+                    self.events.entry(sig).or_default().push(event.clone());
+                }
 
                 // Flatten errors from all ABIs
-                abi.errors().for_each(|error| {
-                    let entry = self.errors.errors.entry(error.name.clone()).or_default();
-                    entry.push(error.clone());
-                });
+                for error in abi.errors() {
+                    self.errors.errors.entry(error.name.clone()).or_default().push(error.clone());
+                }
 
                 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,45 @@
 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);
+                    .filter_map(|(id, (abi, known_code))| {
+                        let score = diff_score(known_code, code);
                         if score < 0.1 {
-                            Some((OrderedFloat(score), entry))
+                            Some((OrderedFloat(score), id, abi))
                         } else {
                             None
                         }
                     })
-                    .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

@@ -561,13 +561,39 @@ impl fmt::Display for CallTrace {
 }
 
 /// Specifies the kind of trace.
-#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub enum TraceKind {
     Deployment,
     Setup,
     Execution,
 }
 
+impl TraceKind {
+    /// Returns `true` if the trace kind is [`Deployment`].
+    ///
+    /// [`Deployment`]: TraceKind::Deployment
+    #[must_use]
+    pub fn is_deployment(self) -> bool {
+        matches!(self, Self::Deployment)
+    }
+
+    /// Returns `true` if the trace kind is [`Setup`].
+    ///
+    /// [`Setup`]: TraceKind::Setup
+    #[must_use]
+    pub fn is_setup(self) -> bool {
+        matches!(self, Self::Setup)
+    }
+
+    /// Returns `true` if the trace kind is [`Execution`].
+    ///
+    /// [`Execution`]: TraceKind::Execution
+    #[must_use]
+    pub fn is_execution(self) -> bool {
+        matches!(self, Self::Execution)
+    }
+}
+
 /// Chooses the color of the trace depending on the destination address and status of the call.
 fn trace_color(trace: &CallTrace) -> Color {
     if trace.address == CHEATCODE_ADDRESS {
@@ -584,26 +610,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/evm/src/trace/utils.rs

@@ -79,7 +79,7 @@ pub(crate) fn decode_cheatcode_inputs(
         "serializeBytes32" |
         "serializeString" |
         "serializeBytes" => {
-            if verbosity == 5 {
+            if verbosity >= 5 {
                 None
             } else {
                 let mut decoded = func.decode_input(&data[SELECTOR_LEN..]).ok()?;
@@ -111,10 +111,10 @@ pub(crate) fn decode_cheatcode_outputs(
         // redacts derived private key
         return Some("<pk>".to_string())
     }
-    if func.name == "parseJson" && verbosity != 5 {
+    if func.name == "parseJson" && verbosity < 5 {
         return Some("<encoded JSON value>".to_string())
     }
-    if func.name == "readFile" && verbosity != 5 {
+    if func.name == "readFile" && verbosity < 5 {
         return Some("<file>".to_string())
     }
     None