Bytecode level coverage reporting (#6563)

* feat(forge): add bytecode level coverage report

* include source references in bytecode coverage

* map source ranges to linenumbers in bytecode coverage

* clippy fixes

* format fixes

* nits

---------

Co-authored-by: Enrique Ortiz <hi@enriqueortiz.dev>

Cargo.toml

@@ -168,6 +168,7 @@ serde_json = { version = "1.0", features = ["arbitrary_precision"] }
 toml = "0.8"
 tracing = "0.1"
 tracing-subscriber = "0.3"
+evm-disassembler = "0.3"
 
 axum = "0.6"
 hyper = "0.14"

crates/cast/Cargo.toml

@@ -34,7 +34,7 @@ ethers-core.workspace = true
 ethers-providers.workspace = true
 
 chrono.workspace = true
-evm-disassembler = "0.3"
+evm-disassembler.workspace = true
 eyre.workspace = true
 futures = "0.3"
 hex.workspace = true

crates/evm/coverage/src/lib.rs

@@ -8,13 +8,16 @@
 extern crate tracing;
 
 use alloy_primitives::{Bytes, B256};
+use foundry_compilers::sourcemap::SourceElement;
 use semver::Version;
 use std::{
  collections::{BTreeMap, HashMap},
  fmt::Display,
  ops::{AddAssign, Deref, DerefMut},
 };
 
+use eyre::{Context, Result};
+
 pub mod analysis;
 pub mod anchors;
 
@@ -35,6 +38,10 @@ pub struct CoverageReport {
  pub items: HashMap<Version, Vec<CoverageItem>>,
  /// All item anchors for the codebase, keyed by their contract ID.
  pub anchors: HashMap<ContractId, Vec<ItemAnchor>>,
+ /// All the bytecode hits for the codebase
+ pub bytecode_hits: HashMap<ContractId, HitMap>,
+ /// The bytecode -> source mappings
+ pub source_maps: HashMap<ContractId, (Vec<SourceElement>, Vec<SourceElement>)>,
 }
 
 impl CoverageReport {
@@ -49,6 +56,14 @@ impl CoverageReport {
  self.source_paths_to_ids.get(&(version, path))
  }
 
+ /// Add the source maps
+ pub fn add_source_maps(
+ &mut self,
+ source_maps: HashMap<ContractId, (Vec<SourceElement>, Vec<SourceElement>)>,
+ ) {
+ self.source_maps.extend(source_maps);
+ }
+
  /// Add coverage items to this report
  pub fn add_items(&mut self, version: Version, items: Vec<CoverageItem>) {
  self.items.entry(version).or_default().extend(items);
@@ -109,7 +124,20 @@ impl CoverageReport {
  ///
  /// This function should only be called *after* all the relevant sources have been processed and
  /// added to the map (see [add_source]).
- pub fn add_hit_map(&mut self, contract_id: &ContractId, hit_map: &HitMap) {
+ pub fn add_hit_map(&mut self, contract_id: &ContractId, hit_map: &HitMap) -> Result<()> {
+ // Add bytecode level hits
+ let e = self
+ .bytecode_hits
+ .entry(contract_id.clone())
+ .or_insert_with(|| HitMap::new(hit_map.bytecode.clone()));
+ e.merge(hit_map).context(format!(
+ "contract_id {:?}, hash {}, hash {}",
+ contract_id,
+ e.bytecode.clone(),
+ hit_map.bytecode.clone(),
+ ))?;
+
+ // Add source level hits
  if let Some(anchors) = self.anchors.get(contract_id) {
  for anchor in anchors {
  if let Some(hits) = hit_map.hits.get(&anchor.instruction) {
@@ -121,6 +149,7 @@ impl CoverageReport {
  }
  }
  }
+ Ok(())
  }
 }
 
@@ -174,6 +203,28 @@ impl HitMap {
  pub fn hit(&mut self, pc: usize) {
  *self.hits.entry(pc).or_default() += 1;
  }
+
+ /// Merge another hitmap into this, assuming the bytecode is consistent
+ pub fn merge(&mut self, other: &HitMap) -> Result<(), eyre::Report> {
+ for (pc, hits) in &other.hits {
+ *self.hits.entry(*pc).or_default() += hits;
+ }
+ Ok(())
+ }
+
+ pub fn consistent_bytecode(&self, hm1: &HitMap, hm2: &HitMap) -> bool {
+ // Consider the bytecodes consistent if they are the same out as far as the
+ // recorded hits
+ let len1 = hm1.hits.last_key_value();
+ let len2 = hm2.hits.last_key_value();
+ if let (Some(len1), Some(len2)) = (len1, len2) {
+ let len = std::cmp::max(len1.0, len2.0);
+ let ok = hm1.bytecode.0[..*len] == hm2.bytecode.0[..*len];
+ println!("consistent_bytecode: {}, {}, {}, {}", ok, len1.0, len2.0, len);
+ return ok;
+ }
+ true
+ }
 }
 
 /// A unique identifier for a contract

crates/forge/Cargo.toml

@@ -71,6 +71,7 @@ strum = { version = "0.25", features = ["derive"] }
 thiserror = "1"
 tokio = { version = "1", features = ["time"] }
 watchexec = "2.3.2"
+evm-disassembler.workspace = true
 
 # doc server
 axum = { workspace = true, features = ["ws"] }

crates/forge/bin/cmd/coverage.rs

@@ -4,8 +4,8 @@ use clap::{Parser, ValueEnum, ValueHint};
 use eyre::{Context, Result};
 use forge::{
  coverage::{
- analysis::SourceAnalyzer, anchors::find_anchors, ContractId, CoverageReport,
- CoverageReporter, DebugReporter, ItemAnchor, LcovReporter, SummaryReporter,
+ analysis::SourceAnalyzer, anchors::find_anchors, BytecodeReporter, ContractId,
+ CoverageReport, CoverageReporter, DebugReporter, ItemAnchor, LcovReporter, SummaryReporter,
  },
  inspectors::CheatsConfig,
  opts::EvmOpts,
@@ -161,6 +161,8 @@ impl CoverageArgs {
  let mut versioned_asts: HashMap<Version, HashMap<usize, Ast>> = HashMap::new();
  let mut versioned_sources: HashMap<Version, HashMap<usize, String>> = HashMap::new();
  for (path, mut source_file, version) in sources.into_sources_with_version() {
+ report.add_source(version.clone(), source_file.id as usize, path.clone());
+
  // Filter out dependencies
  if project_paths.has_library_ancestor(std::path::Path::new(&path)) {
  continue
@@ -180,7 +182,6 @@ impl CoverageArgs {
  fs::read_to_string(&file)
  .wrap_err("Could not read source code for analysis")?,
  );
- report.add_source(version, source_file.id as usize, path);
  }
  }
 
@@ -279,6 +280,8 @@ impl CoverageArgs {
  report.add_anchors(anchors);
  }
 
+ report.add_source_maps(source_maps);
+
  Ok(report)
  }
 
@@ -342,7 +345,7 @@ impl CoverageArgs {
  contract_name: artifact_id.name.clone(),
  },
  &hits,
- );
+ )?;
  }
  }
 
@@ -367,6 +370,12 @@ impl CoverageArgs {
  .report(&report)
  }
  }
+ CoverageReportKind::Bytecode => {
+ let destdir = root.join("bytecode-coverage");
+ fs::create_dir_all(&destdir)?;
+ BytecodeReporter::new(root.clone(), destdir).report(&report)?;
+ Ok(())
+ }
  CoverageReportKind::Debug => DebugReporter.report(&report),
  }?;
  }
