cast: Fetch bytecodes in run/call to better match contracts

Without fetching the bytecodes from the current chain, matching the
contracts with `--with-local-artifacts` option only works if the
matching contracts have been deployed in the trace. This is very
limiting when trying to `--debug` an on-chain transaction.

By fetching the contracts' bytecodes, we can increase the matching of
address to source file, by thus providing the runtime bytecode.

Author: ngotchac

crates/cast/src/cmd/call.rs

@@ -34,6 +34,8 @@ use regex::Regex;
 use revm::context::TransactionType;
 use std::{str::FromStr, sync::LazyLock};
 
+use super::run::fetch_contracts_bytecode_from_trace;
+
 // matches override pattern <address>:<slot>:<value>
 // e.g. 0x123:0x1:0x1234
 static OVERRIDE_PATTERN: LazyLock<Regex> =
@@ -301,10 +303,12 @@ impl CallArgs {
                 ),
             };
 
+            let contracts_bytecode = fetch_contracts_bytecode_from_trace(&provider, &trace).await?;
             handle_traces(
                 trace,
                 &config,
                 chain,
+                &contracts_bytecode,
                 labels,
                 with_local_artifacts,
                 debug,

crates/cast/src/cmd/run.rs

@@ -1,6 +1,10 @@
 use alloy_consensus::Transaction;
 use alloy_network::{AnyNetwork, TransactionResponse};
-use alloy_provider::Provider;
+use alloy_primitives::{
+    map::{HashMap, HashSet},
+    Address, Bytes,
+};
+use alloy_provider::{Provider, RootProvider};
 use alloy_rpc_types::BlockTransactions;
 use clap::Parser;
 use eyre::{Result, WrapErr};
@@ -21,7 +25,7 @@ use foundry_config::{
 use foundry_evm::{
     executors::{EvmError, TracingExecutor},
     opts::EvmOpts,
-    traces::{InternalTraceMode, TraceMode},
+    traces::{InternalTraceMode, TraceMode, Traces},
     utils::configure_tx_env,
     Env,
 };
@@ -272,10 +276,12 @@ impl RunArgs {
             }
         };
 
+        let contracts_bytecode = fetch_contracts_bytecode_from_trace(&provider, &result).await?;
         handle_traces(
             result,
             &config,
             chain,
+            &contracts_bytecode,
             self.label,
             self.with_local_artifacts,
             self.debug,
@@ -287,6 +293,47 @@ impl RunArgs {
     }
 }
 
+pub async fn fetch_contracts_bytecode_from_trace(
+    provider: &RootProvider<AnyNetwork>,
+    result: &TraceResult,
+) -> Result<HashMap<Address, Bytes>> {
+    let mut contracts_bytecode = HashMap::default();
+    if let Some(ref traces) = result.traces {
+        let addresses = gather_trace_addresses(traces);
+        let results = futures::future::join_all(addresses.into_iter().map(async |a| {
+            (
+                a,
+                provider.get_code_at(a).await.unwrap_or_else(|e| {
+                    sh_warn!("Failed to fetch code for {a:?}: {e:?}").ok();
+                    Bytes::new()
+                }),
+            )
+        }))
+        .await;
+        for (address, code) in results {
+            if !code.is_empty() {
+                contracts_bytecode.insert(address, code);
+            }
+        }
+    }
+    Ok(contracts_bytecode)
+}
+
+fn gather_trace_addresses(traces: &Traces) -> HashSet<Address> {
+    let mut addresses = HashSet::default();
+    for (_, trace) in traces {
+        for node in trace.arena.nodes() {
+            if !node.trace.address.is_zero() {
+                addresses.insert(node.trace.address);
+            }
+            if !node.trace.caller.is_zero() {
+                addresses.insert(node.trace.caller);
+            }
+        }
+    }
+    addresses
+}
+
 impl figment::Provider for RunArgs {
     fn metadata(&self) -> Metadata {
         Metadata::named("RunArgs")

crates/cli/src/utils/cmd.rs

@@ -1,5 +1,5 @@
 use alloy_json_abi::JsonAbi;
-use alloy_primitives::Address;
+use alloy_primitives::{map::HashMap, Address, Bytes};
 use eyre::{Result, WrapErr};
 use foundry_common::{
     compile::ProjectCompiler, fs, selectors::SelectorKind, shell, ContractsByArtifact,
@@ -330,10 +330,12 @@ impl TryFrom<Result<RawCallResult>> for TraceResult {
 }
 
 /// labels the traces, conditionally prints them or opens the debugger
+#[expect(clippy::too_many_arguments)]
 pub async fn handle_traces(
     mut result: TraceResult,
     config: &Config,
     chain: Option<Chain>,
+    contracts_bytecode: &HashMap<Address, Bytes>,
     labels: Vec<String>,
     with_local_artifacts: bool,
     debug: bool,
@@ -372,7 +374,7 @@ pub async fn handle_traces(
     let mut identifier = TraceIdentifiers::new().with_etherscan(config, chain)?;
     if let Some(contracts) = &known_contracts {
         builder = builder.with_known_contracts(contracts);
-        identifier = identifier.with_local(contracts);
+        identifier = identifier.with_local_and_bytecodes(contracts, contracts_bytecode);
     }
 
     let mut decoder = builder.build();

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

@@ -1,6 +1,7 @@
 use super::{IdentifiedAddress, TraceIdentifier};
 use alloy_dyn_abi::JsonAbiExt;
 use alloy_json_abi::JsonAbi;
+use alloy_primitives::{map::HashMap, Address, Bytes};
 use foundry_common::contracts::{bytecode_diff_score, ContractsByArtifact};
 use foundry_compilers::ArtifactId;
 use revm_inspectors::tracing::types::CallTraceNode;
@@ -12,6 +13,8 @@ pub struct LocalTraceIdentifier<'a> {
     known_contracts: &'a ContractsByArtifact,
     /// Vector of pairs of artifact ID and the runtime code length of the given artifact.
     ordered_ids: Vec<(&'a ArtifactId, usize)>,
+    /// The contracts bytecode.
+    contracts_bytecode: Option<&'a HashMap<Address, Bytes>>,
 }
 
 impl<'a> LocalTraceIdentifier<'a> {
@@ -24,7 +27,12 @@ impl<'a> LocalTraceIdentifier<'a> {
             .map(|(id, bytecode)| (id, bytecode.len()))
             .collect::<Vec<_>>();
         ordered_ids.sort_by_key(|(_, len)| *len);
-        Self { known_contracts, ordered_ids }
+        Self { known_contracts, ordered_ids, contracts_bytecode: None }
+    }
+
+    pub fn with_bytecodes(mut self, contracts_bytecode: &'a HashMap<Address, Bytes>) -> Self {
+        self.contracts_bytecode = Some(contracts_bytecode);
+        self
     }
 
     /// Returns the known contracts.
@@ -161,7 +169,18 @@ impl TraceIdentifier for LocalTraceIdentifier<'_> {
                 let _span =
                     trace_span!(target: "evm::traces::local", "identify", %address).entered();
 
-                let (id, abi) = self.identify_code(runtime_code?, creation_code?)?;
+                // In order to identify the addresses, we need at least the runtime code. It can be
+                // obtained from the trace itself (if it's a CREATE* call), or from the fetched
+                // bytecodes.
+                let (runtime_code, creation_code) = match (runtime_code, creation_code) {
+                    (Some(runtime_code), Some(creation_code)) => (runtime_code, creation_code),
+                    (Some(runtime_code), _) => (runtime_code, &[] as &[u8]),
+                    _ => {
+                        let code = self.contracts_bytecode?.get(&address)?;
+                        (code.as_ref(), &[] as &[u8])
+                    }
+                };
+                let (id, abi) = self.identify_code(runtime_code, creation_code)?;
                 trace!(target: "evm::traces::local", id=%id.identifier(), "identified");
 
                 Some(IdentifiedAddress {

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

@@ -1,5 +1,5 @@
 use alloy_json_abi::JsonAbi;
-use alloy_primitives::Address;
+use alloy_primitives::{map::HashMap, Address, Bytes};
 use foundry_common::ContractsByArtifact;
 use foundry_compilers::ArtifactId;
 use foundry_config::{Chain, Config};
@@ -43,6 +43,8 @@ pub struct TraceIdentifiers<'a> {
     pub local: Option<LocalTraceIdentifier<'a>>,
     /// The optional Etherscan trace identifier.
     pub etherscan: Option<EtherscanIdentifier>,
+    /// The contracts bytecode.
+    pub contracts_bytecode: Option<&'a HashMap<Address, Bytes>>,
 }
 
 impl Default for TraceIdentifiers<'_> {
@@ -70,7 +72,7 @@ impl TraceIdentifier for TraceIdentifiers<'_> {
 impl<'a> TraceIdentifiers<'a> {
     /// Creates a new, empty instance.
     pub const fn new() -> Self {
-        Self { local: None, etherscan: None }
+        Self { local: None, etherscan: None, contracts_bytecode: None }
     }
 
     /// Sets the local identifier.
@@ -79,6 +81,17 @@ impl<'a> TraceIdentifiers<'a> {
         self
     }
 
+    /// Sets the local identifier.
+    pub fn with_local_and_bytecodes(
+        mut self,
+        known_contracts: &'a ContractsByArtifact,
+        contracts_bytecode: &'a HashMap<Address, Bytes>,
+    ) -> Self {
+        self.local =
+            Some(LocalTraceIdentifier::new(known_contracts).with_bytecodes(contracts_bytecode));
+        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)?;