Move coverage region collection to CodegenCx finalization

Moved from `query coverageinfo` (renamed from `query coverage_data`),
as discussed in the PR at:

#73684 (comment)

Author: richkadel

src/libcore/intrinsics.rs

@@ -1958,6 +1958,39 @@ extern "rust-intrinsic" {
     #[lang = "count_code_region"]
     pub fn count_code_region(index: u32, start_byte_pos: u32, end_byte_pos: u32);
 
+    /// Internal marker for code coverage expressions, injected into the MIR when the
+    /// "instrument-coverage" option is enabled. This intrinsic is not converted into a
+    /// backend intrinsic call, but its arguments are extracted during the production of a
+    /// "coverage map", which is injected into the generated code, as additional data.
+    /// This marker identifies a code region and two other counters or counter expressions
+    /// whose sum is the number of times the code region was executed.
+    #[cfg(not(bootstrap))]
+    pub fn coverage_counter_add(
+        index: u32,
+        left_index: u32,
+        right_index: u32,
+        start_byte_pos: u32,
+        end_byte_pos: u32,
+    );
+
+    /// This marker identifies a code region and two other counters or counter expressions
+    /// whose difference is the number of times the code region was executed.
+    /// (See `coverage_counter_add` for more information.)
+    #[cfg(not(bootstrap))]
+    pub fn coverage_counter_subtract(
+        index: u32,
+        left_index: u32,
+        right_index: u32,
+        start_byte_pos: u32,
+        end_byte_pos: u32,
+    );
+
+    /// This marker identifies a code region to be added to the "coverage map" to indicate source
+    /// code that can never be reached.
+    /// (See `coverage_counter_add` for more information.)
+    #[cfg(not(bootstrap))]
+    pub fn coverage_unreachable(start_byte_pos: u32, end_byte_pos: u32);
+
     /// See documentation of `<*const T>::guaranteed_eq` for details.
     #[rustc_const_unstable(feature = "const_raw_ptr_comparison", issue = "53020")]
     #[cfg(not(bootstrap))]

src/librustc_codegen_llvm/context.rs

@@ -1,5 +1,6 @@
 use crate::attributes;
 use crate::callee::get_fn;
+use crate::coverageinfo;
 use crate::debuginfo;
 use crate::llvm;
 use crate::llvm_util;
@@ -77,6 +78,7 @@ pub struct CodegenCx<'ll, 'tcx> {
     pub pointee_infos: RefCell<FxHashMap<(Ty<'tcx>, Size), Option<PointeeInfo>>>,
     pub isize_ty: &'ll Type,
 
+    pub coverage_cx: Option<coverageinfo::CrateCoverageContext<'tcx>>,
     pub dbg_cx: Option<debuginfo::CrateDebugContext<'ll, 'tcx>>,
 
     eh_personality: Cell<Option<&'ll Value>>,
@@ -256,6 +258,13 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
 
         let (llcx, llmod) = (&*llvm_module.llcx, llvm_module.llmod());
 
+        let coverage_cx = if tcx.sess.opts.debugging_opts.instrument_coverage {
+            let covctx = coverageinfo::CrateCoverageContext::new();
+            Some(covctx)
+        } else {
+            None
+        };
+
         let dbg_cx = if tcx.sess.opts.debuginfo != DebugInfo::None {
             let dctx = debuginfo::CrateDebugContext::new(llmod);
             debuginfo::metadata::compile_unit_metadata(tcx, &codegen_unit.name().as_str(), &dctx);
@@ -285,6 +294,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
             scalar_lltypes: Default::default(),
             pointee_infos: Default::default(),
             isize_ty,
+            coverage_cx,
             dbg_cx,
             eh_personality: Cell::new(None),
             rust_try_fn: Cell::new(None),
@@ -296,6 +306,11 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
     crate fn statics_to_rauw(&self) -> &RefCell<Vec<(&'ll Value, &'ll Value)>> {
         &self.statics_to_rauw
     }
+
+    #[inline]
+    pub fn coverage_context(&'a self) -> &'a coverageinfo::CrateCoverageContext<'tcx> {
+        self.coverage_cx.as_ref().unwrap()
+    }
 }
 
 impl MiscMethods<'tcx> for CodegenCx<'ll, 'tcx> {
@@ -749,8 +764,6 @@ impl CodegenCx<'b, 'tcx> {
         ifn!("llvm.lifetime.start.p0i8", fn(t_i64, i8p) -> void);
         ifn!("llvm.lifetime.end.p0i8", fn(t_i64, i8p) -> void);
 
-        ifn!("llvm.instrprof.increment", fn(i8p, t_i64, t_i32, t_i32) -> void);
-
         ifn!("llvm.expect.i1", fn(i1, i1) -> i1);
         ifn!("llvm.eh.typeid.for", fn(i8p) -> t_i32);
         ifn!("llvm.localescape", fn(...) -> void);
@@ -765,6 +778,10 @@ impl CodegenCx<'b, 'tcx> {
         ifn!("llvm.va_end", fn(i8p) -> void);
         ifn!("llvm.va_copy", fn(i8p, i8p) -> void);
 
+        if self.sess().opts.debugging_opts.instrument_coverage {
+            ifn!("llvm.instrprof.increment", fn(i8p, t_i64, t_i32, t_i32) -> void);
+        }
+
         if self.sess().opts.debuginfo != DebugInfo::None {
             ifn!("llvm.dbg.declare", fn(self.type_metadata(), self.type_metadata()) -> void);
             ifn!("llvm.dbg.value", fn(self.type_metadata(), t_i64, self.type_metadata()) -> void);

src/librustc_codegen_llvm/coverageinfo/mod.rs

@@ -1,55 +1,61 @@
+use crate::builder::Builder;
 use crate::common::CodegenCx;
 use log::debug;
-use rustc_codegen_ssa::traits::CoverageInfoMethods;
-use rustc_hir::def_id::LOCAL_CRATE;
+use rustc_codegen_ssa::coverageinfo::map::*;
+use rustc_codegen_ssa::traits::{CoverageInfoBuilderMethods, CoverageInfoMethods};
+use rustc_data_structures::fx::FxHashMap;
+use rustc_middle::ty::Instance;
 
-// FIXME(richkadel): This might not be where I put the coverage map generation, but it's a
-// placeholder for now.
-//
-// I need to inject the coverage map into the LLVM IR. This is one possible place to do it.
-//
-// The code below is here, at least temporarily, to verify that the `CoverageRegion`s are
-// available, to produce the coverage map.
+use std::cell::RefCell;
+
+/// A context object for maintaining all state needed by the coverageinfo module.
+pub struct CrateCoverageContext<'tcx> {
+    // Coverage region data for each instrumented function identified by DefId.
+    pub(crate) coverage_regions: RefCell<FxHashMap<Instance<'tcx>, FunctionCoverageRegions>>,
+}
+
+impl<'tcx> CrateCoverageContext<'tcx> {
+    pub fn new() -> Self {
+        Self { coverage_regions: Default::default() }
+    }
+}
 
 /// Generates and exports the Coverage Map.
+// FIXME(richkadel): Actually generate and export the coverage map to LLVM.
+// The current implementation is actually just debug messages to show the data is available.
 pub fn finalize(cx: &CodegenCx<'_, '_>) {
-    let instances = &*cx.instances.borrow();
-    for instance in instances.keys() {
-        if instance.def_id().krate == LOCAL_CRATE {
-            // FIXME(richkadel): Is this check `krate == LOCAL_CRATE` right?
-            //
-            // NOTE: I'm considering at least one alternative to this loop on `instances`,
-            // But if I do go with this, make sure the check on LOCAL_CRATE works for all cases.
-            //
-            //
-            // Without it, I was crashing on some of the instances, with:
-            //
-            //     src/librustc_metadata/rmeta/decoder.rs:1127:17: get_optimized_mir: missing MIR for `DefId(1:2867 ~ std[70cc]::io[0]::stdio[0]::_print[0])`
-            //
-            // This check avoids it, but still works for the very limited testing I've done so far.
-            let coverageinfo = cx.tcx.coverageinfo(instance.def_id());
-            if coverageinfo.num_counters > 0 {
-                debug!(
-                    "Generate coverage map for: {}, hash: {:?}, num_counters: {}\n{}",
-                    instance,
-                    coverageinfo.hash,
-                    coverageinfo.num_counters,
-                    {
-                        let regions = &coverageinfo.coverage_regions;
-                        (0..regions.len() as u32)
-                            .map(|counter_index| {
-                                let region = &regions[counter_index];
-                                format!(
-                                    "  counter_index {} byte range: {}..{}",
-                                    counter_index, region.start_byte_pos, region.end_byte_pos
-                                )
-                            })
-                            .collect::<Vec<String>>()
-                            .join("\n")
-                    }
-                );
+    let coverage_regions = &*cx.coverage_context().coverage_regions.borrow();
+    for instance in coverage_regions.keys() {
+        let coverageinfo = cx.tcx.coverageinfo(instance.def_id());
+        debug_assert!(coverageinfo.num_counters > 0);
+        debug!(
+            "Generate coverage map for: {:?}, hash: {}, num_counters: {}",
+            instance, coverageinfo.hash, coverageinfo.num_counters
+        );
+        let function_coverage_regions = &coverage_regions[instance];
+        for (index, region) in function_coverage_regions.indexed_regions() {
+            match region.kind {
+                CoverageKind::Counter => debug!(
+                    "  Counter {}, for {}..{}",
+                    index, region.coverage_span.start_byte_pos, region.coverage_span.end_byte_pos
+                ),
+                CoverageKind::CounterExpression(lhs, op, rhs) => debug!(
+                    "  CounterExpression {} = {} {:?} {}, for {}..{}",
+                    index,
+                    lhs,
+                    op,
+                    rhs,
+                    region.coverage_span.start_byte_pos,
+                    region.coverage_span.end_byte_pos
+                ),
             }
         }
+        for unreachable in function_coverage_regions.unreachable_regions() {
+            debug!(
+                "  Unreachable code region: {}..{}",
+                unreachable.start_byte_pos, unreachable.end_byte_pos
+            );
+        }
     }
 }
 
@@ -58,3 +64,49 @@ impl CoverageInfoMethods for CodegenCx<'ll, 'tcx> {
         finalize(self)
     }
 }
+
+impl CoverageInfoBuilderMethods<'tcx> for Builder<'a, 'll, 'tcx> {
+    fn new_counter_region(
+        &mut self,
+        instance: Instance<'tcx>,
+        index: u32,
+        start_byte_pos: u32,
+        end_byte_pos: u32,
+    ) {
+        let mut coverage_regions = self.coverage_context().coverage_regions.borrow_mut();
+        coverage_regions
+            .entry(instance)
+            .or_insert_with(|| FunctionCoverageRegions::new())
+            .new_counter(index, start_byte_pos, end_byte_pos);
+    }
+
+    fn new_counter_expression_region(
+        &mut self,
+        instance: Instance<'tcx>,
+        index: u32,
+        lhs: u32,
+        op: CounterOp,
+        rhs: u32,
+        start_byte_pos: u32,
+        end_byte_pos: u32,
+    ) {
+        let mut coverage_regions = self.coverage_context().coverage_regions.borrow_mut();
+        coverage_regions
+            .entry(instance)
+            .or_insert_with(|| FunctionCoverageRegions::new())
+            .new_counter_expression(index, lhs, op, rhs, start_byte_pos, end_byte_pos);
+    }
+
+    fn new_unreachable_region(
+        &mut self,
+        instance: Instance<'tcx>,
+        start_byte_pos: u32,
+        end_byte_pos: u32,
+    ) {
+        let mut coverage_regions = self.coverage_context().coverage_regions.borrow_mut();
+        coverage_regions
+            .entry(instance)
+            .or_insert_with(|| FunctionCoverageRegions::new())
+            .new_unreachable(start_byte_pos, end_byte_pos);
+    }
+}

src/librustc_codegen_llvm/intrinsic.rs

@@ -141,7 +141,6 @@ impl IntrinsicCallMethods<'tcx> for Builder<'a, 'll, 'tcx> {
                 self.call(llfn, &[], None)
             }
             "count_code_region" => {
-                use coverage::count_code_region_args::*;
                 // FIXME(richkadel): The current implementation assumes the MIR for the given
                 // caller_instance represents a single function. Validate and/or correct if inlining
                 // and/or monomorphization invalidates these assumptions.
@@ -150,15 +149,11 @@ impl IntrinsicCallMethods<'tcx> for Builder<'a, 'll, 'tcx> {
                 let (mangled_fn_name, _len_val) = self.const_str(mangled_fn.name);
                 let hash = self.const_u64(coverageinfo.hash);
                 let num_counters = self.const_u32(coverageinfo.num_counters);
+                use coverage::count_code_region_args::*;
                 let index = args[COUNTER_INDEX].immediate();
                 debug!(
-                    "count_code_region to LLVM intrinsic instrprof.increment(fn_name={}, hash={:?}, num_counters={:?}, index={:?}), byte range {:?}..{:?}",
-                    mangled_fn.name,
-                    hash,
-                    num_counters,
-                    index,
-                    args[START_BYTE_POS].immediate(),
-                    args[END_BYTE_POS].immediate(),
+                    "count_code_region to LLVM intrinsic instrprof.increment(fn_name={}, hash={:?}, num_counters={:?}, index={:?})",
+                    mangled_fn.name, hash, num_counters, index,
                 );
                 self.instrprof_increment(mangled_fn_name, hash, num_counters, index)
             }

src/librustc_codegen_ssa/coverageinfo/map.rs

@@ -0,0 +1,86 @@
+use rustc_data_structures::fx::FxHashMap;
+use std::collections::hash_map;
+use std::slice;
+
+#[derive(Copy, Clone, Debug)]
+pub enum CounterOp {
+    Add,
+    Subtract,
+}
+
+pub enum CoverageKind {
+    Counter,
+    CounterExpression(u32, CounterOp, u32),
+}
+
+pub struct CoverageSpan {
+    pub start_byte_pos: u32,
+    pub end_byte_pos: u32,
+}
+
+pub struct CoverageRegion {
+    pub kind: CoverageKind,
+    pub coverage_span: CoverageSpan,
+}
+
+/// Collects all of the coverage regions associated with (a) injected counters, (b) counter
+/// expressions (additions or subtraction), and (c) unreachable regions (always counted as zero),
+/// for a given Function. Counters and counter expressions are indexed because they can be operands
+/// in an expression.
+///
+/// Note, it's important to distinguish the `unreachable` region type from what LLVM's refers to as
+/// a "gap region" (or "gap area"). A gap region is a code region within a counted region (either
+/// counter or expression), but the line or lines in the gap region are not executable (such as
+/// lines with only whitespace or comments). According to LLVM Code Coverage Mapping documentation,
+/// "A count for a gap area is only used as the line execution count if there are no other regions
+/// on a line."
+pub struct FunctionCoverageRegions {
+    indexed: FxHashMap<u32, CoverageRegion>,
+    unreachable: Vec<CoverageSpan>,
+}
+
+impl FunctionCoverageRegions {
+    pub fn new() -> Self {
+        Self { indexed: FxHashMap::default(), unreachable: Default::default() }
+    }
+
+    pub fn new_counter(&mut self, index: u32, start_byte_pos: u32, end_byte_pos: u32) {
+        self.indexed.insert(
+            index,
+            CoverageRegion {
+                kind: CoverageKind::Counter,
+                coverage_span: CoverageSpan { start_byte_pos, end_byte_pos },
+            },
+        );
+    }
+
+    pub fn new_counter_expression(
+        &mut self,
+        index: u32,
+        lhs: u32,
+        op: CounterOp,
+        rhs: u32,
+        start_byte_pos: u32,
+        end_byte_pos: u32,
+    ) {
+        self.indexed.insert(
+            index,
+            CoverageRegion {
+                kind: CoverageKind::CounterExpression(lhs, op, rhs),
+                coverage_span: CoverageSpan { start_byte_pos, end_byte_pos },
+            },
+        );
+    }
+
+    pub fn new_unreachable(&mut self, start_byte_pos: u32, end_byte_pos: u32) {
+        self.unreachable.push(CoverageSpan { start_byte_pos, end_byte_pos });
+    }
+
+    pub fn indexed_regions(&self) -> hash_map::Iter<'_, u32, CoverageRegion> {
+        self.indexed.iter()
+    }
+
+    pub fn unreachable_regions(&self) -> slice::Iter<'_, CoverageSpan> {
+        self.unreachable.iter()
+    }
+}

src/librustc_codegen_ssa/coverageinfo/mod.rs

@@ -0,0 +1,3 @@
+pub mod map;
+
+pub use map::CounterOp;

src/librustc_codegen_ssa/lib.rs

@@ -34,6 +34,7 @@ use std::path::{Path, PathBuf};
 pub mod back;
 pub mod base;
 pub mod common;
+pub mod coverageinfo;
 pub mod debuginfo;
 pub mod glue;
 pub mod meth;