coverage: Store fn_sig_span and body_span in the expansion tree

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/expansion.rs

@@ -71,6 +71,13 @@ 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.
@@ -95,6 +102,9 @@ impl ExpnNode {
             call_site,
             call_site_expn_id,
 
+            fn_sig_span: None,
+            body_span: None,
+
             spans: vec![],
             child_expn_ids: FxIndexSet::default(),
 
@@ -142,6 +152,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 {

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)