Rollup merge of rust-lang#73011 - richkadel:llvm-count-from-mir-pass, r=tmandry

first stage of implementing LLVM code coverage

This PR replaces rust-lang#70680 (WIP toward LLVM Code Coverage for Rust) since I am re-implementing the Rust LLVM code coverage feature in a different part of the compiler (in MIR pass(es) vs AST).

This PR updates rustc with `-Zinstrument-coverage` option that injects the llvm intrinsic `instrprof.increment()` for code generation.

This initial version only injects counters at the top of each function, and does not yet implement the required coverage map.

Upcoming PRs will add the coverage map, and add more counters and/or counter expressions for each conditional code branch.

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

***[I put together some development notes here, under a separate branch.](https://github.com/richkadel/rust/blob/cfa0b21d34ee64e4ebee226101bd2ef0c6757865/src/test/codegen/coverage-experiments/README-THIS-IS-TEMPORARY.md)***

Author: Dylan-DPC

src/libcore/intrinsics.rs

@@ -1941,6 +1941,13 @@ extern "rust-intrinsic" {
     ///
     /// Perma-unstable: do not use.
     pub fn miri_start_panic(payload: *mut u8) -> !;
+
+    /// Internal placeholder for injecting code coverage counters when the "instrument-coverage"
+    /// option is enabled. The placeholder is replaced with `llvm.instrprof.increment` during code
+    /// generation.
+    #[cfg(not(bootstrap))]
+    #[lang = "count_code_region"]
+    pub fn count_code_region(_index: u32);
 }
 
 // Some functions are defined here because they accidentally got made

src/librustc_codegen_llvm/builder.rs

@@ -997,6 +997,33 @@ impl BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         self.call_lifetime_intrinsic("llvm.lifetime.end.p0i8", ptr, size);
     }
 
+    fn instrprof_increment(
+        &mut self,
+        fn_name: &'ll Value,
+        hash: &'ll Value,
+        num_counters: &'ll Value,
+        index: &'ll Value,
+    ) -> &'ll Value {
+        debug!(
+            "instrprof_increment() with args ({:?}, {:?}, {:?}, {:?})",
+            fn_name, hash, num_counters, index
+        );
+
+        let llfn = unsafe { llvm::LLVMRustGetInstrprofIncrementIntrinsic(self.cx().llmod) };
+        let args = &[fn_name, hash, num_counters, index];
+        let args = self.check_call("call", llfn, args);
+
+        unsafe {
+            llvm::LLVMRustBuildCall(
+                self.llbuilder,
+                llfn,
+                args.as_ptr() as *const &llvm::Value,
+                args.len() as c_uint,
+                None,
+            )
+        }
+    }
+
     fn call(
         &mut self,
         llfn: &'ll Value,

src/librustc_codegen_llvm/context.rs

@@ -749,6 +749,8 @@ 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);

src/librustc_codegen_llvm/intrinsic.rs

@@ -7,6 +7,8 @@ use crate::type_of::LayoutLlvmExt;
 use crate::va_arg::emit_va_arg;
 use crate::value::Value;
 
+use log::debug;
+
 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;
@@ -21,6 +23,7 @@ use rustc_middle::ty::layout::{FnAbiExt, HasTyCtxt};
 use rustc_middle::ty::{self, Ty};
 use rustc_middle::{bug, span_bug};
 use rustc_span::Span;
+use rustc_span::Symbol;
 use rustc_target::abi::{self, HasDataLayout, LayoutOf, Primitive};
 use rustc_target::spec::PanicStrategy;
 
