coverage: Extract hole spans from HIR instead of MIR

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
+}

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));
 

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,43 +14,34 @@ 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,
 ) -> ExtractedCovspans {
     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.
         let fn_sig_span = hir_info.fn_sig_span_extended.unwrap_or_else(|| body_span.shrink_to_lo());
         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> {

tests/coverage-run-rustdoc/doctest.coverage

@@ -34,10 +34,10 @@ $DIR/doctest.rs:
    LL|       |//!
    LL|       |//! doctest returning a result:
    LL|      1|//! ```
-   LL|      1|//! #[derive(Debug, PartialEq)]
-   LL|      1|//! struct SomeError {
-   LL|      1|//!     msg: String,
-   LL|      1|//! }
+   LL|       |//! #[derive(Debug, PartialEq)]
+   LL|       |//! struct SomeError {
+   LL|       |//!     msg: String,
+   LL|       |//! }
    LL|      1|//! let mut res = Err(SomeError { msg: String::from("a message") });
    LL|      1|//! if res.is_ok() {
    LL|      0|//!     res?;

tests/coverage/async.cov-map

@@ -167,15 +167,16 @@ Number of file 0 mappings: 14
     = ((c6 + c7) + c8)
 
 Function name: async::j
-Raw bytes (53): 0x[01, 01, 02, 07, 0d, 05, 09, 09, 01, 35, 01, 13, 0c, 05, 14, 09, 00, 0a, 01, 00, 0e, 00, 1b, 05, 00, 1f, 00, 27, 09, 01, 09, 00, 0a, 11, 00, 0e, 00, 1a, 09, 00, 1e, 00, 20, 0d, 01, 0e, 00, 10, 03, 02, 01, 00, 02]
+Raw bytes (58): 0x[01, 01, 02, 07, 0d, 05, 09, 0a, 01, 35, 01, 00, 0d, 01, 0b, 0b, 00, 0c, 05, 01, 09, 00, 0a, 01, 00, 0e, 00, 1b, 05, 00, 1f, 00, 27, 09, 01, 09, 00, 0a, 11, 00, 0e, 00, 1a, 09, 00, 1e, 00, 20, 0d, 01, 0e, 00, 10, 03, 02, 01, 00, 02]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 2
 - expression 0 operands: lhs = Expression(1, Add), rhs = Counter(3)
 - expression 1 operands: lhs = Counter(1), rhs = Counter(2)
-Number of file 0 mappings: 9
-- Code(Counter(0)) at (prev + 53, 1) to (start + 19, 12)
-- Code(Counter(1)) at (prev + 20, 9) to (start + 0, 10)
+Number of file 0 mappings: 10
+- Code(Counter(0)) at (prev + 53, 1) to (start + 0, 13)
+- Code(Counter(0)) at (prev + 11, 11) to (start + 0, 12)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 10)
 - Code(Counter(0)) at (prev + 0, 14) to (start + 0, 27)
 - Code(Counter(1)) at (prev + 0, 31) to (start + 0, 39)
 - Code(Counter(2)) at (prev + 1, 9) to (start + 0, 10)
@@ -186,48 +187,48 @@ Number of file 0 mappings: 9
     = ((c1 + c2) + c3)
 
 Function name: async::j::c
-Raw bytes (26): 0x[01, 01, 01, 01, 05, 04, 01, 37, 05, 01, 12, 05, 02, 0d, 00, 0e, 02, 0a, 0d, 00, 0e, 01, 02, 05, 00, 06]
+Raw bytes (26): 0x[01, 01, 01, 01, 05, 04, 01, 37, 05, 01, 12, 05, 02, 0d, 00, 0e, 02, 02, 0d, 00, 0e, 01, 02, 05, 00, 06]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 1
 - expression 0 operands: lhs = Counter(0), rhs = Counter(1)
 Number of file 0 mappings: 4
 - Code(Counter(0)) at (prev + 55, 5) to (start + 1, 18)
 - Code(Counter(1)) at (prev + 2, 13) to (start + 0, 14)
-- Code(Expression(0, Sub)) at (prev + 10, 13) to (start + 0, 14)
+- Code(Expression(0, Sub)) at (prev + 2, 13) to (start + 0, 14)
     = (c0 - c1)
 - Code(Counter(0)) at (prev + 2, 5) to (start + 0, 6)
 
 Function name: async::j::d
-Raw bytes (9): 0x[01, 01, 00, 01, 01, 46, 05, 00, 17]
+Raw bytes (9): 0x[01, 01, 00, 01, 01, 3e, 05, 00, 17]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 0
 Number of file 0 mappings: 1
-- Code(Counter(0)) at (prev + 70, 5) to (start + 0, 23)
+- Code(Counter(0)) at (prev + 62, 5) to (start + 0, 23)
 
 Function name: async::j::f
-Raw bytes (9): 0x[01, 01, 00, 01, 01, 47, 05, 00, 17]
+Raw bytes (9): 0x[01, 01, 00, 01, 01, 3f, 05, 00, 17]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 0
 Number of file 0 mappings: 1
-- Code(Counter(0)) at (prev + 71, 5) to (start + 0, 23)
+- Code(Counter(0)) at (prev + 63, 5) to (start + 0, 23)
 
 Function name: async::k (unused)
-Raw bytes (29): 0x[01, 01, 00, 05, 00, 4f, 01, 01, 0c, 00, 02, 0e, 00, 10, 00, 01, 0e, 00, 10, 00, 01, 0e, 00, 10, 00, 02, 01, 00, 02]
+Raw bytes (29): 0x[01, 01, 00, 05, 00, 47, 01, 01, 0c, 00, 02, 0e, 00, 10, 00, 01, 0e, 00, 10, 00, 01, 0e, 00, 10, 00, 02, 01, 00, 02]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 0
 Number of file 0 mappings: 5
-- Code(Zero) at (prev + 79, 1) to (start + 1, 12)
+- Code(Zero) at (prev + 71, 1) to (start + 1, 12)
 - Code(Zero) at (prev + 2, 14) to (start + 0, 16)
 - Code(Zero) at (prev + 1, 14) to (start + 0, 16)
 - Code(Zero) at (prev + 1, 14) to (start + 0, 16)
 - Code(Zero) at (prev + 2, 1) to (start + 0, 2)
 
 Function name: async::l
-Raw bytes (37): 0x[01, 01, 04, 01, 07, 05, 09, 0f, 02, 09, 05, 05, 01, 57, 01, 01, 0c, 02, 02, 0e, 00, 10, 05, 01, 0e, 00, 10, 09, 01, 0e, 00, 10, 0b, 02, 01, 00, 02]
+Raw bytes (37): 0x[01, 01, 04, 01, 07, 05, 09, 0f, 02, 09, 05, 05, 01, 4f, 01, 01, 0c, 02, 02, 0e, 00, 10, 05, 01, 0e, 00, 10, 09, 01, 0e, 00, 10, 0b, 02, 01, 00, 02]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 4
@@ -236,7 +237,7 @@ Number of expressions: 4
 - expression 2 operands: lhs = Expression(3, Add), rhs = Expression(0, Sub)
 - expression 3 operands: lhs = Counter(2), rhs = Counter(1)
 Number of file 0 mappings: 5
-- Code(Counter(0)) at (prev + 87, 1) to (start + 1, 12)
+- Code(Counter(0)) at (prev + 79, 1) to (start + 1, 12)
 - Code(Expression(0, Sub)) at (prev + 2, 14) to (start + 0, 16)
     = (c0 - (c1 + c2))
 - Code(Counter(1)) at (prev + 1, 14) to (start + 0, 16)
@@ -245,26 +246,26 @@ Number of file 0 mappings: 5
     = ((c2 + c1) + (c0 - (c1 + c2)))
 
 Function name: async::m
-Raw bytes (9): 0x[01, 01, 00, 01, 01, 5f, 01, 00, 19]
+Raw bytes (9): 0x[01, 01, 00, 01, 01, 57, 01, 00, 19]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 0
 Number of file 0 mappings: 1
-- Code(Counter(0)) at (prev + 95, 1) to (start + 0, 25)
+- Code(Counter(0)) at (prev + 87, 1) to (start + 0, 25)
 
 Function name: async::m::{closure#0} (unused)
-Raw bytes (9): 0x[01, 01, 00, 01, 00, 5f, 19, 00, 22]
+Raw bytes (9): 0x[01, 01, 00, 01, 00, 57, 19, 00, 22]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 0
 Number of file 0 mappings: 1
-- Code(Zero) at (prev + 95, 25) to (start + 0, 34)
+- Code(Zero) at (prev + 87, 25) to (start + 0, 34)
 
 Function name: async::main
-Raw bytes (9): 0x[01, 01, 00, 01, 01, 61, 01, 08, 02]
+Raw bytes (9): 0x[01, 01, 00, 01, 01, 59, 01, 08, 02]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 0
 Number of file 0 mappings: 1
-- Code(Counter(0)) at (prev + 97, 1) to (start + 8, 2)
+- Code(Counter(0)) at (prev + 89, 1) to (start + 8, 2)