1 #![feature(rustc_private)]
3 //! This program implements a rustc driver that retrieves MIR bodies with
4 //! borrowck information. This cannot be done in a straightforward way because
5 //! `get_body_with_borrowck_facts`–the function for retrieving a MIR body with
6 //! borrowck facts–can panic if the body is stolen before it is invoked.
7 //! Therefore, the driver overrides `mir_borrowck` query (this is done in the
8 //! `config` callback), which retrieves the body that is about to be borrow
9 //! checked and stores it in a thread local `MIR_BODIES`. Then, `after_analysis`
10 //! callback triggers borrow checking of all MIR bodies by retrieving
11 //! `optimized_mir` and pulls out the MIR bodies with the borrowck information
12 //! from the thread local storage.
14 extern crate rustc_borrowck;
15 extern crate rustc_driver;
16 extern crate rustc_hir;
17 extern crate rustc_interface;
18 extern crate rustc_middle;
19 extern crate rustc_session;
21 use rustc_borrowck::consumers::BodyWithBorrowckFacts;
22 use rustc_driver::Compilation;
23 use rustc_hir::def_id::LocalDefId;
24 use rustc_hir::def::DefKind;
25 use rustc_interface::interface::Compiler;
26 use rustc_interface::{Config, Queries};
27 use rustc_middle::ty::query::query_values::mir_borrowck;
28 use rustc_middle::ty::query::{ExternProviders, Providers};
29 use rustc_middle::ty::{self, TyCtxt};
30 use rustc_session::Session;
31 use std::cell::RefCell;
32 use std::collections::HashMap;
33 use std::thread_local;
36 let exit_code = rustc_driver::catch_with_exit_code(move || {
37 let mut rustc_args: Vec<_> = std::env::args().collect();
38 // We must pass -Zpolonius so that the borrowck information is computed.
39 rustc_args.push("-Zpolonius".to_owned());
40 let mut callbacks = CompilerCalls::default();
41 // Call the Rust compiler with our callbacks.
42 rustc_driver::RunCompiler::new(&rustc_args, &mut callbacks).run()
44 std::process::exit(exit_code);
48 pub struct CompilerCalls;
50 impl rustc_driver::Callbacks for CompilerCalls {
51 // In this callback we override the mir_borrowck query.
52 fn config(&mut self, config: &mut Config) {
53 assert!(config.override_queries.is_none());
54 config.override_queries = Some(override_queries);
57 // In this callback we trigger borrow checking of all functions and obtain
59 fn after_analysis<'tcx>(
62 queries: &'tcx Queries<'tcx>,
64 compiler.session().abort_if_errors();
65 queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
66 // Collect definition ids of MIR bodies.
68 let mut bodies = Vec::new();
70 let crate_items = tcx.hir_crate_items(());
71 for id in crate_items.items() {
72 if matches!(tcx.def_kind(id.owner_id), DefKind::Fn) {
73 bodies.push(id.owner_id);
77 for id in crate_items.trait_items() {
78 if matches!(tcx.def_kind(id.owner_id), DefKind::AssocFn) {
79 let trait_item = hir.trait_item(id);
80 if let rustc_hir::TraitItemKind::Fn(_, trait_fn) = &trait_item.kind {
81 if let rustc_hir::TraitFn::Provided(_) = trait_fn {
82 bodies.push(trait_item.owner_id);
88 for id in crate_items.impl_items() {
89 if matches!(tcx.def_kind(id.owner_id), DefKind::AssocFn) {
90 bodies.push(id.owner_id);
94 // Trigger borrow checking of all bodies.
95 for def_id in bodies {
96 let _ = tcx.optimized_mir(def_id);
99 // See what bodies were borrow checked.
100 let mut bodies = get_bodies(tcx);
101 bodies.sort_by(|(def_id1, _), (def_id2, _)| def_id1.cmp(def_id2));
102 println!("Bodies retrieved for:");
103 for (def_id, body) in bodies {
104 println!("{}", def_id);
105 assert!(body.input_facts.cfg_edge.len() > 0);
109 Compilation::Continue
113 fn override_queries(_session: &Session, local: &mut Providers, _external: &mut ExternProviders) {
114 local.mir_borrowck = mir_borrowck;
117 // Since mir_borrowck does not have access to any other state, we need to use a
118 // thread-local for storing the obtained MIR bodies.
120 // Note: We are using 'static lifetime here, which is in general unsound.
121 // Unfortunately, that is the only lifetime allowed here. Our use is safe
122 // because we cast it back to `'tcx` before using.
124 pub static MIR_BODIES:
125 RefCell<HashMap<LocalDefId, BodyWithBorrowckFacts<'static>>> =
126 RefCell::new(HashMap::new());
129 fn mir_borrowck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> mir_borrowck<'tcx> {
130 let body_with_facts = rustc_borrowck::consumers::get_body_with_borrowck_facts(
132 ty::WithOptConstParam::unknown(def_id),
134 // SAFETY: The reader casts the 'static lifetime to 'tcx before using it.
135 let body_with_facts: BodyWithBorrowckFacts<'static> =
136 unsafe { std::mem::transmute(body_with_facts) };
137 MIR_BODIES.with(|state| {
138 let mut map = state.borrow_mut();
139 assert!(map.insert(def_id, body_with_facts).is_none());
141 let mut providers = Providers::default();
142 rustc_borrowck::provide(&mut providers);
143 let original_mir_borrowck = providers.mir_borrowck;
144 original_mir_borrowck(tcx, def_id)
147 /// Pull MIR bodies stored in the thread-local.
148 fn get_bodies<'tcx>(tcx: TyCtxt<'tcx>) -> Vec<(String, BodyWithBorrowckFacts<'tcx>)> {
149 MIR_BODIES.with(|state| {
150 let mut map = state.borrow_mut();
152 .map(|(def_id, body)| {
153 let def_path = tcx.def_path(def_id.to_def_id());
154 // SAFETY: For soundness we need to ensure that the bodies have
155 // the same lifetime (`'tcx`), which they had before they were
156 // stored in the thread local.
157 (def_path.to_string_no_crate_verbose(), unsafe { std::mem::transmute(body) })