chore(traces): add a trace identifier stack/builder (#7338)

crates/chisel/src/dispatcher.rs

@@ -17,7 +17,7 @@ use foundry_config::{Config, RpcEndpoint};
 use foundry_evm::{
     decode::decode_console_logs,
     traces::{
-        identifier::{EtherscanIdentifier, SignaturesIdentifier},
+        identifier::{SignaturesIdentifier, TraceIdentifiers},
         render_trace_arena, CallTraceDecoder, CallTraceDecoderBuilder, TraceKind,
     },
 };
@@ -893,11 +893,6 @@ impl ChiselDispatcher {
         result: &mut ChiselResult,
         // known_contracts: &ContractsByArtifact,
     ) -> eyre::Result<CallTraceDecoder> {
-        let mut etherscan_identifier = EtherscanIdentifier::new(
-            &session_config.foundry_config,
-            session_config.evm_opts.get_remote_chain_id(),
-        )?;
-
         let mut decoder = CallTraceDecoderBuilder::new()
             .with_labels(result.labeled_addresses.clone())
             .with_signature_identifier(SignaturesIdentifier::new(
@@ -906,9 +901,14 @@ impl ChiselDispatcher {
             )?)
             .build();
 
-        for (_, trace) in &mut result.traces {
-            // decoder.identify(trace, &mut local_identifier);
-            decoder.identify(trace, &mut etherscan_identifier);
+        let mut identifier = TraceIdentifiers::new().with_etherscan(
+            &session_config.foundry_config,
+            session_config.evm_opts.get_remote_chain_id(),
+        )?;
+        if !identifier.is_empty() {
+            for (_, trace) in &mut result.traces {
+                decoder.identify(trace, &mut identifier);
+            }
         }
         Ok(decoder)
     }

crates/cli/src/utils/cmd.rs

@@ -397,12 +397,18 @@ pub async fn handle_traces(
         .build();
 
     let mut etherscan_identifier = EtherscanIdentifier::new(config, chain)?;
-    for (_, trace) in &mut result.traces {
-        decoder.identify(trace, &mut etherscan_identifier);
+    if let Some(etherscan_identifier) = &mut etherscan_identifier {
+        for (_, trace) in &mut result.traces {
+            decoder.identify(trace, etherscan_identifier);
+        }
     }
 
     if debug {
-        let sources = etherscan_identifier.get_compiled_contracts().await?;
+        let sources = if let Some(etherscan_identifier) = etherscan_identifier {
+            etherscan_identifier.get_compiled_contracts().await?
+        } else {
+            Default::default()
+        };
         let mut debugger = Debugger::builder()
             .debug_arena(&result.debug)
             .decoder(&decoder)

crates/evm/traces/src/decoder/mod.rs

@@ -7,7 +7,9 @@ use crate::{
 use alloy_dyn_abi::{DecodedEvent, DynSolValue, EventExt, FunctionExt, JsonAbiExt};
 use alloy_json_abi::{Error, Event, Function, JsonAbi};
 use alloy_primitives::{Address, LogData, Selector, B256};
-use foundry_common::{abi::get_indexed_event, fmt::format_token, SELECTOR_LEN};
+use foundry_common::{
+    abi::get_indexed_event, fmt::format_token, ContractsByArtifact, SELECTOR_LEN,
+};
 use foundry_evm_core::{
     abi::{Console, HardhatConsole, Vm, HARDHAT_CONSOLE_SELECTOR_PATCHES},
     constants::{
@@ -50,17 +52,22 @@ impl CallTraceDecoderBuilder {
         self
     }
 
-    /// Add known contracts to the decoder from a `LocalTraceIdentifier`.
+    /// Add known contracts to the decoder.
     #[inline]
-    pub fn with_local_identifier_abis(mut self, identifier: &LocalTraceIdentifier<'_>) -> Self {
-        let contracts = identifier.contracts();
-        trace!(target: "evm::traces", len=contracts.len(), "collecting local identifier ABIs");
+    pub fn with_known_contracts(mut self, contracts: &ContractsByArtifact) -> Self {
+        trace!(target: "evm::traces", len=contracts.len(), "collecting known contract ABIs");
         for (abi, _) in contracts.values() {
             self.decoder.collect_abi(abi, None);
         }
         self
     }
 
+    /// Add known contracts to the decoder from a `LocalTraceIdentifier`.
+    #[inline]
+    pub fn with_local_identifier_abis(self, identifier: &LocalTraceIdentifier<'_>) -> Self {
+        self.with_known_contracts(identifier.contracts())
+    }
+
     /// Sets the verbosity level of the decoder.
     #[inline]
     pub fn with_verbosity(mut self, level: u8) -> Self {
@@ -225,7 +232,7 @@ impl CallTraceDecoder {
     fn addresses<'a>(
         &'a self,
         arena: &'a CallTraceArena,
-    ) -> impl Iterator<Item = (&'a Address, Option<&'a [u8]>)> + 'a {
+    ) -> impl Iterator<Item = (&'a Address, Option<&'a [u8]>)> + Clone + 'a {
         arena
             .nodes()
             .iter()

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

@@ -25,10 +25,9 @@ use std::{
 use tokio::time::{Duration, Interval};
 
 /// A trace identifier that tries to identify addresses using Etherscan.
-#[derive(Default)]
 pub struct EtherscanIdentifier {
     /// The Etherscan client
-    client: Option<Arc<foundry_block_explorers::Client>>,
+    client: Arc<foundry_block_explorers::Client>,
     /// Tracks whether the API key provides was marked as invalid
     ///
     /// After the first [EtherscanError::InvalidApiKey] this will get set to true, so we can
@@ -40,22 +39,21 @@ pub struct EtherscanIdentifier {
 
 impl EtherscanIdentifier {
     /// Creates a new Etherscan identifier with the given client
-    pub fn new(config: &Config, chain: Option<Chain>) -> eyre::Result<Self> {
+    pub fn new(config: &Config, chain: Option<Chain>) -> eyre::Result<Option<Self>> {
+        // In offline mode, don't use Etherscan.
         if config.offline {
-            // offline mode, don't use etherscan
-            return Ok(Default::default())
-        }
-        if let Some(config) = config.get_etherscan_config_with_chain(chain)? {
-            trace!(target: "etherscanidentifier", chain=?config.chain, url=?config.api_url, "using etherscan identifier");
-            Ok(Self {
-                client: Some(Arc::new(config.into_client()?)),
-                invalid_api_key: Arc::new(Default::default()),
-                contracts: BTreeMap::new(),
-                sources: BTreeMap::new(),
-            })
-        } else {
-            Ok(Default::default())
+            return Ok(None);
         }
+        let Some(config) = config.get_etherscan_config_with_chain(chain)? else {
+            return Ok(None);
+        };
+        trace!(target: "traces::etherscan", chain=?config.chain, url=?config.api_url, "using etherscan identifier");
+        Ok(Some(Self {
+            client: Arc::new(config.into_client()?),
+            invalid_api_key: Arc::new(AtomicBool::new(false)),
+            contracts: BTreeMap::new(),
+            sources: BTreeMap::new(),
+        }))
     }
 
     /// Goes over the list of contracts we have pulled from the traces, clones their source from
@@ -101,18 +99,13 @@ impl TraceIdentifier for EtherscanIdentifier {
     {
         trace!(target: "evm::traces", "identify {:?} addresses", addresses.size_hint().1);
 
-        let Some(client) = self.client.clone() else {
-            // no client was configured
-            return Vec::new()
-        };
-
         if self.invalid_api_key.load(Ordering::Relaxed) {
             // api key was marked as invalid
             return Vec::new()
         }
 
         let mut fetcher = EtherscanFetcher::new(
-            client,
+            self.client.clone(),
             Duration::from_secs(1),
             5,
             Arc::clone(&self.invalid_api_key),
@@ -191,16 +184,13 @@ impl EtherscanFetcher {
 
     fn queue_next_reqs(&mut self) {
         while self.in_progress.len() < self.concurrency {
-            if let Some(addr) = self.queue.pop() {
-                let client = Arc::clone(&self.client);
-                trace!(target: "etherscanidentifier", "fetching info for {:?}", addr);
-                self.in_progress.push(Box::pin(async move {
-                    let res = client.contract_source_code(addr).await;
-                    (addr, res)
-                }));
-            } else {
-                break
-            }
+            let Some(addr) = self.queue.pop() else { break };
+            let client = Arc::clone(&self.client);
+            self.in_progress.push(Box::pin(async move {
+                trace!(target: "traces::etherscan", ?addr, "fetching info");
+                let res = client.contract_source_code(addr).await;
+                (addr, res)
+            }));
         }
     }
 }
@@ -234,24 +224,24 @@ impl Stream for EtherscanFetcher {
                             }
                         }
                         Err(EtherscanError::RateLimitExceeded) => {
-                            warn!(target: "etherscanidentifier", "rate limit exceeded on attempt");
+                            warn!(target: "traces::etherscan", "rate limit exceeded on attempt");
                             pin.backoff = Some(tokio::time::interval(pin.timeout));
                             pin.queue.push(addr);
                         }
                         Err(EtherscanError::InvalidApiKey) => {
-                            warn!(target: "etherscanidentifier", "invalid api key");
+                            warn!(target: "traces::etherscan", "invalid api key");
                             // mark key as invalid
                             pin.invalid_api_key.store(true, Ordering::Relaxed);
                             return Poll::Ready(None)
                         }
                         Err(EtherscanError::BlockedByCloudflare) => {
-                            warn!(target: "etherscanidentifier", "blocked by cloudflare");
+                            warn!(target: "traces::etherscan", "blocked by cloudflare");
                             // mark key as invalid
                             pin.invalid_api_key.store(true, Ordering::Relaxed);
                             return Poll::Ready(None)
                         }
                         Err(err) => {
-                            warn!(target: "etherscanidentifier", "could not get etherscan info: {:?}", err);
+                            warn!(target: "traces::etherscan", "could not get etherscan info: {:?}", err);
                         }
                     }
                 }

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

@@ -1,6 +1,8 @@
 use alloy_json_abi::JsonAbi;
 use alloy_primitives::Address;
+use foundry_common::ContractsByArtifact;
 use foundry_compilers::ArtifactId;
+use foundry_config::{Chain, Config};
 use std::borrow::Cow;
 
 mod local;
@@ -33,5 +35,59 @@ pub trait TraceIdentifier {
     /// Attempts to identify an address in one or more call traces.
     fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity<'_>>
     where
-        A: Iterator<Item = (&'a Address, Option<&'a [u8]>)>;
+        A: Iterator<Item = (&'a Address, Option<&'a [u8]>)> + Clone;
+}
+
+/// A collection of trace identifiers.
+pub struct TraceIdentifiers<'a> {
+    /// The local trace identifier.
+    pub local: Option<LocalTraceIdentifier<'a>>,
+    /// The optional Etherscan trace identifier.
+    pub etherscan: Option<EtherscanIdentifier>,
+}
+
+impl Default for TraceIdentifiers<'_> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl TraceIdentifier for TraceIdentifiers<'_> {
+    fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity<'_>>
+    where
+        A: Iterator<Item = (&'a Address, Option<&'a [u8]>)> + Clone,
+    {
+        let mut identities = Vec::new();
+        if let Some(local) = &mut self.local {
+            identities.extend(local.identify_addresses(addresses.clone()));
+        }
+        if let Some(etherscan) = &mut self.etherscan {
+            identities.extend(etherscan.identify_addresses(addresses));
+        }
+        identities
+    }
+}
+
+impl<'a> TraceIdentifiers<'a> {
+    /// Creates a new, empty instance.
+    pub const fn new() -> Self {
+        Self { local: None, etherscan: None }
+    }
+
+    /// Sets the local identifier.
+    pub fn with_local(mut self, known_contracts: &'a ContractsByArtifact) -> Self {
+        self.local = Some(LocalTraceIdentifier::new(known_contracts));
+        self
+    }
+
+    /// Sets the etherscan identifier.
+    pub fn with_etherscan(mut self, config: &Config, chain: Option<Chain>) -> eyre::Result<Self> {
+        self.etherscan = EtherscanIdentifier::new(config, chain)?;
+        Ok(self)
+    }
+
+    /// Returns `true` if there are no set identifiers.
+    pub fn is_empty(&self) -> bool {
+        self.local.is_none() && self.etherscan.is_none()
+    }
 }

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

@@ -13,8 +13,8 @@ use forge::{
     decode::decode_console_logs,
     opts::EvmOpts,
     traces::{
-        identifier::{EtherscanIdentifier, LocalTraceIdentifier, SignaturesIdentifier},
-        render_trace_arena, CallTraceDecoder, CallTraceDecoderBuilder, TraceKind, Traces,
+        identifier::SignaturesIdentifier, render_trace_arena, CallTraceDecoder,
+        CallTraceDecoderBuilder, TraceKind, Traces,
     },
 };
 use forge_verify::RetryArgs;
@@ -42,6 +42,7 @@ use foundry_evm::{
     constants::DEFAULT_CREATE2_DEPLOYER,
     decode::RevertDecoder,
     inspectors::cheatcodes::{BroadcastableTransaction, BroadcastableTransactions},
+    traces::identifier::TraceIdentifiers,
 };
 use foundry_wallets::MultiWalletOpts;
 use futures::future;
@@ -198,33 +199,29 @@ impl ScriptArgs {
         result: &mut ScriptResult,
         known_contracts: &ContractsByArtifact,
     ) -> Result<CallTraceDecoder> {
-        let verbosity = script_config.evm_opts.verbosity;
-        let mut etherscan_identifier = EtherscanIdentifier::new(
-            &script_config.config,
-            script_config.evm_opts.get_remote_chain_id(),
-        )?;
-
-        let mut local_identifier = LocalTraceIdentifier::new(known_contracts);
         let mut decoder = CallTraceDecoderBuilder::new()
             .with_labels(result.labeled_addresses.clone())
-            .with_verbosity(verbosity)
-            .with_local_identifier_abis(&local_identifier)
+            .with_verbosity(script_config.evm_opts.verbosity)
+            .with_known_contracts(known_contracts)
             .with_signature_identifier(SignaturesIdentifier::new(
                 Config::foundry_cache_dir(),
                 script_config.config.offline,
             )?)
             .build();
+        let mut identifier = TraceIdentifiers::new()
+            .with_local(known_contracts)
+            .with_etherscan(&script_config.config, script_config.evm_opts.get_remote_chain_id())?;
 
         // Decoding traces using etherscan is costly as we run into rate limits,
         // causing scripts to run for a very long time unnecessarily.
         // Therefore, we only try and use etherscan if the user has provided an API key.
         let should_use_etherscan_traces = script_config.config.etherscan_api_key.is_some();
+        if !should_use_etherscan_traces {
+            identifier.etherscan = None;
+        }
 
         for (_, trace) in &mut result.traces {
-            decoder.identify(trace, &mut local_identifier);
-            if should_use_etherscan_traces {
-                decoder.identify(trace, &mut etherscan_identifier);
-            }
+            decoder.identify(trace, &mut identifier);
         }
         Ok(decoder)
     }