coverage: Unexpand spans with find_ancestor_inside_same_ctxt

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -23,7 +23,7 @@ use rustc_middle::mir::{
 use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::LocalDefId;
 use rustc_span::source_map::SourceMap;
-use rustc_span::{ExpnKind, Span, Symbol};
+use rustc_span::{DesugaringKind, ExpnKind, Span, Symbol};
 
 /// Inserts `StatementKind::Coverage` statements that either instrument the binary with injected
 /// counters, via intrinsic `llvm.instrprof.increment`, and/or inject metadata used during codegen
@@ -346,20 +346,15 @@ fn get_body_span<'tcx>(
     let mut body_span = hir_body.value.span;
 
     if tcx.is_closure(def_id.to_def_id()) {
-        // If the MIR function is a closure, and if the closure body span
-        // starts from a macro, but it's content is not in that macro, try
-        // to find a non-macro callsite, and instrument the spans there
-        // instead.
-        loop {
-            let expn_data = body_span.ctxt().outer_expn_data();
-            if expn_data.is_root() {
-                break;
-            }
-            if let ExpnKind::Macro { .. } = expn_data.kind {
-                body_span = expn_data.call_site;
-            } else {
-                break;
-            }
+        // If the current function is a closure, and its "body" span was created
+        // by macro expansion or async desugaring, try to walk backwards to the
+        // pre-expansion call site or body.
+        while let expn_data = body_span.ctxt().outer_expn_data()
+            && !expn_data.is_root()
+            && let ExpnKind::Macro { .. } | ExpnKind::Desugaring(DesugaringKind::Async) =
+                expn_data.kind
+        {
+            body_span = expn_data.call_site;
         }
     }
 

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

@@ -204,10 +204,5 @@ fn filtered_terminator_span(terminator: &Terminator<'_>) -> Option<Span> {
 /// etc.).
 #[inline]
 fn unexpand_into_body_span(span: Span, body_span: Span) -> Option<Span> {
-    use rustc_span::source_map::original_sp;
-
-    // FIXME(#118525): Consider switching from `original_sp` to `Span::find_ancestor_inside`,
-    // which is similar but gives slightly different results in some edge cases.
-    let original_span = original_sp(span, body_span).with_ctxt(body_span.ctxt());
-    body_span.contains(original_span).then_some(original_span)
+    span.find_ancestor_inside_same_ctxt(body_span)
 }