coverage: Extract hole spans from HIR instead of MIR

This makes it possible to treat more kinds of nested item/code as holes,
instead of being restricted to closures.

3 files changed, 87 insertions, 41 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index d55bde311c1..2efca40d180 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -8,6 +8,10 @@ mod spans;
 mod tests;
 mod unexpand;
 
+use rustc_hir as hir;
+use rustc_hir::intravisit::{walk_expr, Visitor};
+use rustc_middle::hir::map::Map;
+use rustc_middle::hir::nested_filter;
 use rustc_middle::mir::coverage::{
     CodeRegion, CoverageKind, DecisionInfo, FunctionCoverageInfo, Mapping, MappingKind,
 };
@@ -465,6 +469,9 @@ struct ExtractedHirInfo {
     /// Must have the same context and filename as the body span.
     fn_sig_span_extended: Option<Span>,
     body_span: Span,
+    /// "Holes" are regions within the body span that should not be included in
+    /// coverage spans for this function (e.g. closures and nested items).
+    hole_spans: Vec<Span>,
 }
 
 fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHirInfo {
@@ -480,7 +487,7 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
 
     let mut body_span = hir_body.value.span;
 
-    use rustc_hir::{Closure, Expr, ExprKind, Node};
+    use hir::{Closure, Expr, ExprKind, Node};
     // Unexpand a closure's body span back to the context of its declaration.
     // This helps with closure bodies that consist of just a single bang-macro,
     // and also with closure bodies produced by async desugaring.
@@ -507,11 +514,78 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
 
     let function_source_hash = hash_mir_source(tcx, hir_body);
 
-    ExtractedHirInfo { function_source_hash, is_async_fn, fn_sig_span_extended, body_span }
+    let hole_spans = extract_hole_spans_from_hir(tcx, body_span, hir_body);
+
+    ExtractedHirInfo {
+        function_source_hash,
+        is_async_fn,
+        fn_sig_span_extended,
+        body_span,
+        hole_spans,
+    }
 }
 
-fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx rustc_hir::Body<'tcx>) -> u64 {
+fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx hir::Body<'tcx>) -> u64 {
     // FIXME(cjgillot) Stop hashing HIR manually here.
     let owner = hir_body.id().hir_id.owner;
     tcx.hir_owner_nodes(owner).opt_hash_including_bodies.unwrap().to_smaller_hash().as_u64()
 }
+
+fn extract_hole_spans_from_hir<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    body_span: Span, // Usually `hir_body.value.span`, but not always
+    hir_body: &hir::Body<'tcx>,
+) -> Vec<Span> {
+    struct HolesVisitor<'hir, F> {
+        hir: Map<'hir>,
+        visit_hole_span: F,
+    }
+
+    impl<'hir, F: FnMut(Span)> Visitor<'hir> for HolesVisitor<'hir, F> {
+        /// - We need `NestedFilter::INTRA = true` so that `visit_item` will be called.
+        /// - Bodies of nested items don't actually get visited, because of the
+        ///   `visit_item` override.
+        /// - For nested bodies that are not part of an item, we do want to visit any
+        ///   items contained within them.
+        type NestedFilter = nested_filter::All;
+
+        fn nested_visit_map(&mut self) -> Self::Map {
+            self.hir
+        }
+
+        fn visit_item(&mut self, item: &'hir hir::Item<'hir>) {
+            (self.visit_hole_span)(item.span);
+            // Having visited this item, we don't care about its children,
+            // so don't call `walk_item`.
+        }
+
+        // We override `visit_expr` instead of the more specific expression
+        // visitors, so that we have direct access to the expression span.
+        fn visit_expr(&mut self, expr: &'hir hir::Expr<'hir>) {
+            match expr.kind {
+                hir::ExprKind::Closure(_) | hir::ExprKind::ConstBlock(_) => {
+                    (self.visit_hole_span)(expr.span);
+                    // Having visited this expression, we don't care about its
+                    // children, so don't call `walk_expr`.
+                }
+
+                // For other expressions, recursively visit as normal.
+                _ => walk_expr(self, expr),
+            }
+        }
+    }
+
+    let mut hole_spans = vec![];
+    let mut visitor = HolesVisitor {
+        hir: tcx.hir(),
+        visit_hole_span: |hole_span| {
+            // Discard any holes that aren't directly visible within the body span.
+            if body_span.contains(hole_span) && body_span.eq_ctxt(hole_span) {
+                hole_spans.push(hole_span);
+            }
+        },
+    };
+
+    visitor.visit_body(hir_body);
+    hole_spans
+}
diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
index 7612c01c52e..dbc26a2808e 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans.rs
@@ -8,7 +8,7 @@ use rustc_span::Span;
 use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
 use crate::coverage::mappings;
 use crate::coverage::spans::from_mir::{
-    extract_covspans_and_holes_from_mir, ExtractedCovspans, Hole, SpanFromMir,
+    extract_covspans_from_mir, ExtractedCovspans, Hole, SpanFromMir,
 };
 use crate::coverage::ExtractedHirInfo;
 