@@ -86,6 +89,7 @@ impl IntrinsicCallMethods<'tcx> for Builder<'a, 'll, 'tcx> {
         args: &[OperandRef<'tcx, &'ll Value>],
         llresult: &'ll Value,
         span: Span,
+        caller_instance: ty::Instance<'tcx>,
     ) {
         let tcx = self.tcx;
         let callee_ty = instance.monomorphic_ty(tcx);
@@ -136,6 +140,28 @@ impl IntrinsicCallMethods<'tcx> for Builder<'a, 'll, 'tcx> {
                 let llfn = self.get_intrinsic(&("llvm.debugtrap"));
                 self.call(llfn, &[], None)
             }
+            "count_code_region" => {
+                if let ty::InstanceDef::Item(fn_def_id) = caller_instance.def {
+                    let caller_fn_path = tcx.def_path_str(fn_def_id);
+                    debug!(
+                        "count_code_region to llvm.instrprof.increment(fn_name={})",
+                        caller_fn_path
+                    );
+
+                    // FIXME(richkadel): (1) Replace raw function name with mangled function name;
+                    // (2) Replace hardcoded `1234` in `hash` with a computed hash (as discussed in)
+                    // the MCP (compiler-team/issues/278); and replace the hardcoded `1` for
+                    // `num_counters` with the actual number of counters per function (when the
+                    // changes are made to inject more than one counter per function).
+                    let (fn_name, _len_val) = self.const_str(Symbol::intern(&caller_fn_path));
+                    let index = args[0].immediate();
+                    let hash = self.const_u64(1234);
+                    let num_counters = self.const_u32(1);
+                    self.instrprof_increment(fn_name, hash, num_counters, index)
+                } else {
+                    bug!("intrinsic count_code_region: no src.instance");
+                }
+            }
             "va_start" => self.va_start(args[0].immediate()),
             "va_end" => self.va_end(args[0].immediate()),
             "va_copy" => {

src/librustc_codegen_llvm/llvm/ffi.rs

@@ -1360,6 +1360,7 @@ extern "C" {
 
     // Miscellaneous instructions
     pub fn LLVMBuildPhi(B: &Builder<'a>, Ty: &'a Type, Name: *const c_char) -> &'a Value;
+    pub fn LLVMRustGetInstrprofIncrementIntrinsic(M: &Module) -> &'a Value;
     pub fn LLVMRustBuildCall(
         B: &Builder<'a>,
         Fn: &'a Value,

src/librustc_codegen_ssa/back/write.rs

@@ -175,6 +175,12 @@ impl ModuleConfig {
                     if sess.opts.debugging_opts.profile && !is_compiler_builtins {
                         passes.push("insert-gcov-profiling".to_owned());
                     }
+
+                    // The rustc option `-Zinstrument_coverage` injects intrinsic calls to
+                    // `llvm.instrprof.increment()`, which requires the LLVM `instrprof` pass.
+                    if sess.opts.debugging_opts.instrument_coverage {
+                        passes.push("instrprof".to_owned());
+                    }
                     passes
                 },
                 vec![]

src/librustc_codegen_ssa/mir/block.rs

@@ -693,6 +693,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 &args,
                 dest,
                 terminator.source_info.span,
+                self.instance,
             );
 
             if let ReturnDest::IndirectOperand(dst, _) = ret_dest {

src/librustc_codegen_ssa/traits/builder.rs

@@ -260,6 +260,14 @@ pub trait BuilderMethods<'a, 'tcx>:
     /// Called for `StorageDead`
     fn lifetime_end(&mut self, ptr: Self::Value, size: Size);
 
+    fn instrprof_increment(
+        &mut self,
+        fn_name: Self::Value,
+        hash: Self::Value,
+        num_counters: Self::Value,
+        index: Self::Value,
+    ) -> Self::Value;
+
     fn call(
         &mut self,
         llfn: Self::Value,

src/librustc_codegen_ssa/traits/intrinsic.rs

@@ -15,6 +15,7 @@ pub trait IntrinsicCallMethods<'tcx>: BackendTypes {
         args: &[OperandRef<'tcx, Self::Value>],
         llresult: Self::Value,
         span: Span,
+        caller_instance: ty::Instance<'tcx>,
     );
 
     fn abort(&mut self);

src/librustc_hir/lang_items.rs

@@ -242,6 +242,8 @@ language_item_table! {
 
     StartFnLangItem,             "start",              start_fn,                Target::Fn;
 
+    CountCodeRegionFnLangItem,   "count_code_region",  count_code_region_fn,    Target::Fn;
+
     EhPersonalityLangItem,       "eh_personality",     eh_personality,          Target::Fn;
     EhCatchTypeinfoLangItem,     "eh_catch_typeinfo",  eh_catch_typeinfo,       Target::Static;
 

src/librustc_interface/tests.rs

@@ -548,6 +548,7 @@ fn test_debugging_options_tracking_hash() {
     tracked!(human_readable_cgu_names, true);
     tracked!(inline_in_all_cgus, Some(true));
     tracked!(insert_sideeffect, true);
+    tracked!(instrument_coverage, true);
     tracked!(instrument_mcount, true);
     tracked!(link_only, true);
     tracked!(merge_functions, Some(MergeFunctions::Disabled));

src/librustc_middle/mir/mod.rs

@@ -29,6 +29,7 @@ use rustc_macros::HashStable;
 use rustc_serialize::{Decodable, Encodable};
 use rustc_span::symbol::Symbol;
 use rustc_span::{Span, DUMMY_SP};
+use rustc_target::abi;
 use rustc_target::asm::InlineAsmRegOrRegClass;
 use std::borrow::Cow;
 use std::fmt::{self, Debug, Display, Formatter, Write};
@@ -2218,6 +2219,33 @@ impl<'tcx> Operand<'tcx> {
         })
     }
 
+    /// Convenience helper to make a literal-like constant from a given scalar value.
+    /// Since this is used to synthesize MIR, assumes `user_ty` is None.
+    pub fn const_from_scalar(
+        tcx: TyCtxt<'tcx>,
+        ty: Ty<'tcx>,
+        val: Scalar,
+        span: Span,
+    ) -> Operand<'tcx> {
+        debug_assert!({
+            let param_env_and_ty = ty::ParamEnv::empty().and(ty);
+            let type_size = tcx
+                .layout_of(param_env_and_ty)
+                .unwrap_or_else(|e| panic!("could not compute layout for {:?}: {:?}", ty, e))
+                .size;
+            let scalar_size = abi::Size::from_bytes(match val {
+                Scalar::Raw { size, .. } => size,
+                _ => panic!("Invalid scalar type {:?}", val),
+            });
+            scalar_size == type_size
+        });
+        Operand::Constant(box Constant {
+            span,
+            user_ty: None,
+            literal: ty::Const::from_scalar(tcx, val, ty),
+        })
+    }
+
     pub fn to_copy(&self) -> Self {
         match *self {
             Operand::Copy(_) | Operand::Constant(_) => self.clone(),

src/librustc_mir/interpret/intrinsics.rs

@@ -389,6 +389,8 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 );
                 self.copy_op(self.operand_index(args[0], index)?, dest)?;
             }
+            // FIXME(#73156): Handle source code coverage in const eval
+            sym::count_code_region => (),
             _ => return Ok(false),
         }
 

src/librustc_mir/transform/instrument_coverage.rs

@@ -0,0 +1,91 @@
+use crate::transform::{MirPass, MirSource};
+use crate::util::patch::MirPatch;
+use rustc_hir::lang_items;
+use rustc_middle::mir::interpret::Scalar;
+use rustc_middle::mir::*;
+use rustc_middle::ty;
+use rustc_middle::ty::TyCtxt;
+use rustc_span::def_id::DefId;
+use rustc_span::Span;
+
+/// Inserts call to count_code_region() as a placeholder to be replaced during code generation with
+/// the intrinsic llvm.instrprof.increment.
+pub struct InstrumentCoverage;
+
+impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut Body<'tcx>) {
+        if tcx.sess.opts.debugging_opts.instrument_coverage {
+            debug!("instrumenting {:?}", src.def_id());
+            instrument_coverage(tcx, body);
+        }
+    }
+}
+
+// The first counter (start of the function) is index zero.
+const INIT_FUNCTION_COUNTER: u32 = 0;
+
+/// Injects calls to placeholder function `count_code_region()`.
+// 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.
+pub fn instrument_coverage<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+    let span = body.span.shrink_to_lo();
+
+    let count_code_region_fn = function_handle(
+        tcx,
+        tcx.require_lang_item(lang_items::CountCodeRegionFnLangItem, None),
+        span,
+    );
+    let counter_index = Operand::const_from_scalar(
+        tcx,
+        tcx.types.u32,
+        Scalar::from_u32(INIT_FUNCTION_COUNTER),
+        span,
+    );
+
+    let mut patch = MirPatch::new(body);
+
+    let new_block = patch.new_block(placeholder_block(SourceInfo::outermost(body.span)));
+    let next_block = START_BLOCK;
+
+    let temp = patch.new_temp(tcx.mk_unit(), body.span);
+    patch.patch_terminator(
+        new_block,
+        TerminatorKind::Call {
+            func: count_code_region_fn,
+            args: vec![counter_index],
+            // new_block will swapped with the next_block, after applying patch
+            destination: Some((Place::from(temp), new_block)),
+            cleanup: None,
+            from_hir_call: false,
+        },
+    );
+
+    patch.add_statement(new_block.start_location(), StatementKind::StorageLive(temp));
+    patch.add_statement(next_block.start_location(), StatementKind::StorageDead(temp));
+
+    patch.apply(body);
+
+    // To insert the `new_block` in front of the first block in the counted branch (for example,
+    // the START_BLOCK, at the top of the function), just swap the indexes, leaving the rest of the
+    // graph unchanged.
+    body.basic_blocks_mut().swap(next_block, new_block);
+}
+
+fn function_handle<'tcx>(tcx: TyCtxt<'tcx>, fn_def_id: DefId, span: Span) -> Operand<'tcx> {
+    let ret_ty = tcx.fn_sig(fn_def_id).output();
+    let ret_ty = ret_ty.no_bound_vars().unwrap();
+    let substs = tcx.mk_substs(::std::iter::once(ty::subst::GenericArg::from(ret_ty)));
+    Operand::function_handle(tcx, fn_def_id, substs, span)
+}
+
+fn placeholder_block<'tcx>(source_info: SourceInfo) -> BasicBlockData<'tcx> {
+    BasicBlockData {
+        statements: vec![],
+        terminator: Some(Terminator {
+            source_info,
+            // this gets overwritten by the counter Call
+            kind: TerminatorKind::Unreachable,
+        }),
+        is_cleanup: false,
+    }
+}

