Make source-based code coverage compatible with MIR inlining

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

@@ -1397,6 +1397,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,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") {

src/test/run-make-fulldeps/coverage-reports/Makefile

@@ -82,13 +82,13 @@ endif
 %: $(SOURCEDIR)/lib/%.rs
 	# Compile the test library with coverage instrumentation
 	$(RUSTC) $(SOURCEDIR)/lib/$@.rs \
-			$$( grep -q '^\/\/ require-rust-edition-2018' $(SOURCEDIR)/lib/$@.rs && echo "--edition=2018" ) \
+			$$( sed -nE 's#^// compile-flags:(.*)#\1# p' $(SOURCEDIR)/lib/$@.rs) \
 			--crate-type rlib -Zinstrument-coverage
 
 %: $(SOURCEDIR)/%.rs
 	# Compile the test program with coverage instrumentation
 	$(RUSTC) $(SOURCEDIR)/$@.rs \
-			$$( grep -q '^\/\/ require-rust-edition-2018' $(SOURCEDIR)/$@.rs && echo "--edition=2018" ) \
+			$$( sed -nE 's#^// compile-flags:(.*)#\1# p' $(SOURCEDIR)/$@.rs) \
 			-L "$(TMPDIR)" -Zinstrument-coverage
 
 	# Run it in order to generate some profiling data,
@@ -107,7 +107,7 @@ endif
 	# Run it through rustdoc as well to cover doctests
 	LLVM_PROFILE_FILE="$(TMPDIR)"/$@-%p.profraw \
 			$(RUSTDOC) --crate-name workaround_for_79771 --test $(SOURCEDIR)/$@.rs \
-			$$( grep -q '^\/\/ require-rust-edition-2018' $(SOURCEDIR)/$@.rs && echo "--edition=2018" ) \
+			$$( sed -nE 's#^// compile-flags:(.*)#\1# p' $(SOURCEDIR)/$@.rs) \
 			-L "$(TMPDIR)" -Zinstrument-coverage \
 			-Z unstable-options --persist-doctests=$(TMPDIR)/rustdoc-$@
 

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

@@ -1,6 +1,6 @@
     1|       |#![allow(unused_assignments, dead_code)]
     2|       |
-    3|       |// require-rust-edition-2018
+    3|       |// compile-flags: --edition=2018
     4|       |
     5|      1|async fn c(x: u8) -> u8 {
     6|      1|    if x == 8 {

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

@@ -0,0 +1,53 @@
+    1|       |// compile-flags: -Zinline-mir
+    2|       |
+    3|       |use std::fmt::Display;
+    4|       |
+    5|      1|fn main() {
+    6|      1|    permutations(&['a', 'b', 'c']);
+    7|      1|}
+    8|       |
+    9|       |#[inline(always)]
+   10|      1|fn permutations<T: Copy + Display>(xs: &[T]) {
+   11|      1|    let mut ys = xs.to_owned();
+   12|      1|    permutate(&mut ys, 0);
+   13|      1|}
+   14|       |
+   15|     16|fn permutate<T: Copy + Display>(xs: &mut [T], k: usize) {
+   16|     16|    let n = length(xs);
+   17|     16|    if k == n {
+   18|      6|        display(xs);
+   19|     10|    } else if k < n {
+   20|     15|        for i in k..n {
+                               ^10
+   21|     15|            swap(xs, i, k);
+   22|     15|            permutate(xs, k + 1);
+   23|     15|            swap(xs, i, k);
+   24|     15|        }
+   25|      0|    } else {
+   26|      0|        error();
+   27|      0|    }
+   28|     16|}
+   29|       |
+   30|     16|fn length<T>(xs: &[T]) -> usize {
+   31|     16|    xs.len()
+   32|     16|}
+   33|       |
+   34|       |#[inline]
+   35|     30|fn swap<T: Copy>(xs: &mut [T], i: usize, j: usize) {
+   36|     30|    let t = xs[i];
+   37|     30|    xs[i] = xs[j];
+   38|     30|    xs[j] = t;
+   39|     30|}
+   40|       |
+   41|      6|fn display<T: Display>(xs: &[T]) {
+   42|     18|    for x in xs {
+   43|     18|        print!("{}", x);
+   44|     18|    }
+   45|      6|    println!();
+   46|      6|}
+   47|       |
+   48|       |#[inline(always)]
+   49|       |fn error() {
+   50|       |    panic!("error");
+   51|       |}
+