Rollup merge of #149471 - Zalathar:tree, r=oli-obk

coverage: Store signature/body spans and branch spans in the expansion tree

In order to support coverage instrumentation of expansion regions, we need to reduce the amount of code that assumes we're only instrumenting a flat function body. Moving more data into expansion tree nodes is an incremental step in that direction.

There should be no change to compiler output.

Author: matthiaskrgr

compiler/rustc_mir_transform/src/coverage/expansion.rs

@@ -1,6 +1,6 @@
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet, IndexEntry};
 use rustc_middle::mir;
-use rustc_middle::mir::coverage::BasicCoverageBlock;
+use rustc_middle::mir::coverage::{BasicCoverageBlock, BranchSpan};
 use rustc_span::{ExpnId, ExpnKind, Span};
 
 use crate::coverage::from_mir;
@@ -71,11 +71,21 @@ pub(crate) struct ExpnNode {
     /// This links an expansion node to its parent in the tree.
     pub(crate) call_site_expn_id: Option<ExpnId>,
 
+    /// Holds the function signature span, if it belongs to this expansion.
+    /// Used by special-case code in span refinement.
+    pub(crate) fn_sig_span: Option<Span>,
+    /// Holds the function body span, if it belongs to this expansion.
+    /// Used by special-case code in span refinement.
+    pub(crate) body_span: Option<Span>,
+
     /// Spans (and their associated BCBs) belonging to this expansion.
     pub(crate) spans: Vec<SpanWithBcb>,
     /// Expansions whose call-site is in this expansion.
     pub(crate) child_expn_ids: FxIndexSet<ExpnId>,
 
+    /// Branch spans (recorded during MIR building) belonging to this expansion.
+    pub(crate) branch_spans: Vec<BranchSpan>,
+
     /// Hole spans belonging to this expansion, to be carved out from the
     /// code spans during span refinement.
     pub(crate) hole_spans: Vec<Span>,
@@ -95,9 +105,14 @@ impl ExpnNode {
             call_site,
             call_site_expn_id,
 
+            fn_sig_span: None,
+            body_span: None,
+
             spans: vec![],
             child_expn_ids: FxIndexSet::default(),
 
+            branch_spans: vec![],
+
             hole_spans: vec![],
         }
     }
@@ -142,6 +157,20 @@ pub(crate) fn build_expn_tree(
         }
     }
 
+    // If we have a span for the function signature, associate it with the
+    // corresponding expansion tree node.
+    if let Some(fn_sig_span) = hir_info.fn_sig_span
+        && let Some(node) = nodes.get_mut(&fn_sig_span.ctxt().outer_expn())
+    {
+        node.fn_sig_span = Some(fn_sig_span);
+    }
+
+    // Also associate the body span with its expansion tree node.
+    let body_span = hir_info.body_span;
+    if let Some(node) = nodes.get_mut(&body_span.ctxt().outer_expn()) {
+        node.body_span = Some(body_span);
+    }
+
     // Associate each hole span (extracted from HIR) with its corresponding
     // expansion tree node.
     for &hole_span in &hir_info.hole_spans {
@@ -150,5 +179,15 @@ pub(crate) fn build_expn_tree(
         node.hole_spans.push(hole_span);
     }
 
+    // Associate each branch span (recorded during MIR building) with its
+    // corresponding expansion tree node.
+    if let Some(coverage_info_hi) = mir_body.coverage_info_hi.as_deref() {
+        for branch_span in &coverage_info_hi.branch_spans {
+            if let Some(node) = nodes.get_mut(&branch_span.span.ctxt().outer_expn()) {
+                node.branch_spans.push(BranchSpan::clone(branch_span));
+            }
+        }
+    }
+
     ExpnTree { nodes }
 }

compiler/rustc_mir_transform/src/coverage/mappings.rs

@@ -4,12 +4,12 @@ use rustc_middle::mir::coverage::{
 };
 use rustc_middle::mir::{self, BasicBlock, StatementKind};
 use rustc_middle::ty::TyCtxt;
