Rollup merge of rust-lang#73684 - richkadel:llvm-coverage-map-gen-2, r=wesleywiser

add spans to injected coverage counters, extract with CoverageData query

This is the next iteration on the Rust Coverage implementation, and follows PR rust-lang#73488

@tmandry @wesleywiser

I came up with an approach for coverage spans, pushing them through the Call terminator as additional args so they can be extracted by the CoverageData query.

I'm using an IndexVec to store them in CoverageData such that there can be only one per index (even if parts of the MIR get duplicated during optimization).

If this approach works for you, I can quickly expand on this to build a separate IndexVec for counter expressions, using a separate call that will be ignored during code generation, but from which I can extract the counter expression values.

Let me know your thoughts. Thanks!

r? @tmandry

Rust compiler MCP rust-lang/compiler-team#278
Relevant issue: rust-lang#34701 - Implement support for LLVMs code coverage instrumentation

Author: Manishearth

Diff for: src/libcore/intrinsics.rs

@@ -1958,7 +1958,40 @@ extern "rust-intrinsic" {
     /// generation.
     #[cfg(not(bootstrap))]
     #[lang = "count_code_region"]
-    pub fn count_code_region(index: u32);
+    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")]

Diff for: src/librustc_codegen_llvm/base.rs

@@ -150,6 +150,11 @@ pub fn compile_codegen_unit(
                 cx.create_used_variable()
             }
 
+            // Finalize code coverage by injecting the coverage map
+            if cx.sess().opts.debugging_opts.instrument_coverage {
+                cx.coverageinfo_finalize();
+            }
+
             // Finalize debuginfo
             if cx.sess().opts.debuginfo != DebugInfo::None {
                 cx.debuginfo_finalize();

Diff for: 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);

Diff for: src/librustc_codegen_llvm/coverageinfo/mod.rs

@@ -0,0 +1,126 @@
+use crate::builder::Builder;
+use crate::common::CodegenCx;
+use log::debug;
+use rustc_codegen_ssa::coverageinfo::map::*;
+use rustc_codegen_ssa::traits::{CoverageInfoBuilderMethods, CoverageInfoMethods};
+use rustc_data_structures::fx::FxHashMap;
+use rustc_middle::ty::Instance;
+
+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 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
+            );
+        }
+    }
+}
+
+impl CoverageInfoMethods for CodegenCx<'ll, 'tcx> {
+    fn coverageinfo_finalize(&self) {
+        finalize(self)
+    }
+}
+
+impl CoverageInfoBuilderMethods<'tcx> for Builder<'a, 'll, 'tcx> {
+    fn add_counter_region(
+        &mut self,
+        instance: Instance<'tcx>,
+        index: u32,
+        start_byte_pos: u32,
+        end_byte_pos: u32,
+    ) {
+        debug!(
+            "adding counter to coverage map: instance={:?}, index={}, byte range {}..{}",
+            instance, index, start_byte_pos, end_byte_pos,
+        );
+        let mut coverage_regions = self.coverage_context().coverage_regions.borrow_mut();
+        coverage_regions.entry(instance).or_default().add_counter(
+            index,
+            start_byte_pos,
+            end_byte_pos,
+        );
+    }
+
+    fn add_counter_expression_region(
+        &mut self,
+        instance: Instance<'tcx>,
+        index: u32,
+        lhs: u32,
+        op: CounterOp,
+        rhs: u32,
+        start_byte_pos: u32,
+        end_byte_pos: u32,
+    ) {
+        debug!(
+            "adding counter expression to coverage map: instance={:?}, index={}, {} {:?} {}, byte range {}..{}",
+            instance, index, lhs, op, rhs, start_byte_pos, end_byte_pos,
+        );
+        let mut coverage_regions = self.coverage_context().coverage_regions.borrow_mut();
+        coverage_regions.entry(instance).or_default().add_counter_expression(
+            index,
+            lhs,
+            op,
+            rhs,
+            start_byte_pos,
+            end_byte_pos,
+        );
+    }
+
+    fn add_unreachable_region(
+        &mut self,
+        instance: Instance<'tcx>,
+        start_byte_pos: u32,
+        end_byte_pos: u32,
+    ) {
+        debug!(
+            "adding unreachable code to coverage map: instance={:?}, byte range {}..{}",
+            instance, start_byte_pos, end_byte_pos,
+        );
+        let mut coverage_regions = self.coverage_context().coverage_regions.borrow_mut();
+        coverage_regions.entry(instance).or_default().add_unreachable(start_byte_pos, end_byte_pos);
+    }
+}

Diff for: src/librustc_codegen_llvm/intrinsic.rs