@@ -385,6 +394,7 @@ pub enum CoverageReportKind {
  Summary,
  Lcov,
  Debug,
+ Bytecode,
 }
 
 /// Helper function that will link references in unlinked bytecode to the 0 address.

crates/forge/src/coverage.rs

@@ -1,8 +1,14 @@
 //! Coverage reports.
 
 use comfy_table::{presets::ASCII_MARKDOWN, Attribute, Cell, Color, Row, Table};
+use evm_disassembler::disassemble_bytes;
+use foundry_common::fs;
 pub use foundry_evm::coverage::*;
-use std::io::Write;
+use std::{
+ collections::{hash_map, HashMap},
+ io::Write,
+ path::PathBuf,
+};
 
 /// A coverage reporter.
 pub trait CoverageReporter {
@@ -170,3 +176,104 @@ impl CoverageReporter for DebugReporter {
  Ok(())
  }
 }
+
+pub struct BytecodeReporter {
+ root: PathBuf,
+ destdir: PathBuf,
+}
+
+impl BytecodeReporter {
+ pub fn new(root: PathBuf, destdir: PathBuf) -> BytecodeReporter {
+ Self { root, destdir }
+ }
+}
+
+impl CoverageReporter for BytecodeReporter {
+ fn report(self, report: &CoverageReport) -> eyre::Result<()> {
+ use std::fmt::Write;
+
+ let no_source_elements = Vec::new();
+ let mut line_number_cache = LineNumberCache::new(self.root.clone());
+
+ for (contract_id, hits) in &report.bytecode_hits {
+ let ops = disassemble_bytes(hits.bytecode.to_vec())?;
+ let mut formatted = String::new();
+
+ let source_elements =
+ report.source_maps.get(contract_id).map(|sm| &sm.1).unwrap_or(&no_source_elements);
+
+ for (code, source_element) in std::iter::zip(ops.iter(), source_elements) {
+ let hits = hits
+ .hits
+ .get(&(code.offset as usize))
+ .map(|h| format!("[{:03}]", h))
+ .unwrap_or(" ".to_owned());
+ let source_id = source_element.index;
+ let source_path = source_id.and_then(|i| {
+ report.source_paths.get(&(contract_id.version.clone(), i as usize))
+ });
+
+ let code = format!("{:?}", code);
+ let start = source_element.offset;
+ let end = source_element.offset + source_element.length;
+
+ if let Some(source_path) = source_path {
+ let (sline, spos) = line_number_cache.get_position(source_path, start)?;
+ let (eline, epos) = line_number_cache.get_position(source_path, end)?;
+ writeln!(
+ formatted,
+ "{} {:40} // {}: {}:{}-{}:{} ({}-{})",
+ hits, code, source_path, sline, spos, eline, epos, start, end
+ )?;
+ } else if let Some(source_id) = source_id {
+ writeln!(
+ formatted,
+ "{} {:40} // SRCID{}: ({}-{})",
+ hits, code, source_id, start, end
+ )?;
+ } else {
+ writeln!(formatted, "{} {:40}", hits, code)?;
+ }
+ }
+ fs::write(
+ &self.destdir.join(contract_id.contract_name.clone()).with_extension("asm"),
+ formatted,
+ )?;
+ }
+
+ Ok(())
+ }
+}
+
+/// Cache line number offsets for source files
+struct LineNumberCache {
+ root: PathBuf,
+ line_offsets: HashMap<String, Vec<usize>>,
+}
+
+impl LineNumberCache {
+ pub fn new(root: PathBuf) -> Self {
+ LineNumberCache { root, line_offsets: HashMap::new() }
+ }
+
+ pub fn get_position(&mut self, path: &str, offset: usize) -> eyre::Result<(usize, usize)> {
+ let line_offsets = match self.line_offsets.entry(path.to_owned()) {
+ hash_map::Entry::Occupied(o) => o.into_mut(),
+ hash_map::Entry::Vacant(v) => {
+ let text = fs::read_to_string(self.root.join(path))?;
+ let mut line_offsets = vec![0];
+ for line in text.lines() {
+ let line_offset = line.as_ptr() as usize - text.as_ptr() as usize;
+ line_offsets.push(line_offset);
+ }
+ v.insert(line_offsets)
+ }
+ };
+ let lo = match line_offsets.binary_search(&offset) {
+ Ok(lo) => lo,
+ Err(lo) => lo - 1,
+ };
+ let pos = offset - line_offsets.get(lo).unwrap() + 1;
+ Ok((lo, pos))
+ }
+}