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

first stage of implementing LLVM code coverage

This PR replaces #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: #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: RalfJung

config.toml.example

@@ -209,7 +209,8 @@
 # Build the sanitizer runtimes
 #sanitizers = false
 
-# Build the profiler runtime
+# Build the profiler runtime (required when compiling with options that depend
+# on this runtime, such as `-C profile-generate` or `-Z instrument-coverage`).
 #profiler = false
 
 # Indicates whether the native libraries linked into Cargo will be statically

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_metadata/creader.rs

@@ -706,7 +706,9 @@ impl<'a> CrateLoader<'a> {
     }
 
     fn inject_profiler_runtime(&mut self) {
-        if (self.sess.opts.debugging_opts.profile || self.sess.opts.cg.profile_generate.enabled())
+        if (self.sess.opts.debugging_opts.instrument_coverage
+            || self.sess.opts.debugging_opts.profile
+            || self.sess.opts.cg.profile_generate.enabled())
             && !self.sess.opts.debugging_opts.no_profiler_runtime
         {
             info!("loading profiler");

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,92 @@
+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,
+            fn_span: span,
+        },
+    );
+
+    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,
         ]],
     );