[code coverage] Fix missing dead code in modules that are never called

The issue here is that the logic used to determine which CGU to put the
dead function stubs in doesn't handle cases where a module is never
assigned to a CGU.

The partitioning logic also caused issues in #85461 where inline
functions were duplicated into multiple CGUs resulting in duplicate
symbols.

This commit fixes the issue by removing the complex logic used to assign
dead code stubs to CGUs and replaces it with a much simplier model: we
pick one CGU to hold all the dead code stubs. We pick a CGU which has
exported items which increases the likelihood the linker won't throw
away our dead functions and we pick the smallest to minimize the impact
on compilation times for crates with very large CGUs.

Fixes #86177
Fixes #85718
Fixes #79622

Author: wesleywiser

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -5,12 +5,13 @@ use crate::llvm;
 use llvm::coverageinfo::CounterMappingRegion;
 use rustc_codegen_ssa::coverageinfo::map::{Counter, CounterExpression};
 use rustc_codegen_ssa::traits::{ConstMethods, CoverageInfoMethods};
-use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexSet};
-use rustc_hir::def_id::{DefId, DefIdSet};
+use rustc_data_structures::fx::FxIndexSet;
+use rustc_hir::def::DefKind;
+use rustc_hir::def_id::DefIdSet;
 use rustc_llvm::RustString;
+use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::coverage::CodeRegion;
 use rustc_middle::ty::TyCtxt;
-use rustc_span::Symbol;
 
 use std::ffi::CString;
 
