feat(evm): compile etherscan sources and use them on the debugger

cli/src/cmd/cast/run.rs

@@ -4,7 +4,7 @@ use cast::trace::CallTraceDecoder;
 use clap::Parser;
 use ethers::{
     abi::Address,
-    prelude::{Middleware, Provider},
+    prelude::{artifacts::ContractBytecodeSome, ArtifactId, Middleware, Provider},
     types::H256,
 };
 use forge::{
@@ -14,11 +14,7 @@ use forge::{
 };
 use foundry_config::Config;
 use foundry_utils::RuntimeOrHandle;
-use std::{
-    collections::{BTreeMap, HashMap},
-    str::FromStr,
-    time::Duration,
-};
+use std::{collections::BTreeMap, str::FromStr, time::Duration};
 use ui::{TUIExitReason, Tui, Ui};
 
 #[derive(Debug, Clone, Parser)]
@@ -131,11 +127,12 @@ impl RunArgs {
                 }
             };
 
-            let etherscan_identifier = EtherscanIdentifier::new(
+            let mut etherscan_identifier = EtherscanIdentifier::new(
                 evm_opts.get_remote_chain_id(),
                 config.etherscan_api_key,
                 Config::foundry_etherscan_cache_dir(evm_opts.get_chain_id()),
                 Duration::from_secs(24 * 60 * 60),
+                true,
             );
 
             let labeled_addresses: BTreeMap<Address, String> = self
@@ -156,11 +153,13 @@ impl RunArgs {
             let mut decoder = CallTraceDecoderBuilder::new().with_labels(labeled_addresses).build();
 
             for (_, trace) in &mut result.traces {
-                decoder.identify(trace, &etherscan_identifier);
+                decoder.identify(trace, &mut etherscan_identifier);
             }
 
+            let (sources, bytecode) =
+                etherscan_identifier.get_compiled_contracts_with_sources().await?;
             if self.debug {
-                run_debugger(result, decoder)?;
+                run_debugger(result, decoder, bytecode, sources)?;
             } else {
                 print_traces(&mut result, decoder)?;
             }
@@ -169,12 +168,29 @@ impl RunArgs {
     }
 }
 
-fn run_debugger(result: RunResult, decoder: CallTraceDecoder) -> eyre::Result<()> {
-    // TODO Get source from etherscan
-    let source_code: BTreeMap<u32, String> = BTreeMap::new();
+fn run_debugger(
+    result: RunResult,
+    decoder: CallTraceDecoder,
+    known_contracts: BTreeMap<ArtifactId, ContractBytecodeSome>,
+    sources: BTreeMap<ArtifactId, String>,
+) -> eyre::Result<()> {
     let calls: Vec<DebugArena> = vec![result.debug];
     let flattened = calls.last().expect("we should have collected debug info").flatten(0);
-    let tui = Tui::new(flattened, 0, decoder.contracts, HashMap::new(), source_code)?;
+
+    let tui = Tui::new(
+        flattened,
+        0,
+        decoder.contracts,
+        known_contracts.into_iter().map(|(id, artifact)| (id.name, artifact)).collect(),
+        sources
+            .into_iter()
+            .map(|(id, source)| {
+                let mut sources = BTreeMap::new();
+                sources.insert(0, source);
+                (id.name, sources)
+            })
+            .collect(),
+    )?;
     match tui.start().expect("Failed to start tui") {
         TUIExitReason::CharExit => Ok(()),
     }

cli/src/cmd/forge/run.rs

@@ -161,13 +161,13 @@ impl Cmd for RunArgs {
         // TODO: Could we use the Etherscan identifier here? Main issue: Pulling source code and
         // bytecode. Might be better to wait for an interactive debugger where we can do this on
         // the fly while retaining access to the database?
-        let local_identifier = LocalTraceIdentifier::new(&known_contracts);
+        let mut local_identifier = LocalTraceIdentifier::new(&known_contracts);
         let mut decoder = CallTraceDecoderBuilder::new()
             .with_labels(result.labeled_addresses.clone())
             .with_events(local_identifier.events())
             .build();
         for (_, trace) in &mut result.traces {
-            decoder.identify(trace, &local_identifier);
+            decoder.identify(trace, &mut local_identifier);
         }
 
         if self.debug {
@@ -195,10 +195,14 @@ impl Cmd for RunArgs {
                 0,
                 decoder.contracts,
                 highlevel_known_contracts
+                    .clone()
                     .into_iter()
                     .map(|(id, artifact)| (id.name, artifact))
                     .collect(),
-                source_code,
+                highlevel_known_contracts
+                    .into_iter()
+                    .map(|(id, _)| (id.name, source_code.clone()))
+                    .collect(),
             )?;
             match tui.start().expect("Failed to start tui") {
                 TUIExitReason::CharExit => return Ok(()),

cli/src/cmd/forge/test.rs

@@ -513,16 +513,17 @@ fn test(
         Ok(TestOutcome::new(results, allow_failure))
     } else {
         // Set up identifiers
-        let local_identifier = LocalTraceIdentifier::new(&runner.known_contracts);
+        let mut local_identifier = LocalTraceIdentifier::new(&runner.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.
         // TODO: Make this configurable.
         let cache_ttl = Duration::from_secs(24 * 60 * 60);
-        let etherscan_identifier = EtherscanIdentifier::new(
+        let mut etherscan_identifier = EtherscanIdentifier::new(
             remote_chain_id,
             config.etherscan_api_key,
             remote_chain_id.and_then(Config::foundry_etherscan_cache_dir),
             cache_ttl,
+            false,
         );
 
         // Set up test reporter channel
@@ -570,8 +571,8 @@ fn test(
                     // Decode the traces
                     let mut decoded_traces = Vec::new();
                     for (kind, trace) in &mut result.traces {
-                        decoder.identify(trace, &local_identifier);
-                        decoder.identify(trace, &etherscan_identifier);
+                        decoder.identify(trace, &mut local_identifier);
+                        decoder.identify(trace, &mut etherscan_identifier);
 
                         let should_include = match kind {
                             // At verbosity level 3, we only display traces for failed tests

evm/src/trace/decoder.rs

@@ -176,7 +176,7 @@ impl CallTraceDecoder {
     /// Identify unknown addresses in the specified call trace using the specified identifier.
     ///
     /// Unknown contracts are contracts that either lack a label or an ABI.
-    pub fn identify(&mut self, trace: &CallTraceArena, identifier: &impl TraceIdentifier) {
+    pub fn identify(&mut self, trace: &CallTraceArena, identifier: &mut impl TraceIdentifier) {
         let unidentified_addresses = trace
             .addresses()
             .into_iter()

evm/src/trace/identifier/etherscan.rs

@@ -2,22 +2,29 @@ use super::{AddressIdentity, TraceIdentifier};
 use ethers::{
     abi::{Abi, Address},
     etherscan,
-    prelude::{contract::ContractMetadata, errors::EtherscanError},
+    prelude::{
+        artifacts::ContractBytecodeSome, contract::ContractMetadata, errors::EtherscanError,
+        ArtifactId,
+    },
     types::Chain,
 };
 use futures::{
     future::Future,
     stream::{FuturesUnordered, Stream, StreamExt},
     task::{Context, Poll},
 };
-use std::{borrow::Cow, path::PathBuf, pin::Pin};
+use std::{borrow::Cow, collections::BTreeMap, path::PathBuf, pin::Pin};
 use tokio::time::{Duration, Interval};
 use tracing::{trace, warn};
 
 /// A trace identifier that tries to identify addresses using Etherscan.
 pub struct EtherscanIdentifier {
     /// The Etherscan client
     client: Option<etherscan::Client>,
+    // cache_path: Option<PathBuf>,
+    pub contracts: BTreeMap<String, (String, bool, u32, String)>,
+    pub sources: BTreeMap<u32, String>,
+    compiles: bool,
 }
 
 impl EtherscanIdentifier {
@@ -29,6 +36,7 @@ impl EtherscanIdentifier {
         etherscan_api_key: Option<String>,
         cache_path: Option<PathBuf>,
         ttl: Duration,
+        compiles: bool,
     ) -> Self {
         if let Some(cache_path) = &cache_path {
             if let Err(err) = std::fs::create_dir_all(cache_path.join("sources")) {
@@ -39,16 +47,57 @@ impl EtherscanIdentifier {
         Self {
             client: chain.and_then(|chain| {
                 etherscan_api_key.and_then(|key| {
-                    etherscan::Client::new_cached(chain.into(), key, cache_path, ttl).ok()
+                    etherscan::Client::new_cached(chain.into(), key, cache_path.clone(), ttl).ok()
                 })
             }),
+            contracts: BTreeMap::new(),
+            // cache_path,
+            compiles,
+            sources: BTreeMap::new(),
+        }
+    }
+
+    pub async fn get_compiled_contracts_with_sources(
+        &self,
+    ) -> eyre::Result<(BTreeMap<ArtifactId, String>, BTreeMap<ArtifactId, ContractBytecodeSome>)>
+    {
+        if self.compiles {
+            let mut compiled_contracts = BTreeMap::new();
+            let mut sources = BTreeMap::new();
+
+            for (label, (source_code, optimization, runs, version)) in &self.contracts {
+                // todo skip multi-file, v0.4 and vyper for now
+                if source_code.starts_with("{{") ||
+                    version.starts_with("v0.4") ||
+                    version.starts_with("vyper")
+                {
+                    continue
+                }
+
+                println!("Compiling {label}");
+                let compiled = foundry_utils::compile_contract_source(
+                    label.to_string(),
+                    source_code.clone(),
+                    *optimization,
+                    *runs,
+                    version.to_string(),
+                )
+                .await?;
+
+                compiled_contracts.insert(compiled.0.clone(), compiled.1.to_owned());
+                sources.insert(compiled.0.to_owned(), source_code.to_owned());
+            }
+
+            Ok((sources, compiled_contracts))
+        } else {
+            Ok((BTreeMap::new(), BTreeMap::new()))
         }
     }
 }
 
 impl TraceIdentifier for EtherscanIdentifier {
     fn identify_addresses(
-        &self,
+        &mut self,
         addresses: Vec<(&Address, Option<&Vec<u8>>)>,
     ) -> Vec<AddressIdentity> {
         self.client.as_ref().map_or(Default::default(), |client| {
@@ -57,13 +106,17 @@ impl TraceIdentifier for EtherscanIdentifier {
             for (addr, _) in addresses {
                 fetcher.push(*addr);
             }
-
             let fut = fetcher
-                .map(|(address, label, abi)| AddressIdentity {
-                    address,
-                    label: Some(label.clone()),
-                    contract: Some(label),
-                    abi: Some(Cow::Owned(abi)),
+                .map(|(address, label, abi, source_code, optimization, runs, version)| {
+                    self.contracts
+                        .insert(label.clone(), (source_code, optimization, runs, version));
+
+                    AddressIdentity {
+                        address,
+                        label: Some(label.clone()),
+                        contract: Some(label),
+                        abi: Some(Cow::Owned(abi)),
+                    }
                 })
                 .collect();
 
@@ -126,7 +179,7 @@ impl EtherscanFetcher {
 }
 
 impl Stream for EtherscanFetcher {
-    type Item = (Address, String, Abi);
+    type Item = (Address, String, Abi, String, bool, u32, String);
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let pin = self.get_mut();
@@ -151,7 +204,15 @@ impl Stream for EtherscanFetcher {
                         Ok(mut metadata) => {
                             if let Some(item) = metadata.items.pop() {
                                 if let Ok(abi) = serde_json::from_str(&item.abi) {
-                                    return Poll::Ready(Some((addr, item.contract_name, abi)))
+                                    return Poll::Ready(Some((
+                                        addr,
+                                        item.contract_name,
+                                        abi,
+                                        item.source_code,
+                                        item.optimization_used.eq("1"),
+                                        item.runs.parse::<u32>().expect("runs parse error"),
+                                        item.compiler_version,
+                                    )))
                                 }
                             }
                         }

evm/src/trace/identifier/local.rs

@@ -30,7 +30,7 @@ impl LocalTraceIdentifier {
 
 impl TraceIdentifier for LocalTraceIdentifier {
     fn identify_addresses(
-        &self,
+        &mut self,
         addresses: Vec<(&Address, Option<&Vec<u8>>)>,
     ) -> Vec<AddressIdentity> {
         addresses

evm/src/trace/identifier/mod.rs

@@ -27,7 +27,7 @@ pub trait TraceIdentifier {
     /// Attempts to identify an address in one or more call traces.
     #[allow(clippy::type_complexity)]
     fn identify_addresses(
-        &self,
+        &mut self,
         addresses: Vec<(&Address, Option<&Vec<u8>>)>,
     ) -> Vec<AddressIdentity>;
 }