src/librustc_mir/transform/mod.rs

@@ -28,6 +28,7 @@ pub mod elaborate_drops;
 pub mod generator;
 pub mod inline;
 pub mod instcombine;
+pub mod instrument_coverage;
 pub mod no_landing_pads;
 pub mod nrvo;
 pub mod promote_consts;
@@ -288,6 +289,10 @@ fn mir_validated(
             // What we need to run borrowck etc.
             &promote_pass,
             &simplify::SimplifyCfg::new("qualify-consts"),
+            // If the `instrument-coverage` option is enabled, analyze the CFG, identify each
+            // conditional branch, construct a coverage map to be passed to LLVM, and inject counters
+            // where needed.
+            &instrument_coverage::InstrumentCoverage,
         ]],
     );
 

src/librustc_passes/weak_lang_items.rs

@@ -5,10 +5,12 @@ use rustc_errors::struct_span_err;
 use rustc_hir as hir;
 use rustc_hir::intravisit::{self, NestedVisitorMap, Visitor};
 use rustc_hir::lang_items;
+use rustc_hir::lang_items::ITEM_REFS;
 use rustc_hir::weak_lang_items::WEAK_ITEMS_REFS;
 use rustc_middle::middle::lang_items::whitelisted;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::config::CrateType;
+use rustc_span::symbol::sym;
 use rustc_span::symbol::Symbol;
 use rustc_span::Span;
 