@@ -20,8 +20,8 @@ pub(super) fn extract_refined_covspans(
     basic_coverage_blocks: &CoverageGraph,
     code_mappings: &mut impl Extend<mappings::CodeMapping>,
 ) {
-    let ExtractedCovspans { mut covspans, mut holes } =
-        extract_covspans_and_holes_from_mir(mir_body, hir_info, basic_coverage_blocks);
+    let ExtractedCovspans { mut covspans } =
+        extract_covspans_from_mir(mir_body, hir_info, basic_coverage_blocks);
 
     // First, perform the passes that need macro information.
     covspans.sort_by(|a, b| basic_coverage_blocks.cmp_in_dominator_order(a.bcb, b.bcb));
@@ -45,6 +45,7 @@ pub(super) fn extract_refined_covspans(
     covspans.dedup_by(|b, a| a.span.source_equal(b.span));
 
     // Sort the holes, and merge overlapping/adjacent holes.
+    let mut holes = hir_info.hole_spans.iter().map(|&span| Hole { span }).collect::<Vec<_>>();
     holes.sort_by(|a, b| compare_spans(a.span, b.span));
     holes.dedup_by(|b, a| a.merge_if_overlapping_or_adjacent(b));
 
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 a0f8f580b1d..32bd25bf4b9 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
@@ -1,8 +1,7 @@
 use rustc_middle::bug;
 use rustc_middle::mir::coverage::CoverageKind;
 use rustc_middle::mir::{
-    self, AggregateKind, FakeReadCause, Rvalue, Statement, StatementKind, Terminator,
-    TerminatorKind,
+    self, FakeReadCause, Statement, StatementKind, Terminator, TerminatorKind,
 };
 use rustc_span::{Span, Symbol};
 
@@ -15,13 +14,12 @@ use crate::coverage::ExtractedHirInfo;
 
 pub(crate) struct ExtractedCovspans {
     pub(crate) covspans: Vec<SpanFromMir>,
-    pub(crate) holes: Vec<Hole>,
 }
 
 /// Traverses the MIR body to produce an initial collection of coverage-relevant
 /// spans, each associated with a node in the coverage graph (BCB) and possibly
 /// other metadata.
-pub(crate) fn extract_covspans_and_holes_from_mir(
+pub(crate) fn extract_covspans_from_mir(
     mir_body: &mir::Body<'_>,
     hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
@@ -29,21 +27,13 @@ pub(crate) fn extract_covspans_and_holes_from_mir(
     let &ExtractedHirInfo { body_span, .. } = hir_info;
 
     let mut covspans = vec![];
-    let mut holes = vec![];
 
     for (bcb, bcb_data) in basic_coverage_blocks.iter_enumerated() {
-        bcb_to_initial_coverage_spans(
-            mir_body,
-            body_span,
-            bcb,
-            bcb_data,
-            &mut covspans,
-            &mut holes,
-        );
+        bcb_to_initial_coverage_spans(mir_body, body_span, bcb, bcb_data, &mut covspans);
     }
 
     // Only add the signature span if we found at least one span in the body.
-    if !covspans.is_empty() || !holes.is_empty() {
+    if !covspans.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.
@@ -51,7 +41,7 @@ pub(crate) fn extract_covspans_and_holes_from_mir(
         covspans.push(SpanFromMir::for_fn_sig(fn_sig_span));
     }
 
-    ExtractedCovspans { covspans, holes }
+    ExtractedCovspans { covspans }
 }
 
 // Generate a set of coverage spans from the filtered set of `Statement`s and `Terminator`s of
@@ -65,7 +55,6 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
     bcb: BasicCoverageBlock,
     bcb_data: &'a BasicCoverageBlockData,
     initial_covspans: &mut Vec<SpanFromMir>,
-    holes: &mut Vec<Hole>,
 ) {
     for &bb in &bcb_data.basic_blocks {
         let data = &mir_body[bb];
@@ -81,13 +70,7 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
             let expn_span = filtered_statement_span(statement)?;
             let (span, visible_macro) = unexpand(expn_span)?;
 
-            // A statement that looks like the assignment of a closure expression
-            // is treated as a "hole" span, to be carved out of other spans.
-            if is_closure_like(statement) {
-                holes.push(Hole { span });
-            } else {
-                initial_covspans.push(SpanFromMir::new(span, visible_macro, bcb));
-            }
+            initial_covspans.push(SpanFromMir::new(span, visible_macro, bcb));
             Some(())
         };
         for statement in data.statements.iter() {
@@ -105,18 +88,6 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
     }
 }
 
-fn is_closure_like(statement: &Statement<'_>) -> bool {
-    match statement.kind {
-        StatementKind::Assign(box (_, Rvalue::Aggregate(box ref agg_kind, _))) => match agg_kind {
-            AggregateKind::Closure(_, _)
-            | AggregateKind::Coroutine(_, _)
-            | AggregateKind::CoroutineClosure(..) => true,
-            _ => false,
-        },
-        _ => false,
-    }
-}
-
 /// If the MIR `Statement` has a span contributive to computing coverage spans,
 /// return it; otherwise return `None`.
 fn filtered_statement_span(statement: &Statement<'_>) -> Option<Span> {