[WIP] Added spans to count_code_region() Call terminator

and added an index of CoverageRegion (start/end byte pos from the spans)
to CoverageData.

This should give me what I need to build coverage maps for basic
counters.

I also added FIXME notes on how I plan to include regions with counter
expressions instead of counters.

Author: richkadel

src/libcore/intrinsics.rs

@@ -1956,7 +1956,7 @@ 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);
 
     /// See documentation of `<*const T>::guaranteed_eq` for details.
     #[rustc_const_unstable(feature = "const_raw_ptr_comparison", issue = "53020")]

src/librustc_codegen_llvm/intrinsic.rs

@@ -140,6 +140,9 @@ impl IntrinsicCallMethods<'tcx> for Builder<'a, 'll, 'tcx> {
                 self.call(llfn, &[], None)
             }
             "count_code_region" => {
+                use rustc_middle::mir::count_code_region_args::{
+                    COUNTER_INDEX, END_BYTE_POS, START_BYTE_POS,
+                };
                 // 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.
@@ -148,10 +151,16 @@ 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(coverage_data.hash);
                 let num_counters = self.const_u32(coverage_data.num_counters);
-                let index = args[0].immediate();
+                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
+                    "count_code_region to LLVM intrinsic instrprof.increment(fn_name={}, hash={:?}, num_counters={:?}, index={:?}), byte range {:?}..{:?}, coverage_regions: {:?}",
+                    mangled_fn.name,
+                    hash,
+                    num_counters,
+                    index,
+                    args[START_BYTE_POS].immediate(),
+                    args[END_BYTE_POS].immediate(),
+                    coverage_data.coverage_regions,
                 );
                 self.instrprof_increment(mangled_fn_name, hash, num_counters, index)
             }

src/librustc_llvm/build.rs

