feat: full if-statement detection

foundry-rs · Jun 27, 2022 · 0b385d8 · 0b385d8
1 parent 8345057
commit 0b385d8
Show file tree

Hide file tree

Showing 3 changed files with 209 additions and 72 deletions.
diff --git a/cli/src/cmd/forge/coverage.rs b/cli/src/cmd/forge/coverage.rs
@@ -10,7 +10,7 @@ use crate::{
 use cast::trace::identifier::TraceIdentifier;
 use clap::{AppSettings, ArgEnum, Parser};
 use ethers::{
-    prelude::{Artifact, Project, ProjectCompileOutput},
+    prelude::{Artifact, Bytes, Project, ProjectCompileOutput},
     solc::{artifacts::contract::CompactContractBytecode, sourcemap::SourceMap, ArtifactId},
 };
 use forge::{
@@ -24,7 +24,7 @@ use forge::{
 };
 use foundry_common::{evm::EvmArgs, fs};
 use foundry_config::{figment::Figment, Config};
-use std::{collections::HashMap, path::PathBuf, sync::mpsc::channel, thread};
+use std::{borrow::Cow, collections::HashMap, path::PathBuf, sync::mpsc::channel, thread};
 
 // Loads project's figment and merges the build cli arguments into it
 foundry_config::impl_figment_convert!(CoverageArgs, opts, evm_opts);
@@ -119,15 +119,19 @@ impl CoverageArgs {
 
     /// Builds the coverage map.
     fn prepare(&self, output: ProjectCompileOutput) -> eyre::Result<(CoverageMap, SourceMaps)> {
-        // Get sources and source maps
+        // Extract artifacts
         let (artifacts, sources) = output.into_artifacts_with_sources();
-
-        let source_maps: SourceMaps = artifacts
+        let artifacts: HashMap<ArtifactId, CompactContractBytecode> = artifacts
             .into_iter()
             .map(|(id, artifact)| (id, CompactContractBytecode::from(artifact)))
-            .filter_map(|(id, artifact): (ArtifactId, CompactContractBytecode)| {
+            .collect();
+
+        // Get source maps
+        let source_maps: SourceMaps = artifacts
+            .iter()
+            .filter_map(|(id, artifact)| {
                 Some((
-                    id,
+                    id.clone(),
                     (
                         artifact.get_source_map()?.ok()?,
                         artifact
@@ -142,6 +146,17 @@ impl CoverageArgs {
             })
             .collect();
 
+        // Get bytecodes
+        let bytecodes: HashMap<ArtifactId, (Cow<Bytes>, Cow<Bytes>)> = artifacts
+            .iter()
+            .filter_map(|(id, artifact)| {
+                Some((
+                    id.clone(),
+                    (artifact.get_bytecode_bytes()?, artifact.get_deployed_bytecode_bytes()?),
+                ))
+            })
+            .collect();
+
         let mut map = CoverageMap::default();
         for (path, versioned_sources) in sources.0.into_iter() {
             // TODO: Make these checks robust
@@ -165,13 +180,24 @@ impl CoverageArgs {
                                 id.source == PathBuf::from(&path)
                         })
                         .map(|(id, (_, source_map))| {
-                            // TODO: Deploy source map too?
+                            // TODO: Deploy source map too
                             (id.name.clone(), source_map.clone())
                         })
                         .collect();
+                    let bytecodes: HashMap<String, Cow<Bytes>> = bytecodes
+                        .iter()
+                        .filter(|(id, _)| {
+                            id.version == versioned_source.version &&
+                                id.source == PathBuf::from(&path)
+                        })
+                        .map(|(id, (_, bytecode))| {
+                            // TODO: Deploy bytecode too
+                            (id.name.clone(), bytecode.clone())
+                        })
+                        .collect();
 
-                    let items =
-                        Visitor::new(fs::read_to_string(&path)?, source_maps).visit_ast(ast)?;
+                    let items = Visitor::new(fs::read_to_string(&path)?, source_maps, bytecodes)
+                        .visit_ast(ast)?;
 
                     if items.is_empty() {
                         continue

diff --git a/evm/src/coverage/mod.rs b/evm/src/coverage/mod.rs
@@ -213,8 +213,6 @@ pub enum CoverageItem {
         ///
         /// The first path has ID 0, the next ID 1, and so on.
         path_id: usize,
-        /// The branch kind.
-        kind: BranchKind,
         /// The number of times this item was hit.
         hits: u64,
     },
@@ -279,11 +277,8 @@ impl Display for CoverageItem {
             CoverageItem::Statement { loc, anchor, hits } => {
                 write!(f, "Statement (location: {loc}, anchor: {anchor}, hits: {hits})")
             }
-            CoverageItem::Branch { loc, anchor, hits, branch_id, path_id, kind } => {
-                write!(f, "{} Branch (branch: {branch_id}, path: {path_id}) (location: {loc}, anchor: {anchor}, hits: {hits})", match kind {
-                    BranchKind::True => "True",
-                    BranchKind::False => "False",
-                })
+            CoverageItem::Branch { loc, anchor, hits, branch_id, path_id } => {
+                write!(f, "Branch (branch: {branch_id}, path: {path_id}) (location: {loc}, anchor: {anchor}, hits: {hits})")
             }
             CoverageItem::Function { loc, anchor, hits, name } => {
                 write!(f, r#"Function "{name}" (location: {loc}, anchor: {anchor}, hits: {hits})"#)
@@ -314,14 +309,6 @@ impl Display for SourceLocation {
     }
 }
 
-#[derive(Debug, Clone)]
-pub enum BranchKind {
-    /// A false branch
-    True,
-    /// A true branch
-    False,
-}
-
 /// Coverage summary for a source file.
 #[derive(Default, Debug, Clone)]
 pub struct CoverageSummary {