coverage: Hoist the splitting of visible macro invocations

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -159,10 +159,6 @@ impl CoverageSpan {
 ///    execution
 ///  * Carve out (leave uncovered) any span that will be counted by another MIR (notably, closures)
 struct CoverageSpansGenerator<'a> {
-    /// A `Span` covering the function body of the MIR (typically from left curly brace to right
-    /// curly brace).
-    body_span: Span,
-
     /// The BasicCoverageBlock Control Flow Graph (BCB CFG).
     basic_coverage_blocks: &'a CoverageGraph,
 
@@ -239,7 +235,6 @@ impl<'a> CoverageSpansGenerator<'a> {
         );
 
         let coverage_spans = Self {
-            body_span: hir_info.body_span,
             basic_coverage_blocks,
             sorted_spans_iter: sorted_spans.into_iter(),
             some_curr: None,
@@ -261,7 +256,6 @@ impl<'a> CoverageSpansGenerator<'a> {
             // span-processing steps don't make sense yet.
             if self.some_prev.is_none() {
                 debug!("  initial span");
-                self.maybe_push_macro_name_span();
                 continue;
             }
 
@@ -273,15 +267,13 @@ impl<'a> CoverageSpansGenerator<'a> {
                 debug!("  same bcb (and neither is a closure), merge with prev={prev:?}");
                 let prev = self.take_prev();
                 self.curr_mut().merge_from(&prev);
-                self.maybe_push_macro_name_span();
             // Note that curr.span may now differ from curr_original_span
             } else if prev.span.hi() <= curr.span.lo() {
                 debug!(
                     "  different bcbs and disjoint spans, so keep curr for next iter, and add prev={prev:?}",
                 );
                 let prev = self.take_prev();
                 self.refined_spans.push(prev);
-                self.maybe_push_macro_name_span();
             } else if prev.is_closure {
                 // drop any equal or overlapping span (`curr`) and keep `prev` to test again in the
                 // next iter
@@ -297,7 +289,6 @@ impl<'a> CoverageSpansGenerator<'a> {
                 self.update_pending_dups();
             } else {
                 self.cutoff_prev_at_overlapping_curr();
-                self.maybe_push_macro_name_span();
             }
         }
 
@@ -332,41 +323,6 @@ impl<'a> CoverageSpansGenerator<'a> {
         self.refined_spans
     }
 
-    /// If `curr` is part of a new macro expansion, carve out and push a separate
-    /// span that ends just after the macro name and its subsequent `!`.
-    fn maybe_push_macro_name_span(&mut self) {
-        let curr = self.curr();
-
-        let Some(visible_macro) = curr.visible_macro(self.body_span) else { return };
-        if let Some(prev) = &self.some_prev
-            && prev.expn_span.eq_ctxt(curr.expn_span)
-        {
-            return;
-        }
-
-        // The split point is relative to `curr_original_span`,
-        // because `curr.span` may have been merged with preceding spans.
-        let split_point_after_macro_bang = self.curr_original_span.lo()
-            + BytePos(visible_macro.as_str().len() as u32)
-            + BytePos(1); // add 1 for the `!`
-        debug_assert!(split_point_after_macro_bang <= curr.span.hi());
-        if split_point_after_macro_bang > curr.span.hi() {
-            // Something is wrong with the macro name span;
-            // return now to avoid emitting malformed mappings (e.g. #117788).
-            return;
-        }
-
-        let mut macro_name_cov = curr.clone();
-        macro_name_cov.span = macro_name_cov.span.with_hi(split_point_after_macro_bang);
-        self.curr_mut().span = curr.span.with_lo(split_point_after_macro_bang);
-
-        debug!(
-            "  and curr starts a new macro expansion, so add a new span just for \
-            the macro `{visible_macro}!`, new span={macro_name_cov:?}",
-        );
-        self.refined_spans.push(macro_name_cov);
-    }
-
     #[track_caller]
     fn curr(&self) -> &CoverageSpan {
         self.some_curr.as_ref().unwrap_or_else(|| bug!("some_curr is None (curr)"))

compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs

@@ -39,6 +39,7 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
 
     initial_spans.sort_by(|a, b| basic_coverage_blocks.cmp_in_dominator_order(a.bcb, b.bcb));
     remove_unwanted_macro_spans(&mut initial_spans);
+    split_visible_macro_spans(&mut initial_spans, hir_info);
 
     initial_spans.sort_by(|a, b| {
         // First sort by span start.
@@ -80,6 +81,41 @@ fn remove_unwanted_macro_spans(initial_spans: &mut Vec<CoverageSpan>) {
     });
 }
 
+/// When a span corresponds to a macro invocation that is visible from the
+/// function body, split it into two parts. The first part covers just the
+/// macro name plus `!`, and the second part covers the rest of the macro
+/// invocation. This seems to give better results for code that uses macros.
+fn split_visible_macro_spans(initial_spans: &mut Vec<CoverageSpan>, hir_info: &ExtractedHirInfo) {
+    let mut extra_spans = vec![];
+
+    initial_spans.retain(|covspan| {
+        if covspan.is_closure {
+            return true;
+        }
+
+        let Some(visible_macro) = covspan.visible_macro(hir_info.body_span) else { return true };
+
+        let split_len = visible_macro.as_str().len() as u32 + 1;
+        let (before, after) = covspan.span.split_at(split_len);
+        if !covspan.span.contains(before) || !covspan.span.contains(after) {
+            // Something is unexpectedly wrong with the split point.
+            // The debug assertion in `split_at` will have already caught this,
+            // but in release builds it's safer to do nothing and maybe get a
+            // bug report for unexpected coverage, rather than risk an ICE.
+            return true;
+        }
+
+        assert!(!covspan.is_closure);
+        extra_spans.push(CoverageSpan::new(before, covspan.expn_span, covspan.bcb, false));
+        extra_spans.push(CoverageSpan::new(after, covspan.expn_span, covspan.bcb, false));
+        false // Discard the original covspan that we just split.
+    });
+
+    // The newly-split spans are added at the end, so any previous sorting
+    // is not preserved.
+    initial_spans.extend(extra_spans);
+}
+
 // Generate a set of `CoverageSpan`s from the filtered set of `Statement`s and `Terminator`s of
 // the `BasicBlock`(s) in the given `BasicCoverageBlockData`. One `CoverageSpan` is generated
 // for each `Statement` and `Terminator`. (Note that subsequent stages of coverage analysis will