Skip to content

Commit

Permalink
Rollup merge of #118929 - Zalathar:look-hir, r=cjgillot
Browse files Browse the repository at this point in the history 
coverage: Tidy up early parts of the instrumentor pass

This is extracted from #118237, which needed to be manually rebased anyway.

Unlike that PR, this one only affects the coverage instrumentor, and doesn't attempt to move any code into the MIR builder. That can be left to a future version of #118305, which can still benefit from these improvements.

So this is now mostly a refactoring of some internal parts of the instrumentor.
  • Loading branch information
@matthiaskrgr
matthiaskrgr authored Dec 15, 2023
2 parents d5af762 + 684b9ea commit 6659b5e
Showing 1 changed file with 77 additions and 70 deletions.
147 changes: 77 additions & 70 deletions compiler/rustc_mir_transform/src/coverage/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ use self::spans::CoverageSpans;

use crate::MirPass;

use rustc_data_structures::sync::Lrc;
use rustc_middle::hir;
use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
use rustc_middle::mir::coverage::*;
Expand All @@ -22,9 +21,9 @@ use rustc_middle::mir::{
TerminatorKind,
};
use rustc_middle::ty::TyCtxt;
use rustc_span::def_id::DefId;
use rustc_span::def_id::LocalDefId;
use rustc_span::source_map::SourceMap;
use rustc_span::{ExpnKind, SourceFile, Span, Symbol};
use rustc_span::{ExpnKind, Span, Symbol};

/// Inserts `StatementKind::Coverage` statements that either instrument the binary with injected
/// counters, via intrinsic `llvm.instrprof.increment`, and/or inject metadata used during codegen
Expand All @@ -39,31 +38,19 @@ impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
fn run_pass(&self, tcx: TyCtxt<'tcx>, mir_body: &mut mir::Body<'tcx>) {
let mir_source = mir_body.source;

// If the InstrumentCoverage pass is called on promoted MIRs, skip them.
// See: https://github.com/rust-lang/rust/pull/73011#discussion_r438317601
if mir_source.promoted.is_some() {
trace!(
"InstrumentCoverage skipped for {:?} (already promoted for Miri evaluation)",
mir_source.def_id()
);
return;
}
// This pass runs after MIR promotion, but before promoted MIR starts to
// be transformed, so it should never see promoted MIR.
assert!(mir_source.promoted.is_none());

let def_id = mir_source.def_id().expect_local();

let is_fn_like =
tcx.hir_node_by_def_id(mir_source.def_id().expect_local()).fn_kind().is_some();

// Only instrument functions, methods, and closures (not constants since they are evaluated
// at compile time by Miri).
// FIXME(#73156): Handle source code coverage in const eval, but note, if and when const
// expressions get coverage spans, we will probably have to "carve out" space for const
// expressions from coverage spans in enclosing MIR's, like we do for closures. (That might
// be tricky if const expressions have no corresponding statements in the enclosing MIR.
// Closures are carved out by their initial `Assign` statement.)
if !is_fn_like {
trace!("InstrumentCoverage skipped for {:?} (not an fn-like)", mir_source.def_id());
if !is_eligible_for_coverage(tcx, def_id) {
trace!("InstrumentCoverage skipped for {def_id:?} (not eligible)");
return;
}

// An otherwise-eligible function is still skipped if its start block
// is known to be unreachable.
match mir_body.basic_blocks[mir::START_BLOCK].terminator().kind {
TerminatorKind::Unreachable => {
trace!("InstrumentCoverage skipped for unreachable `START_BLOCK`");
Expand All @@ -72,21 +59,15 @@ impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
_ => {}
}

let codegen_fn_attrs = tcx.codegen_fn_attrs(mir_source.def_id());
if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
return;
}

trace!("InstrumentCoverage starting for {:?}", mir_source.def_id());
trace!("InstrumentCoverage starting for {def_id:?}");
Instrumentor::new(tcx, mir_body).inject_counters();
trace!("InstrumentCoverage done for {:?}", mir_source.def_id());
trace!("InstrumentCoverage done for {def_id:?}");
}
}

