Auto merge of #79109 - richkadel:llvm-coverage-counters-2.0.5, r=tmandry

Coverage tests for remaining TerminatorKinds and async, improve Assert

Tested and validate results for panic unwind, panic abort, assert!()
macro, TerminatorKind::Assert (for example, numeric overflow), and
async/await.

Implemented a previous documented idea to change Assert handling to be
the same as FalseUnwind and Goto, so it doesn't get its own
BasicCoverageBlock anymore. This changed a couple of coverage regions,
but I validated those changes are not any worse than the prior results,
and probably help assure some consistency (even if some people might
disagree with how the code region is consistently computed).

Fixed issue with async/await. AggregateKind::Generator needs to be
handled like AggregateKind::Closure; coverage span for the outer async
function should not "cover" the async body, which is actually executed
in a separate "closure" MIR.

compiler/rustc_ast_lowering/src/expr.rs

@@ -347,7 +347,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         // `_ => else_block` where `else_block` is `{}` if there's `None`:
         let else_pat = self.pat_wild(span);
         let (else_expr, contains_else_clause) = match else_opt {
-            None => (self.expr_block_empty(span), false),
+            None => (self.expr_block_empty(span.shrink_to_hi()), false),
             Some(els) => (self.lower_expr(els), true),
         };
         let else_arm = self.arm(else_pat, else_expr);

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -127,9 +127,6 @@ fn set_probestack(cx: &CodegenCx<'ll, '_>, llfn: &'ll Value) {
         return;
     }
 
-    // FIXME(richkadel): Make sure probestack plays nice with `-Z instrument-coverage`
-    // or disable it if not, similar to above early exits.
-
     // Flag our internal `__rust_probestack` function as the stack probe symbol.
     // This is defined in the `compiler-builtins` crate for each architecture.
     llvm::AddFunctionAttrStringValue(

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -3,11 +3,14 @@ use crate::coverageinfo;
 use crate::llvm;
 
 use llvm::coverageinfo::CounterMappingRegion;
-use rustc_codegen_ssa::coverageinfo::map::{Counter, CounterExpression};
+use rustc_codegen_ssa::coverageinfo::map::{Counter, CounterExpression, FunctionCoverage};
 use rustc_codegen_ssa::traits::ConstMethods;
-use rustc_data_structures::fx::FxIndexSet;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexSet};
+use rustc_hir::def_id::{DefId, DefIdSet, LOCAL_CRATE};
 use rustc_llvm::RustString;
 use rustc_middle::mir::coverage::CodeRegion;
+use rustc_middle::ty::{Instance, TyCtxt};
+use rustc_span::Symbol;
 
 use std::ffi::CString;
 