@@ -46,7 +47,7 @@ pub fn finalize<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>) {
     // functions exist. Generate synthetic functions with a (required) single counter, and add the
     // MIR `Coverage` code regions to the `function_coverage_map`, before calling
     // `ctx.take_function_coverage_map()`.
-    if !tcx.sess.instrument_coverage_except_unused_functions() {
+    if cx.codegen_unit.is_code_coverage_dead_code_cgu() {
         add_unused_functions(cx);
     }
 
@@ -271,17 +272,12 @@ fn save_function_record(
 /// `DefId`s (`tcx` query `mir_keys`) minus the codegenned `DefId`s (`tcx` query
 /// `codegened_and_inlined_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 unused functions's `CodeRegion`s
-/// more than once, so we have to pick a CGUs `function_coverage_map` into which the unused
-/// function will be inserted.
+/// These unused functions are then codegen'd in one of the CGUs which is marked as the
+/// "code coverage dead code cgu" during the partitioning process.
 fn add_unused_functions<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>) {
-    let tcx = cx.tcx;
+    assert!(cx.codegen_unit.is_code_coverage_dead_code_cgu());
 
-    // 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)?
+    let tcx = cx.tcx;
 
     let ignore_unused_generics = tcx.sess.instrument_coverage_except_unused_generics();
 
@@ -299,79 +295,24 @@ fn add_unused_functions<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>) {
 
     let codegenned_def_ids = tcx.codegened_and_inlined_items(());
 
-    let mut unused_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 (unused) function has at least one MIR
-        // `Coverage` statement with a code region, and return its file name.
-        if let Some(non_codegenned_file_name) = tcx.covered_file_name(non_codegenned_def_id) {
-            let def_ids =
-                unused_def_ids_by_file.entry(*non_codegenned_file_name).or_insert_with(Vec::new);
-            def_ids.push(non_codegenned_def_id);
-        }
-    }
-
-    if unused_def_ids_by_file.is_empty() {
-        // There are no unused functions with file names to add (in any CGU)
-        return;
-    }
-
-    // Each `CodegenUnit` (CGU) has its own function_coverage_map, and generates a specific binary
-    // with its own coverage map.
-    //
-    // Each covered function `Instance` can be included in only one coverage map, produced from a
-    // specific function_coverage_map, from a specific CGU.
-    //
-    // Since unused functions did not generate code, they are not associated with any CGU yet.
-    //
-    // To avoid injecting the unused functions in multiple coverage maps (for multiple CGUs)
-    // determine which function_coverage_map has the responsibility for publishing unreachable
-    // coverage, based on file name: For each unused function, find the CGU that generates the
-    // first function (based on sorted `DefId`) from the same file.
-    //
-    // Add a new `FunctionCoverage` to the `function_coverage_map`, with unreachable code regions
-    // for each region in it's MIR.
-
-    // Convert the `HashSet` of `codegenned_def_ids` to a sortable vector, and sort them.
-    let mut sorted_codegenned_def_ids: Vec<DefId> = codegenned_def_ids.iter().copied().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() {
-        if let Some(covered_file_name) = tcx.covered_file_name(def_id) {
-            // Only add files known to have unused functions
-            if unused_def_ids_by_file.contains_key(covered_file_name) {
-                first_covered_def_id_by_file.entry(*covered_file_name).or_insert(def_id);
+        // `all_def_ids` contains things besides just "functions" such as constants,
+        // statics, etc. We need to filter those out.
+        let kind = tcx.def_kind(non_codegenned_def_id);
+        if matches!(kind, DefKind::Fn | DefKind::AssocFn | DefKind::Closure | DefKind::Generator) {
+            let codegen_fn_attrs = tcx.codegen_fn_attrs(non_codegenned_def_id);
+
+            // If a function is marked `#[no_coverage]`, then skip generating a
+            // dead code stub for it.
+            if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
+                debug!("skipping unused fn marked #[no_coverage]: {:?}", non_codegenned_def_id);
+                continue;
             }
-        }
-    }
-
-    // Get the set of def_ids with coverage regions, known by *this* CoverageContext.
-    let cgu_covered_def_ids: DefIdSet = match cx.coverage_context() {
-        Some(ctx) => ctx
-            .function_coverage_map
-            .borrow()
-            .keys()
-            .map(|&instance| instance.def.def_id())
-            .collect(),
-        None => return,
-    };
-
-    let 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();
 
-    // For each file for which this CGU is responsible for adding unused function coverage,
-    // get the `def_id`s for each unused function (if any), define a synthetic function with a
-    // single LLVM coverage counter, and add the function's coverage `CodeRegion`s. to the
-    // function_coverage_map.
-    for covered_file_name in cgu_covered_files {
-        for def_id in unused_def_ids_by_file.remove(&covered_file_name).into_iter().flatten() {
-            cx.define_unused_fn(def_id);
+            debug!("generating unused fn: {:?}", non_codegenned_def_id);
+            cx.define_unused_fn(non_codegenned_def_id);
+        } else {
+            debug!("skipping unused {:?}: {:?}", kind, non_codegenned_def_id);
         }
     }
 }

compiler/rustc_middle/src/mir/mono.rs

@@ -247,6 +247,9 @@ pub struct CodegenUnit<'tcx> {
     items: FxHashMap<MonoItem<'tcx>, (Linkage, Visibility)>,
     size_estimate: Option<usize>,
     primary: bool,
+    /// True if this is CGU is used to hold code coverage information for dead code,
+    /// false otherwise.
+    is_code_coverage_dead_code_cgu: bool,
 }
 
 /// Specifies the linkage type for a `MonoItem`.
@@ -277,7 +280,13 @@ pub enum Visibility {
 impl<'tcx> CodegenUnit<'tcx> {
     #[inline]
     pub fn new(name: Symbol) -> CodegenUnit<'tcx> {
-        CodegenUnit { name, items: Default::default(), size_estimate: None, primary: false }
+        CodegenUnit {
+            name,
+            items: Default::default(),
+            size_estimate: None,
+            primary: false,
+            is_code_coverage_dead_code_cgu: false,
+        }
     }
 
     pub fn name(&self) -> Symbol {
@@ -304,6 +313,15 @@ impl<'tcx> CodegenUnit<'tcx> {
         &mut self.items
     }
 
+    pub fn is_code_coverage_dead_code_cgu(&self) -> bool {
+        self.is_code_coverage_dead_code_cgu
+    }
+
+    /// Marks this CGU as the one used to contain code coverage information for dead code.
+    pub fn make_code_coverage_dead_code_cgu(&mut self) {
+        self.is_code_coverage_dead_code_cgu = true;
+    }
+
     pub fn mangle_name(human_readable_name: &str) -> String {
         // We generate a 80 bit hash from the name. This should be enough to
         // avoid collisions and is still reasonably short for filenames.
@@ -407,9 +425,11 @@ impl<'a, 'tcx> HashStable<StableHashingContext<'a>> for CodegenUnit<'tcx> {
             // The size estimate is not relevant to the hash
             size_estimate: _,
             primary: _,
+            is_code_coverage_dead_code_cgu,
         } = *self;
 
         name.hash_stable(hcx, hasher);
+        is_code_coverage_dead_code_cgu.hash_stable(hcx, hasher);
 
         let mut items: Vec<(Fingerprint, _)> = items
             .iter()

compiler/rustc_middle/src/query/mod.rs

@@ -386,16 +386,6 @@ rustc_queries! {
         storage(ArenaCacheSelector<'tcx>)
     }
 
-    /// 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> {

compiler/rustc_mir_transform/src/coverage/query.rs

@@ -9,7 +9,6 @@ use rustc_span::def_id::DefId;
 /// A `query` provider for retrieving coverage information injected into MIR.
 pub(crate) fn provide(providers: &mut Providers) {
     providers.coverageinfo = |tcx, def_id| coverageinfo(tcx, def_id);
-    providers.covered_file_name = |tcx, def_id| covered_file_name(tcx, def_id);
     providers.covered_code_regions = |tcx, def_id| covered_code_regions(tcx, def_id);
 }
 
@@ -137,25 +136,6 @@ fn coverageinfo<'tcx>(tcx: TyCtxt<'tcx>, instance_def: ty::InstanceDef<'tcx>) ->
     coverage_visitor.info
 }
 
-fn covered_file_name(tcx: TyCtxt<'_>, def_id: DefId) -> Option<Symbol> {
-    if tcx.is_mir_available(def_id) {
-        let body = mir_body(tcx, def_id);
-        for bb_data in body.basic_blocks().iter() {
-            for statement in bb_data.statements.iter() {
-                if let StatementKind::Coverage(box ref coverage) = statement.kind {
-                    if let Some(code_region) = coverage.code_region.as_ref() {
-                        if is_inlined(body, statement) {
-                            continue;
-                        }
-                        return Some(code_region.file_name);
-                    }
-                }
-            }
-        }
-    }
-    return None;
-}
-
 fn covered_code_regions<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Vec<&'tcx CodeRegion> {
     let body = mir_body(tcx, def_id);
     body.basic_blocks()

compiler/rustc_monomorphize/src/partitioning/mod.rs

@@ -201,6 +201,24 @@ pub fn partition<'tcx>(
         partitioner.internalize_symbols(cx, &mut post_inlining);
     }
 
+    let instrument_dead_code =
+        tcx.sess.instrument_coverage() && !tcx.sess.instrument_coverage_except_unused_functions();
+
+    if instrument_dead_code {
+        // Find the smallest CGU that has exported symbols and put the dead
+        // function stubs in that CGU. We look for exported symbols to increase
+        // the likelyhood the linker won't throw away the dead functions.
+        let mut cgus_with_exported_symbols: Vec<_> = post_inlining
+            .codegen_units
+            .iter_mut()
+            .filter(|cgu| cgu.items().iter().any(|(_, (linkage, _))| *linkage == Linkage::External))
+            .collect();
+        cgus_with_exported_symbols.sort_by_key(|cgu| cgu.size_estimate());
+
+        let dead_code_cgu = cgus_with_exported_symbols.last_mut().unwrap();
+        dead_code_cgu.make_code_coverage_dead_code_cgu();
+    }
+
     // Finally, sort by codegen unit name, so that we get deterministic results.
     let PostInliningPartitioning {
         codegen_units: mut result,

src/test/run-make-fulldeps/coverage-reports/expected_show_coverage.unused_mod.txt

@@ -1,3 +1,9 @@
+../coverage/lib/unused_mod_helper.rs:
+    1|      0|pub fn never_called_function() {
+    2|      0|    println!("I am never called");
+    3|      0|}
+
+../coverage/unused_mod.rs:
     1|       |#[path = "lib/unused_mod_helper.rs"]
     2|       |mod unused_module;
     3|       |

src/test/run-make-fulldeps/coverage-reports/expected_show_coverage.uses_crate.txt

@@ -19,12 +19,12 @@
    18|      2|    println!("used_only_from_bin_crate_generic_function with {:?}", arg);
    19|      2|}
   ------------------
-  | used_crate::used_only_from_bin_crate_generic_function::<&alloc::vec::Vec<i32>>:
+  | used_crate::used_only_from_bin_crate_generic_function::<&str>:
   |   17|      1|pub fn used_only_from_bin_crate_generic_function<T: Debug>(arg: T) {
   |   18|      1|    println!("used_only_from_bin_crate_generic_function with {:?}", arg);
   |   19|      1|}
   ------------------
-  | used_crate::used_only_from_bin_crate_generic_function::<&str>:
+  | used_crate::used_only_from_bin_crate_generic_function::<&alloc::vec::Vec<i32>>:
   |   17|      1|pub fn used_only_from_bin_crate_generic_function<T: Debug>(arg: T) {
   |   18|      1|    println!("used_only_from_bin_crate_generic_function with {:?}", arg);
   |   19|      1|}
@@ -36,12 +36,12 @@
    22|      2|    println!("used_only_from_this_lib_crate_generic_function with {:?}", arg);
    23|      2|}
   ------------------
-  | used_crate::used_only_from_this_lib_crate_generic_function::<alloc::vec::Vec<i32>>:
+  | used_crate::used_only_from_this_lib_crate_generic_function::<&str>:
   |   21|      1|pub fn used_only_from_this_lib_crate_generic_function<T: Debug>(arg: T) {
   |   22|      1|    println!("used_only_from_this_lib_crate_generic_function with {:?}", arg);
   |   23|      1|}
   ------------------
-  | used_crate::used_only_from_this_lib_crate_generic_function::<&str>:
+  | used_crate::used_only_from_this_lib_crate_generic_function::<alloc::vec::Vec<i32>>:
   |   21|      1|pub fn used_only_from_this_lib_crate_generic_function<T: Debug>(arg: T) {
   |   22|      1|    println!("used_only_from_this_lib_crate_generic_function with {:?}", arg);
   |   23|      1|}

src/test/run-make-fulldeps/coverage-reports/expected_show_coverage.uses_inline_crate.txt

@@ -42,12 +42,12 @@
    40|      2|    println!("used_only_from_bin_crate_generic_function with {:?}", arg);
    41|      2|}
   ------------------
-  | used_inline_crate::used_only_from_bin_crate_generic_function::<&str>:
+  | used_inline_crate::used_only_from_bin_crate_generic_function::<&alloc::vec::Vec<i32>>:
   |   39|      1|pub fn used_only_from_bin_crate_generic_function<T: Debug>(arg: T) {
   |   40|      1|    println!("used_only_from_bin_crate_generic_function with {:?}", arg);
   |   41|      1|}
   ------------------
-  | used_inline_crate::used_only_from_bin_crate_generic_function::<&alloc::vec::Vec<i32>>:
+  | used_inline_crate::used_only_from_bin_crate_generic_function::<&str>:
   |   39|      1|pub fn used_only_from_bin_crate_generic_function<T: Debug>(arg: T) {
   |   40|      1|    println!("used_only_from_bin_crate_generic_function with {:?}", arg);
   |   41|      1|}
@@ -61,12 +61,12 @@
    46|      4|    println!("used_only_from_this_lib_crate_generic_function with {:?}", arg);
    47|      4|}
   ------------------
-  | used_inline_crate::used_only_from_this_lib_crate_generic_function::<&str>:
+  | used_inline_crate::used_only_from_this_lib_crate_generic_function::<alloc::vec::Vec<i32>>:
   |   45|      2|pub fn used_only_from_this_lib_crate_generic_function<T: Debug>(arg: T) {
   |   46|      2|    println!("used_only_from_this_lib_crate_generic_function with {:?}", arg);
   |   47|      2|}
   ------------------
-  | used_inline_crate::used_only_from_this_lib_crate_generic_function::<alloc::vec::Vec<i32>>:
+  | used_inline_crate::used_only_from_this_lib_crate_generic_function::<&str>:
   |   45|      2|pub fn used_only_from_this_lib_crate_generic_function<T: Debug>(arg: T) {
   |   46|      2|    println!("used_only_from_this_lib_crate_generic_function with {:?}", arg);
   |   47|      2|}