struct Instrumentor<'a, 'tcx> {
tcx: TyCtxt<'tcx>,
mir_body: &'a mut mir::Body<'tcx>,
source_file: Lrc<SourceFile>,
fn_sig_span: Span,
body_span: Span,
function_source_hash: u64,
Expand All @@ -96,37 +77,17 @@ struct Instrumentor<'a, 'tcx> {

impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
fn new(tcx: TyCtxt<'tcx>, mir_body: &'a mut mir::Body<'tcx>) -> Self {
let source_map = tcx.sess.source_map();
let def_id = mir_body.source.def_id();
let (some_fn_sig, hir_body) = fn_sig_and_body(tcx, def_id);
let hir_info @ ExtractedHirInfo { function_source_hash, fn_sig_span, body_span } =
extract_hir_info(tcx, mir_body.source.def_id().expect_local());

let body_span = get_body_span(tcx, hir_body, mir_body);
debug!(?hir_info, "instrumenting {:?}", mir_body.source.def_id());

let source_file = source_map.lookup_source_file(body_span.lo());
let fn_sig_span = match some_fn_sig.filter(|fn_sig| {
fn_sig.span.eq_ctxt(body_span)
&& Lrc::ptr_eq(&source_file, &source_map.lookup_source_file(fn_sig.span.lo()))
}) {
Some(fn_sig) => fn_sig.span.with_hi(body_span.lo()),
None => body_span.shrink_to_lo(),
};

debug!(
"instrumenting {}: {:?}, fn sig span: {:?}, body span: {:?}",
if tcx.is_closure(def_id) { "closure" } else { "function" },
def_id,
fn_sig_span,
body_span
);

let function_source_hash = hash_mir_source(tcx, hir_body);
let basic_coverage_blocks = CoverageGraph::from_mir(mir_body);
let coverage_counters = CoverageCounters::new(&basic_coverage_blocks);

Self {
tcx,
mir_body,
source_file,
fn_sig_span,
body_span,
function_source_hash,
Expand All @@ -136,15 +97,12 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
}

fn inject_counters(&'a mut self) {
let fn_sig_span = self.fn_sig_span;
let body_span = self.body_span;

////////////////////////////////////////////////////
// Compute coverage spans from the `CoverageGraph`.
let coverage_spans = CoverageSpans::generate_coverage_spans(
self.mir_body,
fn_sig_span,
body_span,
self.fn_sig_span,
self.body_span,
&self.basic_coverage_blocks,
);

Expand Down Expand Up @@ -177,9 +135,10 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
let source_map = self.tcx.sess.source_map();
let body_span = self.body_span;

let source_file = source_map.lookup_source_file(body_span.lo());
use rustc_session::RemapFileNameExt;
let file_name =
Symbol::intern(&self.source_file.name.for_codegen(self.tcx.sess).to_string_lossy());
Symbol::intern(&source_file.name.for_codegen(self.tcx.sess).to_string_lossy());

let mut mappings = Vec::new();

Expand Down Expand Up @@ -325,27 +284,75 @@ fn make_code_region(
}
}

fn fn_sig_and_body(
tcx: TyCtxt<'_>,
def_id: DefId,
) -> (Option<&rustc_hir::FnSig<'_>>, &rustc_hir::Body<'_>) {
fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
// Only instrument functions, methods, and closures (not constants since they are evaluated
// at compile time by Miri).
// FIXME(#73156): Handle source code coverage in const eval, but note, if and when const
// expressions get coverage spans, we will probably have to "carve out" space for const
// expressions from coverage spans in enclosing MIR's, like we do for closures. (That might
// be tricky if const expressions have no corresponding statements in the enclosing MIR.
// Closures are carved out by their initial `Assign` statement.)
if !tcx.def_kind(def_id).is_fn_like() {
trace!("InstrumentCoverage skipped for {def_id:?} (not an fn-like)");
return false;
}

if tcx.codegen_fn_attrs(def_id).flags.contains(CodegenFnAttrFlags::NO_COVERAGE) {
return false;
}

true
}

/// Function information extracted from HIR by the coverage instrumentor.
#[derive(Debug)]
struct ExtractedHirInfo {
function_source_hash: u64,
fn_sig_span: Span,
body_span: Span,
}

fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHirInfo {
// FIXME(#79625): Consider improving MIR to provide the information needed, to avoid going back
// to HIR for it.
let hir_node = tcx.hir().get_if_local(def_id).expect("expected DefId is local");

let hir_node = tcx.hir_node_by_def_id(def_id);
let (_, fn_body_id) =
hir::map::associated_body(hir_node).expect("HIR node is a function with body");
(hir_node.fn_sig(), tcx.hir().body(fn_body_id))
let hir_body = tcx.hir().body(fn_body_id);

let body_span = get_body_span(tcx, hir_body, def_id);

// The actual signature span is only used if it has the same context and
// filename as the body, and precedes the body.
let maybe_fn_sig_span = hir_node.fn_sig().map(|fn_sig| fn_sig.span);
let fn_sig_span = maybe_fn_sig_span
.filter(|&fn_sig_span| {
let source_map = tcx.sess.source_map();
let file_idx = |span: Span| source_map.lookup_source_file_idx(span.lo());

fn_sig_span.eq_ctxt(body_span)
&& fn_sig_span.hi() <= body_span.lo()
&& file_idx(fn_sig_span) == file_idx(body_span)
})
// If so, extend it to the start of the body span.
.map(|fn_sig_span| fn_sig_span.with_hi(body_span.lo()))
// Otherwise, create a dummy signature span at the start of the body.
.unwrap_or_else(|| body_span.shrink_to_lo());

let function_source_hash = hash_mir_source(tcx, hir_body);

ExtractedHirInfo { function_source_hash, fn_sig_span, body_span }
}

fn get_body_span<'tcx>(
tcx: TyCtxt<'tcx>,
hir_body: &rustc_hir::Body<'tcx>,
mir_body: &mut mir::Body<'tcx>,
def_id: LocalDefId,
) -> Span {
let mut body_span = hir_body.value.span;
let def_id = mir_body.source.def_id();

if tcx.is_closure(def_id) {
if tcx.is_closure(def_id.to_def_id()) {
// If the MIR function is a closure, and if the closure body span
// starts from a macro, but it's content is not in that macro, try
// to find a non-macro callsite, and instrument the spans there
Expand Down

0 comments on commit 6659b5e

Please sign in to comment.