Rollup merge of #83080 - tmiasko:inline-coverage, r=wesleywiser

Make source-based code coverage compatible with MIR inlining

When codegenning code coverage use the instance that coverage data was
originally generated for, to ensure basic level of compatibility with
MIR inlining.

Fixes #83061

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -254,7 +254,7 @@ fn save_function_record(
 ///
 /// 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
+/// 2. The function to which the unreachable code regions will be added must not be a generic
 ///    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.
@@ -284,7 +284,7 @@ fn add_unreachable_coverage<'tcx>(
     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 codegenned_def_ids = tcx.codegened_and_inlined_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) {

compiler/rustc_codegen_ssa/src/coverageinfo/map.rs

@@ -8,7 +8,7 @@ use rustc_middle::mir::coverage::{
 use rustc_middle::ty::Instance;
 use rustc_middle::ty::TyCtxt;
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct Expression {
     lhs: ExpressionOperandId,
     op: Op,
@@ -64,7 +64,9 @@ impl<'tcx> FunctionCoverage<'tcx> {
 
     /// Adds a code region to be counted by an injected counter intrinsic.
     pub fn add_counter(&mut self, id: CounterValueReference, region: CodeRegion) {
-        self.counters[id].replace(region).expect_none("add_counter called with duplicate `id`");
+        if let Some(previous_region) = self.counters[id].replace(region.clone()) {
+            assert_eq!(previous_region, region, "add_counter: code region for id changed");
+        }
     }
 
     /// Both counters and "counter expressions" (or simply, "expressions") can be operands in other
@@ -94,9 +96,18 @@ impl<'tcx> FunctionCoverage<'tcx> {
             expression_id, lhs, op, rhs, region
         );
         let expression_index = self.expression_index(u32::from(expression_id));
-        self.expressions[expression_index]
-            .replace(Expression { lhs, op, rhs, region })
-            .expect_none("add_counter_expression called with duplicate `id_descending_from_max`");
+        if let Some(previous_expression) = self.expressions[expression_index].replace(Expression {
+            lhs,
+            op,
+            rhs,
+            region: region.clone(),
+        }) {
+            assert_eq!(
+                previous_expression,
+                Expression { lhs, op, rhs, region },
+                "add_counter_expression: expression for id changed"
+            );
+        }
     }
 
     /// Add a region that will be marked as "unreachable", with a constant "zero counter".

compiler/rustc_codegen_ssa/src/mir/coverageinfo.rs

@@ -2,27 +2,38 @@ use crate::traits::*;
 
 use rustc_middle::mir::coverage::*;
 use rustc_middle::mir::Coverage;
+use rustc_middle::mir::SourceScope;
 
 use super::FunctionCx;
 
 impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
-    pub fn codegen_coverage(&self, bx: &mut Bx, coverage: Coverage) {
+    pub fn codegen_coverage(&self, bx: &mut Bx, coverage: Coverage, scope: SourceScope) {
+        // Determine the instance that coverage data was originally generated for.
+        let scope_data = &self.mir.source_scopes[scope];
+        let instance = if let Some((inlined_instance, _)) = scope_data.inlined {
+            self.monomorphize(inlined_instance)
+        } else if let Some(inlined_scope) = scope_data.inlined_parent_scope {
+            self.monomorphize(self.mir.source_scopes[inlined_scope].inlined.unwrap().0)
+        } else {
+            self.instance
+        };
+
         let Coverage { kind, code_region } = coverage;
         match kind {
             CoverageKind::Counter { function_source_hash, id } => {
-                if bx.set_function_source_hash(self.instance, function_source_hash) {
+                if bx.set_function_source_hash(instance, function_source_hash) {
                     // If `set_function_source_hash()` returned true, the coverage map is enabled,
                     // so continue adding the counter.
                     if let Some(code_region) = code_region {
                         // Note: Some counters do not have code regions, but may still be referenced
                         // from expressions. In that case, don't add the counter to the coverage map,
                         // but do inject the counter intrinsic.
-                        bx.add_coverage_counter(self.instance, id, code_region);
+                        bx.add_coverage_counter(instance, id, code_region);
                     }
 
-                    let coverageinfo = bx.tcx().coverageinfo(self.instance.def_id());
+                    let coverageinfo = bx.tcx().coverageinfo(instance.def_id());
 
-                    let fn_name = bx.create_pgo_func_name_var(self.instance);
+                    let fn_name = bx.create_pgo_func_name_var(instance);
                     let hash = bx.const_u64(function_source_hash);
                     let num_counters = bx.const_u32(coverageinfo.num_counters);
                     let index = bx.const_u32(u32::from(id));
@@ -34,11 +45,11 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 }
             }
             CoverageKind::Expression { id, lhs, op, rhs } => {
-                bx.add_coverage_counter_expression(self.instance, id, lhs, op, rhs, code_region);
+                bx.add_coverage_counter_expression(instance, id, lhs, op, rhs, code_region);
             }
             CoverageKind::Unreachable => {
                 bx.add_coverage_unreachable(
-                    self.instance,
+                    instance,
                     code_region.expect("unreachable regions always have code regions"),
                 );
             }

compiler/rustc_codegen_ssa/src/mir/statement.rs

@@ -112,7 +112,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 bx
             }
             mir::StatementKind::Coverage(box ref coverage) => {
-                self.codegen_coverage(&mut bx, coverage.clone());
+                self.codegen_coverage(&mut bx, coverage.clone(), statement.source_info.scope);
                 bx
             }
             mir::StatementKind::CopyNonOverlapping(box mir::CopyNonOverlapping {

compiler/rustc_middle/src/query/mod.rs

@@ -1407,6 +1407,14 @@ rustc_queries! {
     query is_codegened_item(def_id: DefId) -> bool {
         desc { |tcx| "determining whether `{}` needs codegen", tcx.def_path_str(def_id) }
     }
+
+    /// All items participating in code generation together with items inlined into them.
+    query codegened_and_inlined_items(_: CrateNum)
+        -> &'tcx DefIdSet {
+        eval_always
+       desc { "codegened_and_inlined_items" }
+    }
+
     query codegen_unit(_: Symbol) -> &'tcx CodegenUnit<'tcx> {
         desc { "codegen_unit" }
     }

compiler/rustc_mir/src/monomorphize/partitioning/mod.rs

@@ -424,8 +424,33 @@ fn collect_and_partition_mono_items<'tcx>(
     (tcx.arena.alloc(mono_items), codegen_units)
 }
 
+fn codegened_and_inlined_items<'tcx>(tcx: TyCtxt<'tcx>, cnum: CrateNum) -> &'tcx DefIdSet {
+    let (items, cgus) = tcx.collect_and_partition_mono_items(cnum);
+    let mut visited = DefIdSet::default();
+    let mut result = items.clone();
+
+    for cgu in cgus {
+        for (item, _) in cgu.items() {
+            if let MonoItem::Fn(ref instance) = item {
+                let did = instance.def_id();
+                if !visited.insert(did) {
+                    continue;
+                }
+                for scope in &tcx.instance_mir(instance.def).source_scopes {
+                    if let Some((ref inlined, _)) = scope.inlined {
+                        result.insert(inlined.def_id());
+                    }
+                }
+            }
+        }
+    }
+
+    tcx.arena.alloc(result)
+}
+
 pub fn provide(providers: &mut Providers) {
     providers.collect_and_partition_mono_items = collect_and_partition_mono_items;
+    providers.codegened_and_inlined_items = codegened_and_inlined_items;
 
     providers.is_codegened_item = |tcx, def_id| {
         let (all_mono_items, _) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);

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

@@ -1,8 +1,7 @@
 use super::*;
 
 use rustc_middle::mir::coverage::*;
-use rustc_middle::mir::visit::Visitor;
-use rustc_middle::mir::{self, Coverage, CoverageInfo, Location};
+use rustc_middle::mir::{self, Body, Coverage, CoverageInfo};
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::{self, TyCtxt};
 use rustc_span::def_id::DefId;
@@ -85,10 +84,21 @@ impl CoverageVisitor {
             }
         }
     }
-}
 
-impl Visitor<'_> for CoverageVisitor {
-    fn visit_coverage(&mut self, coverage: &Coverage, _location: Location) {
+    fn visit_body(&mut self, body: &Body<'_>) {
+        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 is_inlined(body, statement) {
+                        continue;
+                    }
+                    self.visit_coverage(coverage);
+                }
+            }
+        }
+    }
+
+    fn visit_coverage(&mut self, coverage: &Coverage) {
         if self.add_missing_operands {
             match coverage.kind {
                 CoverageKind::Expression { lhs, rhs, .. } => {
@@ -129,10 +139,14 @@ fn coverageinfo_from_mir<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> CoverageInfo
 }
 
 fn covered_file_name<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Symbol> {
-    for bb_data in mir_body(tcx, def_id).basic_blocks().iter() {
+    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);
                 }
             }
@@ -151,13 +165,17 @@ fn mir_body<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> &'tcx mir::Body<'tcx> {
 }
 
 fn covered_code_regions<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Vec<&'tcx CodeRegion> {
-    mir_body(tcx, def_id)
-        .basic_blocks()
+    let body = mir_body(tcx, def_id);
+    body.basic_blocks()
         .iter()
         .map(|data| {
             data.statements.iter().filter_map(|statement| match statement.kind {
                 StatementKind::Coverage(box ref coverage) => {
-                    coverage.code_region.as_ref() // may be None
+                    if is_inlined(body, statement) {
+                        None
+                    } else {
+                        coverage.code_region.as_ref() // may be None
+                    }
                 }
                 _ => None,
             })
@@ -165,3 +183,8 @@ fn covered_code_regions<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Vec<&'tcx Cod
         .flatten()
         .collect()
 }
+
+fn is_inlined(body: &Body<'_>, statement: &Statement<'_>) -> bool {
+    let scope_data = &body.source_scopes[statement.source_info.scope];
+    scope_data.inlined.is_some() || scope_data.inlined_parent_scope.is_some()
+}

compiler/rustc_mir/src/transform/inline.rs

@@ -39,15 +39,6 @@ struct CallSite<'tcx> {
 
 /// Returns true if MIR inlining is enabled in the current compilation session.
 crate fn is_enabled(tcx: TyCtxt<'_>) -> bool {
-    if tcx.sess.opts.debugging_opts.instrument_coverage {
-        // Since `Inline` happens after `InstrumentCoverage`, the function-specific coverage
-        // counters can be invalidated, such as by merging coverage counter statements from
-        // a pre-inlined function into a different function. This kind of change is invalid,
-        // so inlining must be skipped. Note: This check is performed here so inlining can
-        // be disabled without preventing other optimizations (regardless of `mir_opt_level`).
-        return false;
-    }
-
     if let Some(enabled) = tcx.sess.opts.debugging_opts.inline_mir {
         return enabled;
     }

compiler/rustc_session/src/config.rs

@@ -1937,25 +1937,6 @@ pub fn build_session_options(matches: &getopts::Matches) -> Options {
             }
             Some(SymbolManglingVersion::V0) => {}
         }
-
-        if let Some(mir_opt_level) = debugging_opts.mir_opt_level {
-            if mir_opt_level > 1 {
-                // Functions inlined during MIR transform can, at best, make it impossible to
-                // effectively cover inlined functions, and, at worst, break coverage map generation
-                // during LLVM codegen. For example, function counter IDs are only unique within a
-                // function. Inlining after these counters are injected can produce duplicate counters,
-                // resulting in an invalid coverage map (and ICE); so this option combination is not
-                // allowed.
-                early_warn(
-                    error_format,
-                    &format!(
-                        "`-Z mir-opt-level={}` (or any level > 1) enables function inlining, which \
-                    is incompatible with `-Z instrument-coverage`. Inlining will be disabled.",
-                        mir_opt_level,
-                    ),
-                );
-            }
-        }
     }
 
     if let Ok(graphviz_font) = std::env::var("RUSTC_GRAPHVIZ_FONT") {