@@ -70,11 +72,21 @@ fn verify<'tcx>(tcx: TyCtxt<'tcx>, items: &lang_items::LanguageItems) {
 }
 
 impl<'a, 'tcx> Context<'a, 'tcx> {
-    fn register(&mut self, name: Symbol, span: Span) {
+    fn register(&mut self, name: Symbol, span: Span, hir_id: hir::HirId) {
         if let Some(&item) = WEAK_ITEMS_REFS.get(&name) {
             if self.items.require(item).is_err() {
                 self.items.missing.push(item);
             }
+        } else if name == sym::count_code_region {
+            // `core::intrinsics::code_count_region()` is (currently) the only `extern` lang item
+            // that is never actually linked. It is not a `weak_lang_item` that can be registered
+            // when used, and should be registered here instead.
+            if let Some((item_index, _)) = ITEM_REFS.get(&*name.as_str()).cloned() {
+                if self.items.items[item_index].is_none() {
+                    let item_def_id = self.tcx.hir().local_def_id(hir_id).to_def_id();
+                    self.items.items[item_index] = Some(item_def_id);
+                }
+            }
         } else {
             struct_span_err!(self.tcx.sess, span, E0264, "unknown external lang item: `{}`", name)
                 .emit();
@@ -91,7 +103,7 @@ impl<'a, 'tcx, 'v> Visitor<'v> for Context<'a, 'tcx> {
 
     fn visit_foreign_item(&mut self, i: &hir::ForeignItem<'_>) {
         if let Some((lang_item, _)) = hir::lang_items::extract(&i.attrs) {
-            self.register(lang_item, i.span);
+            self.register(lang_item, i.span, i.hir_id);
         }
         intravisit::walk_foreign_item(self, i)
     }

src/librustc_session/options.rs

@@ -876,6 +876,9 @@ options! {DebuggingOptions, DebuggingSetter, basic_debugging_options,
         "fix undefined behavior when a thread doesn't eventually make progress \
         (such as entering an empty infinite loop) by inserting llvm.sideeffect \
         (default: no)"),
+    instrument_coverage: bool = (false, parse_bool, [TRACKED],
+        "instrument the generated code with LLVM code region counters to \
+        (in the future) generate coverage reports (experimental; default: no)"),
     instrument_mcount: bool = (false, parse_bool, [TRACKED],
         "insert function instrument code for mcount-based tracing (default: no)"),
     keep_hygiene_data: bool = (false, parse_bool, [UNTRACKED],