1 files changed, 54 insertions, 0 deletions
diff --git a/compiler/rustc_mir_transform/src/coverage/unexpand.rs b/compiler/rustc_mir_transform/src/coverage/unexpand.rs
new file mode 100644
index 00000000000..18532b8ee45
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/coverage/unexpand.rs
@@ -0,0 +1,54 @@
+use rustc_span::{ExpnKind, MacroKind, Span, Symbol};
+
+/// Returns an extrapolated span (pre-expansion[^1]) corresponding to a range
+/// within the function's body source. This span is guaranteed to be contained
+/// within, or equal to, the `body_span`. If the extrapolated span is not
+/// contained within the `body_span`, `None` is returned.
+///
+/// [^1]Expansions result from Rust syntax including macros, syntactic sugar,
+/// etc.).
+pub(crate) fn unexpand_into_body_span_with_visible_macro(
+    original_span: Span,
+    body_span: Span,
+) -> Option<(Span, Option<Symbol>)> {
+    let (span, prev) = unexpand_into_body_span_with_prev(original_span, body_span)?;
+
+    let visible_macro = prev
+        .map(|prev| match prev.ctxt().outer_expn_data().kind {
+            ExpnKind::Macro(MacroKind::Bang, name) => Some(name),
+            _ => None,
+        })
+        .flatten();
+
+    Some((span, visible_macro))
+}
+
+/// Walks through the expansion ancestors of `original_span` to find a span that
+/// is contained in `body_span` and has the same [`SyntaxContext`] as `body_span`.
+/// The ancestor that was traversed just before the matching span (if any) is
+/// also returned.
+///
+/// For example, a return value of `Some((ancestor, Some(prev))` means that:
+/// - `ancestor == original_span.find_ancestor_inside_same_ctxt(body_span)`
+/// - `ancestor == prev.parent_callsite()`
+///
+/// [`SyntaxContext`]: rustc_span::SyntaxContext
+fn unexpand_into_body_span_with_prev(
+    original_span: Span,
+    body_span: Span,
+) -> Option<(Span, Option<Span>)> {
+    let mut prev = None;
+    let mut curr = original_span;
+
+    while !body_span.contains(curr) || !curr.eq_ctxt(body_span) {
+        prev = Some(curr);
+        curr = curr.parent_callsite()?;
+    }
+
+    debug_assert_eq!(Some(curr), original_span.find_ancestor_in_same_ctxt(body_span));
+    if let Some(prev) = prev {
+        debug_assert_eq!(Some(curr), prev.parent_callsite());
+    }
+
+    Some((curr, prev))
+}