@@ -26,14 +29,17 @@ use tracing::debug;
 /// undocumented details in Clang's implementation (that may or may not be important) were also
 /// replicated for Rust's Coverage Map.
 pub fn finalize<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>) {
+    let tcx = cx.tcx;
     // Ensure LLVM supports Coverage Map Version 4 (encoded as a zero-based value: 3).
     // If not, the LLVM Version must be less than 11.
     let version = coverageinfo::mapping_version();
     if version != 3 {
-        cx.tcx.sess.fatal("rustc option `-Z instrument-coverage` requires LLVM 11 or higher.");
+        tcx.sess.fatal("rustc option `-Z instrument-coverage` requires LLVM 11 or higher.");
     }
 
-    let function_coverage_map = match cx.coverage_context() {
+    debug!("Generating coverage map for CodegenUnit: `{}`", cx.codegen_unit.name());
+
+    let mut function_coverage_map = match cx.coverage_context() {
         Some(ctx) => ctx.take_function_coverage_map(),
         None => return,
     };
@@ -42,14 +48,15 @@ pub fn finalize<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>) {
         return;
     }
 
+    add_unreachable_coverage(tcx, &mut function_coverage_map);
+
     let mut mapgen = CoverageMapGenerator::new();
 
     // Encode coverage mappings and generate function records
     let mut function_data = Vec::new();
     for (instance, function_coverage) in function_coverage_map {
-        debug!("Generate coverage map for: {:?}", instance);
-
-        let mangled_function_name = cx.tcx.symbol_name(instance).to_string();
+        debug!("Generate function coverage for {}, {:?}", cx.codegen_unit.name(), instance);
+        let mangled_function_name = tcx.symbol_name(instance).to_string();
         let function_source_hash = function_coverage.source_hash();
         let (expressions, counter_regions) =
             function_coverage.get_expressions_and_counter_regions();
@@ -228,3 +235,156 @@ fn save_function_record(
     let is_used = true;
     coverageinfo::save_func_record_to_mod(cx, func_name_hash, func_record_val, is_used);
 }
+
+/// When finalizing the coverage map, `FunctionCoverage` only has the `CodeRegion`s and counters for
+/// the functions that went through codegen; such as public functions and "used" functions
+/// (functions referenced by other "used" or public items). Any other functions considered unused,
+/// or "Unreachable" were still parsed and processed through the MIR stage.
+///
+/// We can find the unreachable functions by the set different of all MIR `DefId`s (`tcx` query
+/// `mir_keys`) minus the codegenned `DefId`s (`tcx` query `collect_and_partition_mono_items`).
+///
+/// *HOWEVER* the codegenned `DefId`s are partitioned across multiple `CodegenUnit`s (CGUs), and
+/// this function is processing a `function_coverage_map` for the functions (`Instance`/`DefId`)
+/// allocated to only one of those CGUs. We must NOT inject any "Unreachable" functions's
+/// `CodeRegion`s more than once, so we have to pick which CGU's `function_coverage_map` to add
+/// each "Unreachable" function to.
+///
+/// Some constraints:
+///
+/// 1. The file name of an "Unreachable" function must match the file name of the existing
+///    codegenned (covered) function to which the unreachable code regions will be added.
+/// 2. The function to which the unreachable code regions will be added must not be a genaric
+///    function (must not have type parameters) because the coverage tools will get confused
+///    if the codegenned function has more than one instantiation and additional `CodeRegion`s
+///    attached to only one of those instantiations.
+fn add_unreachable_coverage<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    function_coverage_map: &mut FxHashMap<Instance<'tcx>, FunctionCoverage<'tcx>>,
+) {
+    // FIXME(#79622): Can this solution be simplified and/or improved? Are there other sources
+    // of compiler state data that might help (or better sources that could be exposed, but
+    // aren't yet)?
+
+    // Note: If the crate *only* defines generic functions, there are no codegenerated non-generic
+    // functions to add any unreachable code to. In this case, the unreachable code regions will
+    // have no coverage, instead of having coverage with zero executions.
+    //
+    // This is probably still an improvement over Clang, which does not generate any coverage
+    // for uninstantiated template functions.
+
+    let has_non_generic_def_ids =
+        function_coverage_map.keys().any(|instance| instance.def.attrs(tcx).len() == 0);
+
+    if !has_non_generic_def_ids {
+        // There are no non-generic functions to add unreachable `CodeRegion`s to
+        return;
+    }
+
+    let all_def_ids: DefIdSet =
+        tcx.mir_keys(LOCAL_CRATE).iter().map(|local_def_id| local_def_id.to_def_id()).collect();
+
+    let (codegenned_def_ids, _) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
+
+    let mut unreachable_def_ids_by_file: FxHashMap<Symbol, Vec<DefId>> = FxHashMap::default();
+    for &non_codegenned_def_id in all_def_ids.difference(codegenned_def_ids) {
+        // Make sure the non-codegenned (unreachable) function has a file_name
+        if let Some(non_codegenned_file_name) = tcx.covered_file_name(non_codegenned_def_id) {
+            let def_ids = unreachable_def_ids_by_file
+                .entry(*non_codegenned_file_name)
+                .or_insert_with(|| Vec::new());
+            def_ids.push(non_codegenned_def_id);
+        }
+    }
+
+    if unreachable_def_ids_by_file.is_empty() {
+        // There are no unreachable functions with file names to add (in any CGU)
+        return;
+    }
+
+    // Since there may be multiple `CodegenUnit`s, some codegenned_def_ids may be codegenned in a
+    // different CGU, and will be added to the function_coverage_map for each CGU. Determine which
+    // function_coverage_map has the responsibility for publishing unreachable coverage
+    // based on file name:
+    //
+    // For each covered file name, sort ONLY the non-generic codegenned_def_ids, and if
+    // covered_def_ids.contains(the first def_id) for a given file_name, add the unreachable code
+    // region in this function_coverage_map. Otherwise, ignore it and assume another CGU's
+    // function_coverage_map will be adding it (because it will be first for one, and only one,
+    // of them).
+    let mut sorted_codegenned_def_ids: Vec<DefId> =
+        codegenned_def_ids.iter().map(|def_id| *def_id).collect();
+    sorted_codegenned_def_ids.sort_unstable();
+
+    let mut first_covered_def_id_by_file: FxHashMap<Symbol, DefId> = FxHashMap::default();
+    for &def_id in sorted_codegenned_def_ids.iter() {
+        // Only consider non-generic functions, to potentially add unreachable code regions
+        if tcx.generics_of(def_id).count() == 0 {
+            if let Some(covered_file_name) = tcx.covered_file_name(def_id) {
+                // Only add files known to have unreachable functions
+                if unreachable_def_ids_by_file.contains_key(covered_file_name) {
+                    first_covered_def_id_by_file.entry(*covered_file_name).or_insert(def_id);
+                }
+            }
+        }
+    }
+
+    // Get the set of def_ids with coverage regions, known by *this* CoverageContext.
+    let cgu_covered_def_ids: DefIdSet =
+        function_coverage_map.keys().map(|instance| instance.def.def_id()).collect();
+
+    let mut cgu_covered_files: FxHashSet<Symbol> = first_covered_def_id_by_file
+        .iter()
+        .filter_map(
+            |(&file_name, def_id)| {
+                if cgu_covered_def_ids.contains(def_id) { Some(file_name) } else { None }
+            },
+        )
+        .collect();
+
+    // Find the first covered, non-generic function (instance) for each cgu_covered_file. Take the
+    // unreachable code regions for that file, and add them to the function.
+    //
+    // There are three `for` loops here, but (a) the lists have already been reduced to the minimum
+    // required values, the lists are further reduced (by `remove()` calls) when elements are no
+    // longer needed, and there are several opportunities to branch out of loops early.
+    for (instance, function_coverage) in function_coverage_map.iter_mut() {
+        if instance.def.attrs(tcx).len() > 0 {
+            continue;
+        }
+        // The covered function is not generic...
+        let covered_def_id = instance.def.def_id();
+        if let Some(covered_file_name) = tcx.covered_file_name(covered_def_id) {
+            if !cgu_covered_files.remove(&covered_file_name) {
+                continue;
+            }
+            // The covered function's file is one of the files with unreachable code regions, so
+            // all of the unreachable code regions for this file will be added to this function.
+            for def_id in
+                unreachable_def_ids_by_file.remove(&covered_file_name).into_iter().flatten()
+            {
+                // Note, this loop adds an unreachable code regions for each MIR-derived region.
+                // Alternatively, we could add a single code region for the maximum span of all
+                // code regions here.
+                //
+                // Observed downsides of this approach are:
+                //
+                // 1. The coverage results will appear inconsistent compared with the same (or
+                //    similar) code in a function that is reached.
+                // 2. If the function is unreachable from one crate but reachable when compiling
+                //    another referencing crate (such as a cross-crate reference to a
+                //    generic function or inlined function), actual coverage regions overlaid
+                //    on a single larger code span of `Zero` coverage can appear confusing or
+                //    wrong. Chaning the unreachable coverage from a `code_region` to a
+                //    `gap_region` can help, but still can look odd with `0` line counts for
+                //    lines between executed (> 0) lines (such as for blank lines or comments).
+                for &region in tcx.covered_code_regions(def_id) {
+                    function_coverage.add_unreachable_region(region.clone());
+                }
+            }
+            if cgu_covered_files.is_empty() {
+                break;
+            }
+        }
+    }
+}

compiler/rustc_codegen_ssa/src/lib.rs

@@ -2,6 +2,7 @@
 #![feature(bool_to_option)]
 #![feature(option_expect_none)]
 #![feature(box_patterns)]
+#![feature(drain_filter)]
 #![feature(try_blocks)]
 #![feature(in_band_lifetimes)]
 #![feature(nll)]

compiler/rustc_middle/src/arena.rs

@@ -32,6 +32,7 @@ macro_rules! arena_types {
             [decode] borrowck_result:
                 rustc_middle::mir::BorrowCheckResult<$tcx>,
             [decode] unsafety_check_result: rustc_middle::mir::UnsafetyCheckResult,
+            [decode] code_region: rustc_middle::mir::coverage::CodeRegion,
             [] const_allocs: rustc_middle::mir::interpret::Allocation,
             // Required for the incremental on-disk cache
             [few] mir_keys: rustc_hir::def_id::DefIdSet,

compiler/rustc_middle/src/hir/map/mod.rs

@@ -47,7 +47,7 @@ fn fn_decl<'hir>(node: Node<'hir>) -> Option<&'hir FnDecl<'hir>> {
     }
 }
 