+use rustc_span::ExpnKind;
 
-use crate::coverage::expansion;
+use crate::coverage::expansion::{self, ExpnTree};
 use crate::coverage::graph::CoverageGraph;
 use crate::coverage::hir_info::ExtractedHirInfo;
 use crate::coverage::spans::extract_refined_covspans;
-use crate::coverage::unexpand::unexpand_into_body_span;
 
 #[derive(Default)]
 pub(crate) struct ExtractedMappings {
@@ -31,7 +31,7 @@ pub(crate) fn extract_mappings_from_mir<'tcx>(
     // Extract ordinary code mappings from MIR statement/terminator spans.
     extract_refined_covspans(tcx, hir_info, graph, &expn_tree, &mut mappings);
 
-    extract_branch_mappings(mir_body, hir_info, graph, &mut mappings);
+    extract_branch_mappings(mir_body, hir_info, graph, &expn_tree, &mut mappings);
 
     ExtractedMappings { mappings }
 }
@@ -57,25 +57,25 @@ fn resolve_block_markers(
     block_markers
 }
 
-pub(super) fn extract_branch_mappings(
+fn extract_branch_mappings(
     mir_body: &mir::Body<'_>,
     hir_info: &ExtractedHirInfo,
     graph: &CoverageGraph,
+    expn_tree: &ExpnTree,
     mappings: &mut Vec<Mapping>,
 ) {
     let Some(coverage_info_hi) = mir_body.coverage_info_hi.as_deref() else { return };
-
     let block_markers = resolve_block_markers(coverage_info_hi, mir_body);
 
-    mappings.extend(coverage_info_hi.branch_spans.iter().filter_map(
-        |&BranchSpan { span: raw_span, true_marker, false_marker }| try {
-            // For now, ignore any branch span that was introduced by
-            // expansion. This makes things like assert macros less noisy.
-            if !raw_span.ctxt().outer_expn_data().is_root() {
-                return None;
-            }
-            let span = unexpand_into_body_span(raw_span, hir_info.body_span)?;
+    // For now, ignore any branch span that was introduced by
+    // expansion. This makes things like assert macros less noisy.
+    let Some(node) = expn_tree.get(hir_info.body_span.ctxt().outer_expn()) else { return };
+    if node.expn_kind != ExpnKind::Root {
+        return;
+    }
 
+    mappings.extend(node.branch_spans.iter().filter_map(
+        |&BranchSpan { span, true_marker, false_marker }| try {
             let bcb_from_marker = |marker: BlockMarkerId| graph.bcb_from_bb(block_markers[marker]?);
 
             let true_bcb = bcb_from_marker(true_marker)?;

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -17,7 +17,6 @@ pub(super) mod query;
 mod spans;
 #[cfg(test)]
 mod tests;
-mod unexpand;
 
 /// Inserts `StatementKind::Coverage` statements that either instrument the binary with injected
 /// counters, via intrinsic `llvm.instrprof.increment`, and/or inject metadata used during codegen

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -26,11 +26,9 @@ pub(super) fn extract_refined_covspans<'tcx>(
         return;
     }
 
-    let &ExtractedHirInfo { body_span, .. } = hir_info;
-
     // If there somehow isn't an expansion tree node corresponding to the
     // body span, return now and don't create any mappings.
-    let Some(node) = expn_tree.get(body_span.ctxt().outer_expn()) else { return };
+    let Some(node) = expn_tree.get(hir_info.body_span.ctxt().outer_expn()) else { return };
 
     let mut covspans = vec![];
 
@@ -47,27 +45,29 @@ pub(super) fn extract_refined_covspans<'tcx>(
         }
     }
 
-    covspans.retain(|covspan: &Covspan| {
-        let covspan_span = covspan.span;
-        // Discard any spans not contained within the function body span.
-        // Also discard any spans that fill the entire body, because they tend
-        // to represent compiler-inserted code, e.g. implicitly returning `()`.
-        if !body_span.contains(covspan_span) || body_span.source_equal(covspan_span) {
-            return false;
-        }
+    if let Some(body_span) = node.body_span {
+        covspans.retain(|covspan: &Covspan| {
+            let covspan_span = covspan.span;
+            // Discard any spans not contained within the function body span.
+            // Also discard any spans that fill the entire body, because they tend
+            // to represent compiler-inserted code, e.g. implicitly returning `()`.
+            if !body_span.contains(covspan_span) || body_span.source_equal(covspan_span) {
+                return false;
+            }
 
-        // Each pushed covspan should have the same context as the body span.
-        // If it somehow doesn't, discard the covspan, or panic in debug builds.
-        if !body_span.eq_ctxt(covspan_span) {
-            debug_assert!(
-                false,
-                "span context mismatch: body_span={body_span:?}, covspan.span={covspan_span:?}"
-            );
-            return false;
-        }
+            // Each pushed covspan should have the same context as the body span.
+            // If it somehow doesn't, discard the covspan, or panic in debug builds.
+            if !body_span.eq_ctxt(covspan_span) {
+                debug_assert!(
+                    false,
+                    "span context mismatch: body_span={body_span:?}, covspan.span={covspan_span:?}"
+                );
+                return false;
+            }
 
-        true
-    });
+            true
+        });
+    }
 
     // Only proceed if we found at least one usable span.
     if covspans.is_empty() {
@@ -78,10 +78,9 @@ pub(super) fn extract_refined_covspans<'tcx>(
     // Otherwise, add a fake span at the start of the body, to avoid an ugly
     // gap between the start of the body and the first real span.
     // FIXME: Find a more principled way to solve this problem.
-    covspans.push(Covspan {
-        span: hir_info.fn_sig_span.unwrap_or_else(|| body_span.shrink_to_lo()),
-        bcb: START_BCB,
-    });
+    if let Some(span) = node.fn_sig_span.or_else(|| try { node.body_span?.shrink_to_lo() }) {
+        covspans.push(Covspan { span, bcb: START_BCB });
+    }
 
     let compare_covspans = |a: &Covspan, b: &Covspan| {
         compare_spans(a.span, b.span)

compiler/rustc_mir_transform/src/coverage/unexpand.rs

@@ -0,0 +1,44 @@
+Function name: fn_in_macro::branch_in_macro
+Raw bytes (19): 0x[01, 01, 00, 03, 01, 22, 01, 00, 15, 01, 0b, 05, 00, 17, 01, 01, 01, 00, 02]
+Number of files: 1
+- file 0 => $DIR/fn-in-macro.rs
+Number of expressions: 0
+Number of file 0 mappings: 3
+- Code(Counter(0)) at (prev + 34, 1) to (start + 0, 21)
+- Code(Counter(0)) at (prev + 11, 5) to (start + 0, 23)
+- Code(Counter(0)) at (prev + 1, 1) to (start + 0, 2)
+Highest counter ID seen: c0
+
+Function name: fn_in_macro::fn_in_macro
+Raw bytes (31): 0x[01, 01, 01, 01, 05, 05, 01, 0c, 09, 00, 19, 01, 01, 10, 00, 25, 05, 00, 2c, 02, 0e, 02, 02, 14, 02, 0e, 01, 03, 09, 00, 0a]
+Number of files: 1
+- file 0 => $DIR/fn-in-macro.rs
+Number of expressions: 1
+- expression 0 operands: lhs = Counter(0), rhs = Counter(1)
+Number of file 0 mappings: 5
+- Code(Counter(0)) at (prev + 12, 9) to (start + 0, 25)
+- Code(Counter(0)) at (prev + 1, 16) to (start + 0, 37)
+- Code(Counter(1)) at (prev + 0, 44) to (start + 2, 14)
+- Code(Expression(0, Sub)) at (prev + 2, 20) to (start + 2, 14)
+    = (c0 - c1)
+- Code(Counter(0)) at (prev + 3, 9) to (start + 0, 10)
+Highest counter ID seen: c1
+
+Function name: fn_in_macro::fn_not_in_macro
+Raw bytes (38): 0x[01, 01, 01, 01, 05, 06, 01, 19, 01, 00, 15, 01, 01, 08, 00, 1d, 20, 05, 02, 00, 08, 00, 23, 05, 00, 24, 02, 06, 02, 02, 0c, 02, 06, 01, 03, 01, 00, 02]
+Number of files: 1
+- file 0 => $DIR/fn-in-macro.rs
+Number of expressions: 1
+- expression 0 operands: lhs = Counter(0), rhs = Counter(1)
+Number of file 0 mappings: 6
+- Code(Counter(0)) at (prev + 25, 1) to (start + 0, 21)
+- Code(Counter(0)) at (prev + 1, 8) to (start + 0, 29)
+- Branch { true: Counter(1), false: Expression(0, Sub) } at (prev + 0, 8) to (start + 0, 35)
+    true  = c1
+    false = (c0 - c1)
+- Code(Counter(1)) at (prev + 0, 36) to (start + 2, 6)
+- Code(Expression(0, Sub)) at (prev + 2, 12) to (start + 2, 6)
+    = (c0 - c1)
+- Code(Counter(0)) at (prev + 3, 1) to (start + 0, 2)
+Highest counter ID seen: c1
+

tests/coverage/branch/fn-in-macro.cov-map

@@ -0,0 +1,62 @@
+   LL|       |#![feature(coverage_attribute)]
+   LL|       |//@ edition: 2024
+   LL|       |//@ compile-flags: -Zcoverage-options=branch
+   LL|       |//@ llvm-cov-flags: --show-branches=count
+   LL|       |
+   LL|       |// Snapshot test demonstrating how branch coverage interacts with code in macros.
+   LL|       |// This test captures current behavior, which is not necessarily "correct".
+   LL|       |
+   LL|       |macro_rules! define_fn {
+   LL|       |    () => {
+   LL|       |        /// Function defined entirely within a macro.
+   LL|      1|        fn fn_in_macro() {
+   LL|      1|            if core::hint::black_box(true) {
+   LL|      1|                say("true");
+   LL|      1|            } else {
+   LL|      0|                say("false");
+   LL|      0|            }
+   LL|      1|        }
+   LL|       |    };
+   LL|       |}
+   LL|       |
+   LL|       |define_fn!();
+   LL|       |
+   LL|       |/// Function not in a macro at all, for comparison.
+   LL|      1|fn fn_not_in_macro() {
+   LL|      1|    if core::hint::black_box(true) {
+  ------------------
+  |  Branch (LL:8): [True: 1, False: 0]
+  ------------------
+   LL|      1|        say("true");
+   LL|      1|    } else {
+   LL|      0|        say("false");
+   LL|      0|    }
+   LL|      1|}
+   LL|       |
+   LL|       |/// Function that is not in a macro, containing a branch that is in a macro.
+   LL|      1|fn branch_in_macro() {
+   LL|       |    macro_rules! macro_with_branch {
+   LL|       |        () => {{
+   LL|       |            if core::hint::black_box(true) {
+   LL|       |                say("true");
+   LL|       |            } else {
+   LL|       |                say("false");
+   LL|       |            }
+   LL|       |        }};
+   LL|       |    }
+   LL|       |
+   LL|      1|    macro_with_branch!();
+   LL|      1|}
+   LL|       |
+   LL|       |#[coverage(off)]
+   LL|       |fn main() {
+   LL|       |    fn_in_macro();
+   LL|       |    fn_not_in_macro();
+   LL|       |    branch_in_macro();
+   LL|       |}
+   LL|       |
+   LL|       |#[coverage(off)]
+   LL|       |fn say(message: &str) {
+   LL|       |    println!("{message}");
+   LL|       |}
+