@@ -13,12 +13,15 @@ use rustc_ast::ast;
 use rustc_codegen_ssa::base::{compare_simd_types, to_immediate, wants_msvc_seh};
 use rustc_codegen_ssa::common::span_invalid_monomorphization_error;
 use rustc_codegen_ssa::common::{IntPredicate, TypeKind};
+use rustc_codegen_ssa::coverageinfo::CounterOp;
 use rustc_codegen_ssa::glue;
 use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
 use rustc_codegen_ssa::mir::place::PlaceRef;
 use rustc_codegen_ssa::traits::*;
 use rustc_codegen_ssa::MemFlags;
 use rustc_hir as hir;
+use rustc_middle::mir::coverage;
+use rustc_middle::mir::Operand;
 use rustc_middle::ty::layout::{FnAbiExt, HasTyCtxt};
 use rustc_middle::ty::{self, Ty};
 use rustc_middle::{bug, span_bug};
@@ -81,6 +84,53 @@ fn get_simple_intrinsic(cx: &CodegenCx<'ll, '_>, name: &str) -> Option<&'ll Valu
 }
 
 impl IntrinsicCallMethods<'tcx> for Builder<'a, 'll, 'tcx> {
+    fn is_codegen_intrinsic(
+        &mut self,
+        intrinsic: &str,
+        args: &Vec<Operand<'tcx>>,
+        caller_instance: ty::Instance<'tcx>,
+    ) -> bool {
+        match intrinsic {
+            "count_code_region" => {
+                use coverage::count_code_region_args::*;
+                self.add_counter_region(
+                    caller_instance,
+                    op_to_u32(&args[COUNTER_INDEX]),
+                    op_to_u32(&args[START_BYTE_POS]),
+                    op_to_u32(&args[END_BYTE_POS]),
+                );
+                true // Also inject the counter increment in the backend
+            }
+            "coverage_counter_add" | "coverage_counter_subtract" => {
+                use coverage::coverage_counter_expression_args::*;
+                self.add_counter_expression_region(
+                    caller_instance,
+                    op_to_u32(&args[COUNTER_EXPRESSION_INDEX]),
+                    op_to_u32(&args[LEFT_INDEX]),
+                    if intrinsic == "coverage_counter_add" {
+                        CounterOp::Add
+                    } else {
+                        CounterOp::Subtract
+                    },
+                    op_to_u32(&args[RIGHT_INDEX]),
+                    op_to_u32(&args[START_BYTE_POS]),
+                    op_to_u32(&args[END_BYTE_POS]),
+                );
+                false // Does not inject backend code
+            }
+            "coverage_unreachable" => {
+                use coverage::coverage_unreachable_args::*;
+                self.add_unreachable_region(
+                    caller_instance,
+                    op_to_u32(&args[START_BYTE_POS]),
+                    op_to_u32(&args[END_BYTE_POS]),
+                );
+                false // Does not inject backend code
+            }
+            _ => true, // Unhandled intrinsics should be passed to `codegen_intrinsic_call()`
+        }
+    }
+
     fn codegen_intrinsic_call(
         &mut self,
         instance: ty::Instance<'tcx>,
@@ -143,15 +193,16 @@ impl IntrinsicCallMethods<'tcx> for Builder<'a, 'll, 'tcx> {
                 // 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.
-                let coverage_data = tcx.coverage_data(caller_instance.def_id());
+                let coverageinfo = tcx.coverageinfo(caller_instance.def_id());
                 let mangled_fn = tcx.symbol_name(caller_instance);
                 let (mangled_fn_name, _len_val) = self.const_str(mangled_fn.name);
-                let hash = self.const_u64(coverage_data.hash);
-                let num_counters = self.const_u32(coverage_data.num_counters);
-                let index = args[0].immediate();
+                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={:?})",
-                    mangled_fn.name, hash, num_counters, index
+                    mangled_fn.name, hash, num_counters, index,
                 );
                 self.instrprof_increment(mangled_fn_name, hash, num_counters, index)
             }
@@ -2131,3 +2182,7 @@ fn float_type_width(ty: Ty<'_>) -> Option<u64> {
         _ => None,
     }
 }
+
+fn op_to_u32<'tcx>(op: &Operand<'tcx>) -> u32 {
+    Operand::scalar_from_const(op).to_u32().expect("Scalar is u32")
+}

Diff for: src/librustc_codegen_llvm/lib.rs

@@ -55,6 +55,7 @@ mod callee;
 mod common;
 mod consts;
 mod context;
+mod coverageinfo;
 mod debuginfo;
 mod declare;
 mod intrinsic;