@@ -169,7 +169,6 @@ fn main() {
     cfg.file("../rustllvm/PassWrapper.cpp")
         .file("../rustllvm/RustWrapper.cpp")
         .file("../rustllvm/ArchiveWrapper.cpp")
-        .file("../rustllvm/CoverageMappingWrapper.cpp")
         .file("../rustllvm/Linker.cpp")
         .cpp(true)
         .cpp_link_stdlib(None) // we handle this below

src/librustc_middle/arena.rs

@@ -94,6 +94,7 @@ macro_rules! arena_types {
             [] upvars_mentioned: rustc_data_structures::fx::FxIndexMap<rustc_hir::HirId, rustc_hir::Upvar>, rustc_data_structures::fx::FxIndexMap<rustc_hir::HirId, rustc_hir::Upvar>;
             [] object_safety_violations: rustc_middle::traits::ObjectSafetyViolation, rustc_middle::traits::ObjectSafetyViolation;
             [] codegen_unit: rustc_middle::mir::mono::CodegenUnit<$tcx>, rustc_middle::mir::mono::CodegenUnit<'_x>;
+            [] coverage_data: rustc_middle::mir::CoverageData, rustc_middle::mir::CoverageData;
             [] attribute: rustc_ast::ast::Attribute, rustc_ast::ast::Attribute;
             [] name_set: rustc_data_structures::fx::FxHashSet<rustc_span::symbol::Symbol>, rustc_data_structures::fx::FxHashSet<rustc_span::symbol::Symbol>;
             [] hir_id_set: rustc_hir::HirIdSet, rustc_hir::HirIdSet;

src/librustc_middle/mir/mod.rs

@@ -2994,6 +2994,22 @@ impl Location {
     }
 }
 
+///////////////////////////////////////////////////////////////////////////
+// Source code coverage
+
+/// Coverage data associated with each function (MIR) instrumented with coverage counters, when
+/// compiled with `-Zinstrument_coverage`. The query `tcx.coverage_data(DefId)` computes these
+/// values on demand (during code generation). This query is only valid after executing the MIR pass
+/// `InstrumentCoverage`.
+#[derive(Clone, RustcEncodable, RustcDecodable, Debug, HashStable)]
+pub struct CoverageRegion {
+    /// The code region's starting position within the source code file.
+    pub start_byte_pos: u32,
+
+    /// The code region's ending position within the source code file.
+    pub end_byte_pos: u32,
+}
+
 /// Coverage data associated with each function (MIR) instrumented with coverage counters, when
 /// compiled with `-Zinstrument_coverage`. The query `tcx.coverage_data(DefId)` computes these
 /// values on demand (during code generation). This query is only valid after executing the MIR pass
@@ -3007,4 +3023,14 @@ pub struct CoverageData {
 
     /// The total number of coverage region counters added to the MIR `Body`.
     pub num_counters: u32,
+
+    /// The start and end positions within a source file for the region of source code counted by
+    /// the given counter index.
+    pub coverage_regions: IndexVec<u32, CoverageRegion>,
+}
+
+pub mod count_code_region_args {
+    pub const COUNTER_INDEX: usize = 0;
+    pub const START_BYTE_POS: usize = 1;
+    pub const END_BYTE_POS: usize = 2;
 }

src/librustc_middle/query/mod.rs

@@ -231,7 +231,7 @@ rustc_queries! {
             cache_on_disk_if { key.is_local() }
         }
 
-        query coverage_data(key: DefId) -> mir::CoverageData {
+        query coverage_data(key: DefId) -> &'tcx mir::CoverageData {
             desc { |tcx| "retrieving coverage data from MIR for `{}`", tcx.def_path_str(key) }
             storage(ArenaCacheSelector<'tcx>)
             cache_on_disk_if { key.is_local() }

src/librustc_mir/transform/inline.rs

@@ -39,7 +39,14 @@ struct CallSite<'tcx> {
 impl<'tcx> MirPass<'tcx> for Inline {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
         if tcx.sess.opts.debugging_opts.mir_opt_level >= 2 {
-            Inliner { tcx, source }.run_pass(body);
+            if tcx.sess.opts.debugging_opts.instrument_coverage {
+                debug!("function inlining is disabled when compiling with `instrument_coverage`");
+            // The current implementation of source code coverage injects code region counters
+            // into the MIR, and assumes a 1-to-1 correspondence between MIR and source-code-
+            // based function.
+            } else {
+                Inliner { tcx, source }.run_pass(body);
+            }
         }
     }
 }

src/librustc_mir/transform/instrument_coverage.rs

@@ -3,19 +3,22 @@ use crate::util::patch::MirPatch;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_hir::lang_items;
+use rustc_index::vec::IndexVec;
 use rustc_middle::hir;
 use rustc_middle::ich::StableHashingContext;
-use rustc_middle::mir::interpret::{ConstValue, Scalar};
+use rustc_middle::mir::interpret::{ConstValue, InterpResult, Scalar};
 use rustc_middle::mir::{
-    self, traversal, BasicBlock, BasicBlockData, CoverageData, Operand, Place, SourceInfo,
-    StatementKind, Terminator, TerminatorKind, START_BLOCK,
+    self, traversal, BasicBlock, BasicBlockData, CoverageData, CoverageRegion, Operand, Place,
+    SourceInfo, StatementKind, Terminator, TerminatorKind, START_BLOCK,
 };
 use rustc_middle::ty;
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::TyCtxt;
 use rustc_middle::ty::{ConstKind, FnDef};
 use rustc_span::def_id::DefId;
-use rustc_span::Span;
+use rustc_span::{Pos, Span};
+
+use rustc_middle::mir::count_code_region_args::{COUNTER_INDEX, END_BYTE_POS, START_BYTE_POS};
 
 /// Inserts call to count_code_region() as a placeholder to be replaced during code generation with
 /// the intrinsic llvm.instrprof.increment.
@@ -24,46 +27,96 @@ pub struct InstrumentCoverage;
 /// The `query` provider for `CoverageData`, requested by `codegen_intrinsic_call()` when
 /// constructing the arguments for `llvm.instrprof.increment`.
 pub(crate) fn provide(providers: &mut Providers<'_>) {
-    providers.coverage_data = |tcx, def_id| {
-        let mir_body = tcx.optimized_mir(def_id);
-        // FIXME(richkadel): The current implementation assumes the MIR for the given DefId
-        // represents a single function. Validate and/or correct if inlining and/or monomorphization
-        // invalidates these assumptions.
-        let count_code_region_fn =
-            tcx.require_lang_item(lang_items::CountCodeRegionFnLangItem, None);
-        let mut num_counters: u32 = 0;
-        // The `num_counters` argument to `llvm.instrprof.increment` is the number of injected
-        // counters, with each counter having an index from `0..num_counters-1`. MIR optimization
-        // may split and duplicate some BasicBlock sequences. Simply counting the calls may not
-        // not work; but computing the num_counters by adding `1` to the highest index (for a given
-        // instrumented function) is valid.
-        for (_, data) in traversal::preorder(mir_body) {
-            if let Some(terminator) = &data.terminator {
-                if let TerminatorKind::Call { func: Operand::Constant(func), args, .. } =
-                    &terminator.kind
-                {
-                    if let FnDef(called_fn_def_id, _) = func.literal.ty.kind {
-                        if called_fn_def_id == count_code_region_fn {
-                            if let Operand::Constant(constant) =
-                                args.get(0).expect("count_code_region has at least one arg")
-                            {
-                                if let ConstKind::Value(ConstValue::Scalar(value)) =
-                                    constant.literal.val
-                                {
-                                    let index = value
-                                        .to_u32()
-                                        .expect("count_code_region index at arg0 is u32");
-                                    num_counters = std::cmp::max(num_counters, index + 1);
-                                }
-                            }
-                        }
-                    }
+    providers.coverage_data = |tcx, def_id| coverage_data_from_mir(tcx, def_id);
+}
+
+fn coverage_data_from_mir<'tcx>(tcx: TyCtxt<'tcx>, mir_def_id: DefId) -> &'tcx CoverageData {
+    let mir_body = tcx.optimized_mir(mir_def_id);
+    // FIXME(richkadel): The current implementation assumes the MIR for the given DefId
+    // represents a single function. Validate and/or correct if inlining and/or monomorphization
+    // invalidates these assumptions.
+    let count_code_region_fn = tcx.require_lang_item(lang_items::CountCodeRegionFnLangItem, None);
+
+    // The `num_counters` argument to `llvm.instrprof.increment` is the number of injected
+    // counters, with each counter having an index from `0..num_counters-1`. MIR optimization
+    // may split and duplicate some BasicBlock sequences. Simply counting the calls may not
+    // not work; but computing the num_counters by adding `1` to the highest index (for a given
+    // instrumented function) is valid.
+    let mut indexed_regions = vec![];
+    let mut num_counters: u32 = 0;
+    //for terminator in count_code_region_terminators(tcx, mir_body) {
+    for terminator in traversal::preorder(mir_body)
+        .map(|(_, data)| (data, count_code_region_fn))
+        .filter_map(count_code_region_terminator_filter)
+    {
+        // FIXME(richkadel): When analyzing the MIR and identifying CodeRegions to count by LLVM
+        // counter_expressions, inject (and here, filter for) an additional call, such as:
+        //
+        //     code_region_counter_expression(
+        //         left_counter: i32, signed_right_counter: i64,
+        //         start_byte_pos, end_byte_pos
+        //     )
+        //
+        // `abs(signed_right_counter)` (perhaps there's a better name?) is the index of the counter's
+        // count to be added to or subtracted from (depending on the sign) the `left_counter`
+        // index.
+        //
+        // Then store the expressions in a different IndexVec.
+        // We may also want to consider "GapRegions" (clang had them... do we need them too?)
+        if let TerminatorKind::Call { args, .. } = &terminator.kind {
+            let (index, start_byte_pos, end_byte_pos) = count_code_region_args(&args);
+            indexed_regions.push((index, CoverageRegion { start_byte_pos, end_byte_pos }));
+            num_counters = std::cmp::max(num_counters, index + 1);
+        }
+    }
+    let hash = if num_counters > 0 { hash_mir_source(tcx, mir_def_id) } else { 0 };
+    let mut coverage_regions = IndexVec::<u32, CoverageRegion>::from_elem_n(
+        CoverageRegion { start_byte_pos: 0, end_byte_pos: 0 },
+        indexed_regions.len(),
+    );
+    for (index, region) in indexed_regions {
+        coverage_regions[index] = region;
+    }
+    tcx.arena.alloc(CoverageData { num_counters, hash, coverage_regions })
+}
+
+fn count_code_region_terminator_filter(
+    (data, count_code_region_fn): (&'tcx BasicBlockData<'tcx>, DefId),
+) -> Option<&'tcx Terminator<'tcx>> {
+    if let Some(terminator) = &data.terminator {
+        if let TerminatorKind::Call { func: Operand::Constant(func), .. } = &terminator.kind {
+            if let FnDef(called_fn_def_id, _) = func.literal.ty.kind {
+                if called_fn_def_id == count_code_region_fn {
+                    return Some(&terminator);
                 }
             }
         }
-        let hash = if num_counters > 0 { hash_mir_source(tcx, def_id) } else { 0 };
-        CoverageData { num_counters, hash }
-    };
+    }
+    None
+}
+
+fn arg<T, F>(to_val: F, args: &Vec<Operand<'tcx>>, pos: usize) -> T
+where
+    F: FnOnce(Scalar) -> InterpResult<'static, T>,
+{
+    match args.get(pos).unwrap_or_else(|| bug!("count_code_region arg{} not found", pos)) {
+        Operand::Constant(constant) => match constant.literal.val {
+            ConstKind::Value(ConstValue::Scalar(scalar)) => to_val(scalar).unwrap_or_else(|err| {
+                bug!("count_code_region arg{}, {:?}: {:?}", pos, scalar, err);
+            }),
+            _ => bug!("count_code_region arg{}: ConstKind::Value expected", pos),
+        },
+        _ => bug!("count_code_region arg{}: Operand::Constant expected", pos),
+    }
+}
+
+fn count_code_region_args(args: &Vec<Operand<'tcx>>) -> (u32, u32, u32) {
+    let to_u32 = |scalar: Scalar| scalar.to_u32();
+    (
+        arg(to_u32, args, COUNTER_INDEX),
+        arg(to_u32, args, START_BYTE_POS),
+        arg(to_u32, args, END_BYTE_POS),
+    )
 }
 
 struct Instrumentor<'tcx> {
@@ -102,17 +155,16 @@ impl<'tcx> Instrumentor<'tcx> {
     fn inject_counters(&mut self, mir_body: &mut mir::Body<'tcx>) {
         // FIXME(richkadel): As a first step, counters are only injected at the top of each
         // function. The complete solution will inject counters at each conditional code branch.
-        let top_of_function = START_BLOCK;
-        let entire_function = mir_body.span;
-
-        self.inject_counter(mir_body, top_of_function, entire_function);
+        let code_region = mir_body.span;
+        let next_block = START_BLOCK;
+        self.inject_counter(mir_body, code_region, next_block);
     }
 
     fn inject_counter(
         &mut self,
         mir_body: &mut mir::Body<'tcx>,
-        next_block: BasicBlock,
         code_region: Span,
+        next_block: BasicBlock,
     ) {
         let injection_point = code_region.shrink_to_lo();
 
@@ -121,12 +173,19 @@ impl<'tcx> Instrumentor<'tcx> {
             self.tcx.require_lang_item(lang_items::CountCodeRegionFnLangItem, None),
             injection_point,
         );
-        let counter_index = Operand::const_from_scalar(
-            self.tcx,
-            self.tcx.types.u32,
-            Scalar::from_u32(self.next_counter()),
-            injection_point,
-        );
+
+        let index = self.next_counter();
+
+        let mut args = Vec::new();
+
+        assert_eq!(COUNTER_INDEX, args.len());
+        args.push(self.const_u32(index, injection_point));
+
+        assert_eq!(START_BYTE_POS, args.len());
+        args.push(self.const_u32(code_region.lo().to_u32(), injection_point));
+
+        assert_eq!(END_BYTE_POS, args.len());
+        args.push(self.const_u32(code_region.hi().to_u32(), injection_point));
 
         let mut patch = MirPatch::new(mir_body);
 
@@ -136,7 +195,7 @@ impl<'tcx> Instrumentor<'tcx> {
             new_block,
             TerminatorKind::Call {
                 func: count_code_region_fn,
-                args: vec![counter_index],
+                args,
                 // new_block will swapped with the next_block, after applying patch
                 destination: Some((Place::from(temp), new_block)),
                 cleanup: None,
@@ -154,6 +213,10 @@ impl<'tcx> Instrumentor<'tcx> {
         // `next_block`), just swap the indexes, leaving the rest of the graph unchanged.
         mir_body.basic_blocks_mut().swap(next_block, new_block);
     }
+
+    fn const_u32(&self, value: u32, span: Span) -> Operand<'tcx> {
+        Operand::const_from_scalar(self.tcx, self.tcx.types.u32, Scalar::from_u32(value), span)
+    }
 }
 
 fn function_handle<'tcx>(tcx: TyCtxt<'tcx>, fn_def_id: DefId, span: Span) -> Operand<'tcx> {

src/librustc_typeck/check/intrinsic.rs

@@ -352,7 +352,9 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
                 return;
             }
 
-            "count_code_region" => (0, vec![tcx.types.u32], tcx.mk_unit()),
+            "count_code_region" => {
+                (0, vec![tcx.types.u32, tcx.types.u32, tcx.types.u32], tcx.mk_unit())
+            }
 
             ref other => {
                 struct_span_err!(