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.

Author: tmiasko

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.

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_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,17 +165,26 @@ 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,
             })
         })
         .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") {