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.

7 files changed, 68 insertions, 51 deletions

diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
index 444a9d4ba04..1cea927f50c 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
+++ b/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.
diff --git a/compiler/rustc_codegen_ssa/src/coverageinfo/map.rs b/compiler/rustc_codegen_ssa/src/coverageinfo/map.rs
index 549b8d41f51..af6482fdbc2 100644
--- a/compiler/rustc_codegen_ssa/src/coverageinfo/map.rs
+++ b/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".
diff --git a/compiler/rustc_codegen_ssa/src/mir/coverageinfo.rs b/compiler/rustc_codegen_ssa/src/mir/coverageinfo.rs
index a115d358666..5ab1baafb57 100644
--- a/compiler/rustc_codegen_ssa/src/mir/coverageinfo.rs
+++ b/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"),
                 );
             }
diff --git a/compiler/rustc_codegen_ssa/src/mir/statement.rs b/compiler/rustc_codegen_ssa/src/mir/statement.rs
index 5523e5f2e86..fe7f6288adb 100644
--- a/compiler/rustc_codegen_ssa/src/mir/statement.rs
+++ b/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 {
diff --git a/compiler/rustc_mir/src/transform/coverage/query.rs b/compiler/rustc_mir/src/transform/coverage/query.rs
index 4b455a6a1ba..de8447f1974 100644
--- a/compiler/rustc_mir/src/transform/coverage/query.rs
+++ b/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()
+}
diff --git a/compiler/rustc_mir/src/transform/inline.rs b/compiler/rustc_mir/src/transform/inline.rs
index cf85503b3d5..26e7555a61d 100644
--- a/compiler/rustc_mir/src/transform/inline.rs
+++ b/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;
     }
diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs
index c1be90efc72..75078a12311 100644
--- a/compiler/rustc_session/src/config.rs
+++ b/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") {