coverage: Make fn_sig_span optional, and note its quirks

Zalathar · Zalathar · commit 23fb8946b037 · 2024-02-02T13:05:41.000+11:00
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -394,7 +394,9 @@ fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
 struct ExtractedHirInfo {
     function_source_hash: u64,
     is_async_fn: bool,
-    fn_sig_span: Span,
+    /// The span of the function's signature, extended to the start of `body_span`.
+    /// Must have the same context and filename as the body span.
+    fn_sig_span_extended: Option<Span>,
     body_span: Span,
 }
 
@@ -407,7 +409,8 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
         hir::map::associated_body(hir_node).expect("HIR node is a function with body");
     let hir_body = tcx.hir().body(fn_body_id);
 
-    let is_async_fn = hir_node.fn_sig().is_some_and(|fn_sig| fn_sig.header.is_async());
+    let maybe_fn_sig = hir_node.fn_sig();
+    let is_async_fn = maybe_fn_sig.is_some_and(|fn_sig| fn_sig.header.is_async());
 
     let mut body_span = hir_body.value.span;
 
@@ -423,8 +426,8 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
 
     // The actual signature span is only used if it has the same context and
     // filename as the body, and precedes the body.
-    let maybe_fn_sig_span = hir_node.fn_sig().map(|fn_sig| fn_sig.span);
-    let fn_sig_span = maybe_fn_sig_span
+    let fn_sig_span_extended = maybe_fn_sig
+        .map(|fn_sig| fn_sig.span)
         .filter(|&fn_sig_span| {
             let source_map = tcx.sess.source_map();
             let file_idx = |span: Span| source_map.lookup_source_file_idx(span.lo());
@@ -434,13 +437,11 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
                 && file_idx(fn_sig_span) == file_idx(body_span)
         })
         // If so, extend it to the start of the body span.
-        .map(|fn_sig_span| fn_sig_span.with_hi(body_span.lo()))
-        // Otherwise, create a dummy signature span at the start of the body.
-        .unwrap_or_else(|| body_span.shrink_to_lo());
+        .map(|fn_sig_span| fn_sig_span.with_hi(body_span.lo()));
 
     let function_source_hash = hash_mir_source(tcx, hir_body);
 
-    ExtractedHirInfo { function_source_hash, is_async_fn, fn_sig_span, body_span }
+    ExtractedHirInfo { function_source_hash, is_async_fn, fn_sig_span_extended, body_span }
 }
 
 fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx rustc_hir::Body<'tcx>) -> u64 {
diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
@@ -51,8 +51,9 @@ pub(super) fn generate_coverage_spans(
         // with the user code wrapped in a closure. Any spans in the desugared
         // outer function will be unhelpful, so just keep the signature span
         // and ignore all of the spans in the MIR body.
-        let span = hir_info.fn_sig_span;
-        mappings.push(BcbMapping { kind: BcbMappingKind::Code(START_BCB), span });
+        if let Some(span) = hir_info.fn_sig_span_extended {
+            mappings.push(BcbMapping { kind: BcbMappingKind::Code(START_BCB), span });
+        }
     } else {
         let sorted_spans = from_mir::mir_to_initial_sorted_coverage_spans(
             mir_body,
diff --git a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
@@ -23,7 +23,7 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
     hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
 ) -> Vec<CoverageSpan> {
-    let &ExtractedHirInfo { fn_sig_span, body_span, .. } = hir_info;
+    let &ExtractedHirInfo { body_span, .. } = hir_info;
 
     let mut initial_spans = vec![];
 
@@ -33,6 +33,10 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
 
     // Only add the signature span if we found at least one span in the body.
     if !initial_spans.is_empty() {
+        // If there is no usable signature span, add a fake one (before refinement)
+        // to avoid an ugly gap between the body start and the first real span.
+        // FIXME: Find a more principled way to solve this problem.
+        let fn_sig_span = hir_info.fn_sig_span_extended.unwrap_or_else(|| body_span.shrink_to_lo());
         initial_spans.push(SpanFromMir::for_fn_sig(fn_sig_span));
     }
 
