Auto merge of #119336 - Zalathar:find-ancestor, r=petrochenkov

coverage: Unexpand spans with `find_ancestor_inside_same_ctxt`

Back in https://github.com/rust-lang/rust/pull/118525#discussion_r1412877621 it was observed that our `unexpand_into_body_span` now looks very similar to `Span::find_ancestor_inside`.

At the time I tried switching over, but doing so resulted in incorrect coverage mappings (or assertion failures), so I left a `FIXME` comment instead.

After some investigation, I identified the two problems with my original approach:
- I should have been using `find_ancestor_inside_same_ctxt` instead, since we want a span that's inside the body and has the same context as the body.
- For async functions, we were actually using the post-expansion body span, which is why we needed to forcibly set the unexpanded span's context to match the body span. For body spans produced by macro-expansion, we already have special-case code to detect this and use the pre-expansion call site as the body span. By making this code also detect async desugaring, I was able to end up with a body span that works properly with `find_ancestor_inside_same_ctxt`, avoiding the need to forcibly change the span context.

2 files changed, 6 insertions, 22 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index c5a3391286a..078612aa59c 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/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::{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,21 +346,10 @@ 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 compiler desugaring, try to walk backwards to
+        // the pre-expansion call site or body.
+        body_span = body_span.source_callsite();
     }
 
     body_span
diff --git a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
index a9c4ea33d0e..474cb205d13 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
+++ b/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)
 }