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
|
#![feature(rustc_private)]
//! This program implements a rustc driver that retrieves MIR bodies with
//! borrowck information. This cannot be done in a straightforward way because
//! `get_body_with_borrowck_facts`–the function for retrieving a MIR body with
//! borrowck facts–can panic if the body is stolen before it is invoked.
//! Therefore, the driver overrides `mir_borrowck` query (this is done in the
//! `config` callback), which retrieves the body that is about to be borrow
//! checked and stores it in a thread local `MIR_BODIES`. Then, `after_analysis`
//! callback triggers borrow checking of all MIR bodies by retrieving
//! `optimized_mir` and pulls out the MIR bodies with the borrowck information
//! from the thread local storage.
extern crate rustc_borrowck;
extern crate rustc_driver;
extern crate rustc_hir;
extern crate rustc_interface;
extern crate rustc_middle;
extern crate rustc_session;
use rustc_borrowck::consumers::{self, BodyWithBorrowckFacts, ConsumerOptions};
use rustc_driver::Compilation;
use rustc_hir::def::DefKind;
use rustc_hir::def_id::LocalDefId;
use rustc_interface::interface::Compiler;
use rustc_interface::{Config, Queries};
use rustc_middle::query::queries::mir_borrowck::ProvidedValue;
use rustc_middle::ty::TyCtxt;
use rustc_middle::util::Providers;
use rustc_session::Session;
use std::cell::RefCell;
use std::collections::HashMap;
use std::thread_local;
fn main() {
let exit_code = rustc_driver::catch_with_exit_code(move || {
let mut rustc_args: Vec<_> = std::env::args().collect();
// We must pass -Zpolonius so that the borrowck information is computed.
rustc_args.push("-Zpolonius".to_owned());
let mut callbacks = CompilerCalls::default();
// Call the Rust compiler with our callbacks.
rustc_driver::RunCompiler::new(&rustc_args, &mut callbacks).run()
});
std::process::exit(exit_code);
}
#[derive(Default)]
pub struct CompilerCalls;
impl rustc_driver::Callbacks for CompilerCalls {
// In this callback we override the mir_borrowck query.
fn config(&mut self, config: &mut Config) {
assert!(config.override_queries.is_none());
config.override_queries = Some(override_queries);
}
// In this callback we trigger borrow checking of all functions and obtain
// the result.
fn after_analysis<'tcx>(
&mut self,
compiler: &Compiler,
queries: &'tcx Queries<'tcx>,
) -> Compilation {
compiler.sess.abort_if_errors();
queries.global_ctxt().unwrap().enter(|tcx| {
// Collect definition ids of MIR bodies.
let hir = tcx.hir();
let mut bodies = Vec::new();
let crate_items = tcx.hir_crate_items(());
for id in crate_items.items() {
if matches!(tcx.def_kind(id.owner_id), DefKind::Fn) {
bodies.push(id.owner_id);
}
}
for id in crate_items.trait_items() {
if matches!(tcx.def_kind(id.owner_id), DefKind::AssocFn) {
let trait_item = hir.trait_item(id);
if let rustc_hir::TraitItemKind::Fn(_, trait_fn) = &trait_item.kind {
if let rustc_hir::TraitFn::Provided(_) = trait_fn {
bodies.push(trait_item.owner_id);
}
}
}
}
for id in crate_items.impl_items() {
if matches!(tcx.def_kind(id.owner_id), DefKind::AssocFn) {
bodies.push(id.owner_id);
}
}
// Trigger borrow checking of all bodies.
for def_id in bodies {
let _ = tcx.optimized_mir(def_id);
}
// See what bodies were borrow checked.
let mut bodies = get_bodies(tcx);
bodies.sort_by(|(def_id1, _), (def_id2, _)| def_id1.cmp(def_id2));
println!("Bodies retrieved for:");
for (def_id, body) in bodies {
println!("{}", def_id);
assert!(body.input_facts.unwrap().cfg_edge.len() > 0);
}
});
Compilation::Continue
}
}
fn override_queries(_session: &Session, local: &mut Providers) {
local.mir_borrowck = mir_borrowck;
}
// Since mir_borrowck does not have access to any other state, we need to use a
// thread-local for storing the obtained MIR bodies.
//
// Note: We are using 'static lifetime here, which is in general unsound.
// Unfortunately, that is the only lifetime allowed here. Our use is safe
// because we cast it back to `'tcx` before using.
thread_local! {
pub static MIR_BODIES:
RefCell<HashMap<LocalDefId, BodyWithBorrowckFacts<'static>>> =
RefCell::new(HashMap::new());
}
fn mir_borrowck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ProvidedValue<'tcx> {
let opts = ConsumerOptions::PoloniusInputFacts;
let body_with_facts = consumers::get_body_with_borrowck_facts(tcx, def_id, opts);
// SAFETY: The reader casts the 'static lifetime to 'tcx before using it.
let body_with_facts: BodyWithBorrowckFacts<'static> =
unsafe { std::mem::transmute(body_with_facts) };
MIR_BODIES.with(|state| {
let mut map = state.borrow_mut();
assert!(map.insert(def_id, body_with_facts).is_none());
});
let mut providers = Providers::default();
rustc_borrowck::provide(&mut providers);
let original_mir_borrowck = providers.mir_borrowck;
original_mir_borrowck(tcx, def_id)
}
/// Pull MIR bodies stored in the thread-local.
fn get_bodies<'tcx>(tcx: TyCtxt<'tcx>) -> Vec<(String, BodyWithBorrowckFacts<'tcx>)> {
MIR_BODIES.with(|state| {
let mut map = state.borrow_mut();
map.drain()
.map(|(def_id, body)| {
let def_path = tcx.def_path(def_id.to_def_id());
// SAFETY: For soundness we need to ensure that the bodies have
// the same lifetime (`'tcx`), which they had before they were
// stored in the thread local.
(def_path.to_string_no_crate_verbose(), unsafe { std::mem::transmute(body) })
})
.collect()
})
}
|
