1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
|
use crate::transform::{MirPass, MirSource};
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_middle::hir;
use rustc_middle::ich::StableHashingContext;
use rustc_middle::mir::coverage::*;
use rustc_middle::mir::interpret::Scalar;
use rustc_middle::mir::CoverageInfo;
use rustc_middle::mir::{
self, traversal, BasicBlock, BasicBlockData, Operand, Place, SourceInfo, StatementKind,
Terminator, TerminatorKind, START_BLOCK,
};
use rustc_middle::ty;
use rustc_middle::ty::query::Providers;
use rustc_middle::ty::FnDef;
use rustc_middle::ty::TyCtxt;
use rustc_span::def_id::DefId;
use rustc_span::{Pos, 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;
/// The `query` provider for `CoverageInfo`, requested by `codegen_intrinsic_call()` when
/// constructing the arguments for `llvm.instrprof.increment`.
pub(crate) fn provide(providers: &mut Providers) {
providers.coverageinfo = |tcx, def_id| coverageinfo_from_mir(tcx, def_id);
}
fn coverageinfo_from_mir<'tcx>(tcx: TyCtxt<'tcx>, mir_def_id: DefId) -> CoverageInfo {
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 (which should be disabled
// if -Zinstrument-coverage is enabled) 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 num_counters: u32 = 0;
for terminator in traversal::preorder(mir_body)
.map(|(_, data)| (data, count_code_region_fn))
.filter_map(terminators_that_call_given_fn)
{
if let TerminatorKind::Call { args, .. } = &terminator.kind {
let index_arg = args.get(count_code_region_args::COUNTER_INDEX).expect("arg found");
let index =
mir::Operand::scalar_from_const(index_arg).to_u32().expect("index arg is u32");
num_counters = std::cmp::max(num_counters, index + 1);
}
}
let hash = if num_counters > 0 { hash_mir_source(tcx, mir_def_id) } else { 0 };
CoverageInfo { num_counters, hash }
}
fn terminators_that_call_given_fn(
(data, fn_def_id): (&'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 == fn_def_id {
return Some(&terminator);
}
}
}
}
None
}
struct Instrumentor<'tcx> {
tcx: TyCtxt<'tcx>,
num_counters: u32,
}
impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, mir_body: &mut mir::Body<'tcx>) {
if tcx.sess.opts.debugging_opts.instrument_coverage {
// If the InstrumentCoverage pass is called on promoted MIRs, skip them.
// See: https://github.com/rust-lang/rust/pull/73011#discussion_r438317601
if src.promoted.is_none() {
debug!(
"instrumenting {:?}, span: {}",
src.def_id(),
tcx.sess.source_map().span_to_string(mir_body.span)
);
Instrumentor::new(tcx).inject_counters(mir_body);
}
}
}
}
impl<'tcx> Instrumentor<'tcx> {
fn new(tcx: TyCtxt<'tcx>) -> Self {
Self { tcx, num_counters: 0 }
}
fn next_counter(&mut self) -> u32 {
let next = self.num_counters;
self.num_counters += 1;
next
}
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 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>,
code_region: Span,
next_block: BasicBlock,
) {
let injection_point = code_region.shrink_to_lo();
let count_code_region_fn = function_handle(
self.tcx,
self.tcx.require_lang_item(lang_items::CountCodeRegionFnLangItem, None),
injection_point,
);
let index = self.next_counter();
let mut args = Vec::new();
use count_code_region_args::*;
debug_assert_eq!(COUNTER_INDEX, args.len());
args.push(self.const_u32(index, injection_point));
debug_assert_eq!(START_BYTE_POS, args.len());
args.push(self.const_u32(code_region.lo().to_u32(), injection_point));
debug_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);
let temp = patch.new_temp(self.tcx.mk_unit(), code_region);
let new_block = patch.new_block(placeholder_block(code_region));
patch.patch_terminator(
new_block,
TerminatorKind::Call {
func: count_code_region_fn,
args,
// 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: injection_point,
},
);
patch.add_statement(new_block.start_location(), StatementKind::StorageLive(temp));
patch.add_statement(next_block.start_location(), StatementKind::StorageDead(temp));
patch.apply(mir_body);
// To insert the `new_block` in front of the first block in the counted branch (the
// `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> {
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(span: Span) -> BasicBlockData<'tcx> {
BasicBlockData {
statements: vec![],
terminator: Some(Terminator {
source_info: SourceInfo::outermost(span),
// this gets overwritten by the counter Call
kind: TerminatorKind::Unreachable,
}),
is_cleanup: false,
}
}
fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> u64 {
let hir_node = tcx.hir().get_if_local(def_id).expect("DefId is local");
let fn_body_id = hir::map::associated_body(hir_node).expect("HIR node is a function with body");
let hir_body = tcx.hir().body(fn_body_id);
let mut hcx = tcx.create_no_span_stable_hashing_context();
hash(&mut hcx, &hir_body.value).to_smaller_hash()
}
fn hash(
hcx: &mut StableHashingContext<'tcx>,
node: &impl HashStable<StableHashingContext<'tcx>>,
) -> Fingerprint {
let mut stable_hasher = StableHasher::new();
node.hash_stable(hcx, &mut stable_hasher);
stable_hasher.finish()
}
|