fix(coverage): clean ups, use normalized source code for locations

crates/evm/coverage/src/analysis.rs

@@ -1,7 +1,10 @@
 use super::{CoverageItem, CoverageItemKind, SourceLocation};
 use alloy_primitives::map::HashMap;
 use foundry_common::TestFunctionExt;
-use foundry_compilers::artifacts::ast::{self, Ast, Node, NodeType};
+use foundry_compilers::artifacts::{
+    ast::{self, Ast, Node, NodeType},
+    Source,
+};
 use rayon::prelude::*;
 use std::sync::Arc;
 
@@ -19,7 +22,7 @@ pub struct ContractVisitor<'a> {
     /// The current branch ID
     branch_id: usize,
     /// Stores the last line we put in the items collection to ensure we don't push duplicate lines
-    last_line: usize,
+    last_line: u32,
 
     /// Coverage items
     pub items: Vec<CoverageItem>,
@@ -47,23 +50,20 @@ impl<'a> ContractVisitor<'a> {
     }
 
     fn visit_function_definition(&mut self, node: &Node) -> eyre::Result<()> {
-        let name: String =
-            node.attribute("name").ok_or_else(|| eyre::eyre!("Function has no name"))?;
+        let Some(body) = &node.body else { return Ok(()) };
 
         let kind: String =
             node.attribute("kind").ok_or_else(|| eyre::eyre!("Function has no kind"))?;
 
-        match &node.body {
-            Some(body) => {
-                // Do not add coverage item for constructors without statements.
-                if kind == "constructor" && !has_statements(body) {
-                    return Ok(())
-                }
-                self.push_item_kind(CoverageItemKind::Function { name }, &node.src);
-                self.visit_block(body)
-            }
-            _ => Ok(()),
+        let name: String =
+            node.attribute("name").ok_or_else(|| eyre::eyre!("Function has no name"))?;
+
+        // Do not add coverage item for constructors without statements.
+        if kind == "constructor" && !has_statements(body) {
+            return Ok(())
         }
+        self.push_item_kind(CoverageItemKind::Function { name }, &node.src);
+        self.visit_block(body)
     }
 
     fn visit_modifier_or_yul_fn_definition(&mut self, node: &Node) -> eyre::Result<()> {
@@ -454,30 +454,34 @@ impl<'a> ContractVisitor<'a> {
     /// collection (plus additional coverage line if item is a statement).
     fn push_item_kind(&mut self, kind: CoverageItemKind, src: &ast::LowFidelitySourceLocation) {
         let item = CoverageItem { kind, loc: self.source_location_for(src), hits: 0 };
-        // Push a line item if we haven't already
-        if matches!(item.kind, CoverageItemKind::Statement | CoverageItemKind::Branch { .. }) &&
-            self.last_line < item.loc.line
-        {
+
+        // Push a line item if we haven't already.
+        debug_assert!(!matches!(item.kind, CoverageItemKind::Line));
+        if self.last_line < item.loc.lines.start {
             self.items.push(CoverageItem {
                 kind: CoverageItemKind::Line,
                 loc: item.loc.clone(),
                 hits: 0,
             });
-            self.last_line = item.loc.line;
+            self.last_line = item.loc.lines.start;
         }
 
         self.items.push(item);
     }
 
     fn source_location_for(&self, loc: &ast::LowFidelitySourceLocation) -> SourceLocation {
-        let loc_start =
-            self.source.char_indices().map(|(i, _)| i).nth(loc.start).unwrap_or_default();
+        let bytes_start = loc.start as u32;
+        let bytes_end = (loc.start + loc.length.unwrap_or(0)) as u32;
+        let bytes = bytes_start..bytes_end;
+
+        let start_line = self.source[..bytes.start as usize].lines().count() as u32;
+        let n_lines = self.source[bytes.start as usize..bytes.end as usize].lines().count() as u32;
+        let lines = start_line..start_line + n_lines;
         SourceLocation {
             source_id: self.source_id,
             contract_name: self.contract_name.clone(),
-            start: loc.start as u32,
-            length: loc.length.map(|x| x as u32),
-            line: self.source[..loc_start].lines().count(),
+            bytes,
+            lines,
         }
     }
 }
@@ -556,7 +560,7 @@ impl<'a> SourceAnalyzer<'a> {
                         .attribute("name")
                         .ok_or_else(|| eyre::eyre!("Contract has no name"))?;
 
-                    let mut visitor = ContractVisitor::new(source_id, source, &name);
+                    let mut visitor = ContractVisitor::new(source_id, &source.content, &name);
                     visitor.visit_contract(node)?;
                     let mut items = visitor.items;
 
@@ -590,7 +594,7 @@ pub struct SourceFiles<'a> {
 #[derive(Debug)]
 pub struct SourceFile<'a> {
     /// The source code.
-    pub source: String,
+    pub source: Source,
     /// The AST of the source code.
     pub ast: &'a Ast,
 }

crates/evm/coverage/src/anchors.rs

@@ -177,14 +177,14 @@ fn is_in_source_range(element: &SourceElement, location: &SourceLocation) -> boo
     }
 
     // Needed because some source ranges in the source map mark the entire contract...
-    let is_within_start = element.offset() >= location.start;
+    let is_within_start = element.offset() >= location.bytes.start;
     if !is_within_start {
         return false;
     }
 
-    let start_of_ranges = location.start.max(element.offset());
-    let end_of_ranges = (location.start + location.length.unwrap_or_default())
-        .min(element.offset() + element.length());
+    let start_of_ranges = location.bytes.start.max(element.offset());
+    let end_of_ranges =
+        (location.bytes.start + location.len()).min(element.offset() + element.length());
     let within_ranges = start_of_ranges <= end_of_ranges;
     if !within_ranges {
         return false;

crates/evm/coverage/src/lib.rs

@@ -18,7 +18,7 @@ use semver::Version;
 use std::{
     collections::BTreeMap,
     fmt::Display,
-    ops::{AddAssign, Deref, DerefMut},
+    ops::{Deref, DerefMut, Range},
     path::{Path, PathBuf},
     sync::Arc,
 };
@@ -82,40 +82,29 @@ impl CoverageReport {
         self.anchors.extend(anchors);
     }
 
-    /// Get coverage summaries by source file path.
-    pub fn summary_by_file(&self) -> impl Iterator<Item = (PathBuf, CoverageSummary)> {
-        let mut summaries = BTreeMap::new();
-
-        for (version, items) in self.items.iter() {
-            for item in items {
-                let Some(path) =
-                    self.source_paths.get(&(version.clone(), item.loc.source_id)).cloned()
-                else {
-                    continue;
-                };
-                *summaries.entry(path).or_default() += item;
-            }
-        }
-
-        summaries.into_iter()
+    /// Returns an iterator over coverage summaries by source file path.
+    pub fn summary_by_file(&self) -> impl Iterator<Item = (&Path, CoverageSummary)> {
+        self.by_file(|summary: &mut CoverageSummary, item| summary.add_item(item))
     }
 
-    /// Get coverage items by source file path.
-    pub fn items_by_source(&self) -> impl Iterator<Item = (PathBuf, Vec<CoverageItem>)> {
-        let mut items_by_source: BTreeMap<_, Vec<_>> = BTreeMap::new();
+    /// Returns an iterator over coverage items by source file path.
+    pub fn items_by_file(&self) -> impl Iterator<Item = (&Path, Vec<&CoverageItem>)> {
+        self.by_file(|list: &mut Vec<_>, item| list.push(item))
+    }
 
-        for (version, items) in self.items.iter() {
+    fn by_file<'a, T: Default>(
+        &'a self,
+        mut f: impl FnMut(&mut T, &'a CoverageItem),
+    ) -> impl Iterator<Item = (&'a Path, T)> {
+        let mut by_file: BTreeMap<&Path, T> = BTreeMap::new();
+        for (version, items) in &self.items {
             for item in items {
-                let Some(path) =
-                    self.source_paths.get(&(version.clone(), item.loc.source_id)).cloned()
-                else {
-                    continue;
-                };
-                items_by_source.entry(path).or_default().push(item.clone());
+                let key = (version.clone(), item.loc.source_id);
+                let Some(path) = self.source_paths.get(&key) else { continue };
+                f(by_file.entry(path).or_default(), item);
             }
         }
-
-        items_by_source.into_iter()
+        by_file.into_iter()
     }
 
     /// Processes data from a [`HitMap`] and sets hit counts for coverage items in this coverage
@@ -345,30 +334,34 @@ impl Display for CoverageItem {
     }
 }
 
+/// A source location.
 #[derive(Clone, Debug)]
 pub struct SourceLocation {
     /// The source ID.
     pub source_id: usize,
     /// The contract this source range is in.
     pub contract_name: Arc<str>,
-    /// Start byte in the source code.
-    pub start: u32,
-    /// Number of bytes in the source code.
-    pub length: Option<u32>,
-    /// The line in the source code.
-    pub line: usize,
+    /// Byte range.
+    pub bytes: Range<u32>,
+    /// Line range. Indices are 1-based.
+    pub lines: Range<u32>,
 }
 
 impl Display for SourceLocation {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "source ID {}, line {}, chars {}-{}",
-            self.source_id,
-            self.line,
-            self.start,
-            self.length.map_or(self.start, |length| self.start + length)
-        )
+        write!(f, "source ID {}, lines {:?}, bytes {:?}", self.source_id, self.lines, self.bytes)
+    }
+}
+
+impl SourceLocation {
+    /// Returns the length of the byte range.
+    pub fn len(&self) -> u32 {
+        self.bytes.len() as u32
+    }
+
+    /// Returns true if the byte range is empty.
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
     }
 }
 
@@ -393,21 +386,43 @@ pub struct CoverageSummary {
     pub function_hits: usize,
 }
 
-impl AddAssign<&Self> for CoverageSummary {
-    fn add_assign(&mut self, other: &Self) {
-        self.line_count += other.line_count;
-        self.line_hits += other.line_hits;
-        self.statement_count += other.statement_count;
-        self.statement_hits += other.statement_hits;
-        self.branch_count += other.branch_count;
-        self.branch_hits += other.branch_hits;
-        self.function_count += other.function_count;
-        self.function_hits += other.function_hits;
+impl CoverageSummary {
+    /// Creates a new, empty coverage summary.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Creates a coverage summary from a collection of coverage items.
+    pub fn from_items<'a>(items: impl IntoIterator<Item = &'a CoverageItem>) -> Self {
+        let mut summary = Self::default();
+        summary.add_items(items);
+        summary
     }
-}
 
-impl AddAssign<&CoverageItem> for CoverageSummary {
-    fn add_assign(&mut self, item: &CoverageItem) {
+    /// Adds another coverage summary to this one.
+    pub fn merge(&mut self, other: &Self) {
+        let Self {
+            line_count,
+            line_hits,
+            statement_count,
+            statement_hits,
+            branch_count,
+            branch_hits,
+            function_count,
+            function_hits,
+        } = self;
+        *line_count += other.line_count;
+        *line_hits += other.line_hits;
+        *statement_count += other.statement_count;
+        *statement_hits += other.statement_hits;
+        *branch_count += other.branch_count;
+        *branch_hits += other.branch_hits;
+        *function_count += other.function_count;
+        *function_hits += other.function_hits;
+    }
+
+    /// Adds a coverage item to this summary.
+    pub fn add_item(&mut self, item: &CoverageItem) {
         match item.kind {
             CoverageItemKind::Line => {
                 self.line_count += 1;
@@ -435,4 +450,11 @@ impl AddAssign<&CoverageItem> for CoverageSummary {
             }
         }
     }
+
+    /// Adds multiple coverage items to this summary.
+    pub fn add_items<'a>(&mut self, items: impl IntoIterator<Item = &'a CoverageItem>) {
+        for item in items {
+            self.add_item(item);
+        }
+    }
 }

crates/forge/bin/cmd/coverage.rs

@@ -16,13 +16,16 @@ use forge::{
 use foundry_cli::utils::{LoadConfig, STATIC_FUZZ_SEED};
 use foundry_common::{compile::ProjectCompiler, fs};
 use foundry_compilers::{
-    artifacts::{sourcemap::SourceMap, CompactBytecode, CompactDeployedBytecode, SolcLanguage},
+    artifacts::{
+        sourcemap::SourceMap, CompactBytecode, CompactDeployedBytecode, SolcLanguage, Source,
+    },
     Artifact, ArtifactId, Project, ProjectCompileOutput,
 };
 use foundry_config::{Config, SolcReq};
 use rayon::prelude::*;
 use semver::Version;
 use std::{
+    io,
     path::{Path, PathBuf},
     sync::Arc,
 };
@@ -153,7 +156,7 @@ impl CoverageArgs {
 
                 let source = SourceFile {
                     ast,
-                    source: fs::read_to_string(&file)
+                    source: Source::read(&file)
                         .wrap_err("Could not read source code for analysis")?,
                 };
                 versioned_sources
@@ -290,19 +293,15 @@ impl CoverageArgs {
             match report_kind {
                 CoverageReportKind::Summary => SummaryReporter::default().report(&report),
                 CoverageReportKind::Lcov => {
-                    if let Some(report_file) = self.report_file {
-                        return LcovReporter::new(&mut fs::create_file(root.join(report_file))?)
-                            .report(&report)
-                    } else {
-                        return LcovReporter::new(&mut fs::create_file(root.join("lcov.info"))?)
-                            .report(&report)
-                    }
+                    let path =
+                        root.join(self.report_file.as_deref().unwrap_or("lcov.info".as_ref()));
+                    let mut file = io::BufWriter::new(fs::create_file(path)?);
+                    LcovReporter::new(&mut file).report(&report)
                 }
                 CoverageReportKind::Bytecode => {
                     let destdir = root.join("bytecode-coverage");
                     fs::create_dir_all(&destdir)?;
-                    BytecodeReporter::new(root.clone(), destdir).report(&report)?;
-                    Ok(())
+                    BytecodeReporter::new(root.clone(), destdir).report(&report)
                 }
                 CoverageReportKind::Debug => DebugReporter.report(&report),
             }?;