feat: map addresses to artifact ids

foundry-rs · Jun 3, 2022 · 0b091f3 · 0b091f3
1 parent c0a15e9
commit 0b091f3
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 11 deletions.
diff --git a/cli/src/cmd/forge/coverage.rs b/cli/src/cmd/forge/coverage.rs
@@ -7,8 +7,10 @@ use crate::{
     compile::ProjectCompiler,
     utils,
 };
+use cast::trace::identifier::TraceIdentifier;
 use clap::{AppSettings, ArgEnum, Parser};
 use ethers::{
+    abi::Address,
     prelude::{Artifact, Project, ProjectCompileOutput},
     solc::{
         artifacts::contract::CompactContractBytecode,
@@ -19,7 +21,7 @@ use ethers::{
 use forge::{
     coverage::{CoverageMap, CoverageReporter, LcovReporter, SummaryReporter, Visitor},
     executor::opts::EvmOpts,
-    trace::{identifier::LocalTraceIdentifier, CallTraceDecoder, CallTraceDecoderBuilder},
+    trace::identifier::LocalTraceIdentifier,
     MultiContractRunnerBuilder, SuiteResult,
 };
 use foundry_common::evm::EvmArgs;
@@ -59,7 +61,7 @@ impl CoverageArgs {
 
     /// Returns the currently configured [Config] and the extracted [EvmOpts] from that config
     pub fn config_and_evm_opts(&self) -> eyre::Result<(Config, EvmOpts)> {
-        // merge all configs
+        // Merge all configs
         let figment: Figment = self.into();
         let evm_opts = figment.extract()?;
         let config = Config::from_provider(figment).sanitized();
@@ -82,6 +84,7 @@ impl Cmd for CoverageArgs {
     }
 }
 
+// The main flow of the command itself
 impl CoverageArgs {
     /// Collects and adjusts configuration.
     fn configure(&self) -> eyre::Result<(Config, EvmOpts)> {
@@ -106,7 +109,6 @@ impl CoverageArgs {
             project
         };
 
-        // TODO: This does not strip file prefixes for `SourceFiles`...
         let output = ProjectCompiler::default()
             .compile(&project)?
             .with_stripped_file_prefixes(project.root());
@@ -206,13 +208,19 @@ impl CoverageArgs {
         let handle = thread::spawn(move || runner.test(&self.filter, Some(tx), false).unwrap());
         for mut result in rx.into_iter().flat_map(|(_, suite)| suite.test_results.into_values()) {
             if let Some(_hit_data) = result.coverage.take() {
-                let mut decoder =
-                    CallTraceDecoderBuilder::new().with_events(local_identifier.events()).build();
+                let mut identities: BTreeMap<Address, ArtifactId> = BTreeMap::new();
                 for (_, trace) in &mut result.traces {
-                    decoder.identify(trace, &local_identifier);
+                    local_identifier
+                        .identify_addresses(trace.addresses().into_iter().collect())
+                        .into_iter()
+                        .for_each(|identity| {
+                            if let Some(artifact_id) = identity.artifact_id {
+                                identities.insert(identity.address, artifact_id);
+                            }
+                        })
                 }
-                // TODO: We need an ArtifactId here for the addresses
-                let CallTraceDecoder { contracts: _, .. } = decoder;
+
+                println!("{identities:?}");
 
                 // ..
             }
@@ -244,5 +252,3 @@ pub enum CoverageReportKind {
     Summary,
     Lcov,
 }
-
-// TODO: Move reporters to own module
diff --git a/evm/src/trace/identifier/etherscan.rs b/evm/src/trace/identifier/etherscan.rs
@@ -65,6 +65,7 @@ impl TraceIdentifier for EtherscanIdentifier {
                     label: Some(label.clone()),
                     contract: Some(label),
                     abi: Some(Cow::Owned(abi)),
+                    artifact_id: None,
                 })
                 .collect();
 

diff --git a/evm/src/trace/identifier/local.rs b/evm/src/trace/identifier/local.rs
@@ -46,6 +46,7 @@ impl TraceIdentifier for LocalTraceIdentifier {
                     contract: Some(id.identifier()),
                     label: Some(id.name.clone()),
                     abi: Some(Cow::Borrowed(abi)),
+                    artifact_id: Some(id.clone()),
                 })
             })
             .collect()

diff --git a/evm/src/trace/identifier/mod.rs b/evm/src/trace/identifier/mod.rs
@@ -4,7 +4,10 @@ pub use local::LocalTraceIdentifier;
 mod etherscan;
 pub use etherscan::EtherscanIdentifier;
 
-use ethers::abi::{Abi, Address};
+use ethers::{
+    abi::{Abi, Address},
+    prelude::ArtifactId,
+};
 use std::borrow::Cow;
 
 /// An address identity
@@ -19,6 +22,8 @@ pub struct AddressIdentity<'a> {
     pub contract: Option<String>,
     /// The ABI of the contract at this address
     pub abi: Option<Cow<'a, Abi>>,
+    /// The artifact ID of the contract, if any.
+    pub artifact_id: Option<ArtifactId>,
 }
 
 /// Trace identifiers figure out what ABIs and labels belong to all the addresses of the trace.