-fn fn_sig<'hir>(node: Node<'hir>) -> Option<&'hir FnSig<'hir>> {
+pub fn fn_sig<'hir>(node: Node<'hir>) -> Option<&'hir FnSig<'hir>> {
     match &node {
         Node::Item(Item { kind: ItemKind::Fn(sig, _, _), .. })
         | Node::TraitItem(TraitItem { kind: TraitItemKind::Fn(sig, _), .. })

compiler/rustc_middle/src/query/mod.rs

@@ -346,6 +346,21 @@ rustc_queries! {
             cache_on_disk_if { key.is_local() }
         }
 
+        /// Returns the name of the file that contains the function body, if instrumented for coverage.
+        query covered_file_name(key: DefId) -> Option<Symbol> {
+            desc { |tcx| "retrieving the covered file name, if instrumented, for `{}`", tcx.def_path_str(key) }
+            storage(ArenaCacheSelector<'tcx>)
+            cache_on_disk_if { key.is_local() }
+        }
+
+        /// Returns the `CodeRegions` for a function that has instrumented coverage, in case the
+        /// function was optimized out before codegen, and before being added to the Coverage Map.
+        query covered_code_regions(key: DefId) -> Vec<&'tcx mir::coverage::CodeRegion> {
+            desc { |tcx| "retrieving the covered `CodeRegion`s, if instrumented, for `{}`", tcx.def_path_str(key) }
+            storage(ArenaCacheSelector<'tcx>)
+            cache_on_disk_if { key.is_local() }
+        }
+
         /// The `DefId` is the `DefId` of the containing MIR body. Promoteds do not have their own
         /// `DefId`. This function returns all promoteds in the specified body. The body references
         /// promoteds by the `DefId` and the `mir::Promoted` index. This is necessary, because

compiler/rustc_middle/src/ty/codec.rs

@@ -162,7 +162,8 @@ encodable_via_deref! {
     ty::Region<'tcx>,
     &'tcx mir::Body<'tcx>,
     &'tcx mir::UnsafetyCheckResult,
-    &'tcx mir::BorrowCheckResult<'tcx>
+    &'tcx mir::BorrowCheckResult<'tcx>,
+    &'tcx mir::coverage::CodeRegion
 }
 
 pub trait TyDecoder<'tcx>: Decoder {
@@ -376,7 +377,8 @@ impl_decodable_via_ref! {
     &'tcx Allocation,
     &'tcx mir::Body<'tcx>,
     &'tcx mir::UnsafetyCheckResult,
-    &'tcx mir::BorrowCheckResult<'tcx>
+    &'tcx mir::BorrowCheckResult<'tcx>,
+    &'tcx mir::coverage::CodeRegion
 }
 
 #[macro_export]

compiler/rustc_mir/src/transform/coverage/graph.rs

@@ -118,18 +118,8 @@ impl CoverageGraph {
 
             match term.kind {
                 TerminatorKind::Return { .. }
-                // FIXME(richkadel): Add test(s) for `Abort` coverage.
                 | TerminatorKind::Abort
-                // FIXME(richkadel): Add test(s) for `Assert` coverage.
-                // Should `Assert` be handled like `FalseUnwind` instead? Since we filter out unwind
-                // branches when creating the BCB CFG, aren't `Assert`s (without unwinds) just like
-                // `FalseUnwinds` (which are kind of like `Goto`s)?
-                | TerminatorKind::Assert { .. }
-                // FIXME(richkadel): Add test(s) for `Yield` coverage, and confirm coverage is
-                // sensible for code using the `yield` keyword.
                 | TerminatorKind::Yield { .. }
-                // FIXME(richkadel): Also add coverage tests using async/await, and threading.
-
                 | TerminatorKind::SwitchInt { .. } => {
                     // The `bb` has more than one _outgoing_ edge, or exits the function. Save the
                     // current sequence of `basic_blocks` gathered to this point, as a new
@@ -147,13 +137,24 @@ impl CoverageGraph {
                     // `Terminator`s `successors()` list) checking the number of successors won't
                     // work.
                 }
+
+                // The following `TerminatorKind`s are either not expected outside an unwind branch,
+                // or they should not (under normal circumstances) branch. Coverage graphs are
+                // simplified by assuring coverage results are accurate for program executions that
+                // don't panic.
+                //
+                // Programs that panic and unwind may record slightly inaccurate coverage results
+                // for a coverage region containing the `Terminator` that began the panic. This
+                // is as intended. (See Issue #78544 for a possible future option to support
+                // coverage in test programs that panic.)
                 TerminatorKind::Goto { .. }
                 | TerminatorKind::Resume
                 | TerminatorKind::Unreachable
                 | TerminatorKind::Drop { .. }
                 | TerminatorKind::DropAndReplace { .. }
                 | TerminatorKind::Call { .. }
                 | TerminatorKind::GeneratorDrop
+                | TerminatorKind::Assert { .. }
                 | TerminatorKind::FalseEdge { .. }
                 | TerminatorKind::FalseUnwind { .. }
                 | TerminatorKind::InlineAsm { .. } => {}
@@ -278,6 +279,7 @@ rustc_index::newtype_index! {
     /// A node in the [control-flow graph][CFG] of CoverageGraph.
     pub(super) struct BasicCoverageBlock {
         DEBUG_FORMAT = "bcb{}",
+        const START_BCB = 0,
     }
 }