|
| 1 | +#![feature(rustc_private)] |
| 2 | + |
| 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. |
| 13 | +
|
| 14 | +extern crate rustc_driver; |
| 15 | +extern crate rustc_hir; |
| 16 | +extern crate rustc_interface; |
| 17 | +extern crate rustc_middle; |
| 18 | +extern crate rustc_mir; |
| 19 | +extern crate rustc_session; |
| 20 | + |
| 21 | +use rustc_driver::Compilation; |
| 22 | +use rustc_hir::def_id::LocalDefId; |
| 23 | +use rustc_hir::itemlikevisit::ItemLikeVisitor; |
| 24 | +use rustc_interface::interface::Compiler; |
| 25 | +use rustc_interface::{Config, Queries}; |
| 26 | +use rustc_middle::ty::query::query_values::mir_borrowck; |
| 27 | +use rustc_middle::ty::query::Providers; |
| 28 | +use rustc_middle::ty::{self, TyCtxt}; |
| 29 | +use rustc_mir::consumers::BodyWithBorrowckFacts; |
| 30 | +use rustc_session::Session; |
| 31 | +use std::cell::RefCell; |
| 32 | +use std::collections::HashMap; |
| 33 | +use std::thread_local; |
| 34 | + |
| 35 | +fn main() { |
| 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() |
| 43 | + }); |
| 44 | + std::process::exit(exit_code); |
| 45 | +} |
| 46 | + |
| 47 | +#[derive(Default)] |
| 48 | +pub struct CompilerCalls; |
| 49 | + |
| 50 | +impl rustc_driver::Callbacks for CompilerCalls { |
| 51 | + |
| 52 | + // In this callback we override the mir_borrowck query. |
| 53 | + fn config(&mut self, config: &mut Config) { |
| 54 | + assert!(config.override_queries.is_none()); |
| 55 | + config.override_queries = Some(override_queries); |
| 56 | + } |
| 57 | + |
| 58 | + // In this callback we trigger borrow checking of all functions and obtain |
| 59 | + // the result. |
| 60 | + fn after_analysis<'tcx>( |
| 61 | + &mut self, |
| 62 | + compiler: &Compiler, |
| 63 | + queries: &'tcx Queries<'tcx>, |
| 64 | + ) -> Compilation { |
| 65 | + compiler.session().abort_if_errors(); |
| 66 | + queries.global_ctxt().unwrap().peek_mut().enter(|tcx| { |
| 67 | + |
| 68 | + // Collect definition ids of MIR bodies. |
| 69 | + let hir = tcx.hir(); |
| 70 | + let krate = hir.krate(); |
| 71 | + let mut visitor = HirVisitor { bodies: Vec::new() }; |
| 72 | + krate.visit_all_item_likes(&mut visitor); |
| 73 | + |
| 74 | + // Trigger borrow checking of all bodies. |
| 75 | + for def_id in visitor.bodies { |
| 76 | + let _ = tcx.optimized_mir(def_id); |
| 77 | + } |
| 78 | + |
| 79 | + // See what bodies were borrow checked. |
| 80 | + let mut bodies = get_bodies(tcx); |
| 81 | + bodies.sort_by(|(def_id1, _), (def_id2, _)| def_id1.cmp(def_id2)); |
| 82 | + println!("Bodies retrieved for:"); |
| 83 | + for (def_id, body) in bodies { |
| 84 | + println!("{}", def_id); |
| 85 | + assert!(body.input_facts.cfg_edge.len() > 0); |
| 86 | + } |
| 87 | + }); |
| 88 | + |
| 89 | + Compilation::Continue |
| 90 | + } |
| 91 | +} |
| 92 | + |
| 93 | +fn override_queries(_session: &Session, local: &mut Providers, external: &mut Providers) { |
| 94 | + local.mir_borrowck = mir_borrowck; |
| 95 | + external.mir_borrowck = mir_borrowck; |
| 96 | +} |
| 97 | + |
| 98 | +// Since mir_borrowck does not have access to any other state, we need to use a |
| 99 | +// thread-local for storing the obtained MIR bodies. |
| 100 | +// |
| 101 | +// Note: We are using 'static lifetime here, which is in general unsound. |
| 102 | +// Unfortunately, that is the only lifetime allowed here. Our use is safe |
| 103 | +// because we cast it back to `'tcx` before using. |
| 104 | +thread_local! { |
| 105 | + pub static MIR_BODIES: |
| 106 | + RefCell<HashMap<LocalDefId, BodyWithBorrowckFacts<'static>>> = |
| 107 | + RefCell::new(HashMap::new()); |
| 108 | +} |
| 109 | + |
| 110 | +fn mir_borrowck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> mir_borrowck<'tcx> { |
| 111 | + let body_with_facts = rustc_mir::consumers::get_body_with_borrowck_facts( |
| 112 | + tcx, |
| 113 | + ty::WithOptConstParam::unknown(def_id), |
| 114 | + ); |
| 115 | + // SAFETY: The reader casts the 'static lifetime to 'tcx before using it. |
| 116 | + let body_with_facts: BodyWithBorrowckFacts<'static> = |
| 117 | + unsafe { std::mem::transmute(body_with_facts) }; |
| 118 | + MIR_BODIES.with(|state| { |
| 119 | + let mut map = state.borrow_mut(); |
| 120 | + assert!(map.insert(def_id, body_with_facts).is_none()); |
| 121 | + }); |
| 122 | + let mut providers = Providers::default(); |
| 123 | + rustc_mir::provide(&mut providers); |
| 124 | + let original_mir_borrowck = providers.mir_borrowck; |
| 125 | + original_mir_borrowck(tcx, def_id) |
| 126 | +} |
| 127 | + |
| 128 | +/// Visitor that collects all body definition ids mentioned in the program. |
| 129 | +struct HirVisitor { |
| 130 | + bodies: Vec<LocalDefId>, |
| 131 | +} |
| 132 | + |
| 133 | +impl<'tcx> ItemLikeVisitor<'tcx> for HirVisitor { |
| 134 | + fn visit_item(&mut self, item: &rustc_hir::Item) { |
| 135 | + if let rustc_hir::ItemKind::Fn(..) = item.kind { |
| 136 | + self.bodies.push(item.def_id); |
| 137 | + } |
| 138 | + } |
| 139 | + |
| 140 | + fn visit_trait_item(&mut self, trait_item: &rustc_hir::TraitItem) { |
| 141 | + if let rustc_hir::TraitItemKind::Fn(_, trait_fn) = &trait_item.kind { |
| 142 | + if let rustc_hir::TraitFn::Provided(_) = trait_fn { |
| 143 | + self.bodies.push(trait_item.def_id); |
| 144 | + } |
| 145 | + } |
| 146 | + } |
| 147 | + |
| 148 | + fn visit_impl_item(&mut self, impl_item: &rustc_hir::ImplItem) { |
| 149 | + if let rustc_hir::ImplItemKind::Fn(..) = impl_item.kind { |
| 150 | + self.bodies.push(impl_item.def_id); |
| 151 | + } |
| 152 | + } |
| 153 | + |
| 154 | + fn visit_foreign_item(&mut self, _foreign_item: &rustc_hir::ForeignItem) {} |
| 155 | +} |
| 156 | + |
| 157 | +/// Pull MIR bodies stored in the thread-local. |
| 158 | +fn get_bodies<'tcx>(tcx: TyCtxt<'tcx>) -> Vec<(String, BodyWithBorrowckFacts<'tcx>)> { |
| 159 | + MIR_BODIES.with(|state| { |
| 160 | + let mut map = state.borrow_mut(); |
| 161 | + map.drain() |
| 162 | + .map(|(def_id, body)| { |
| 163 | + let def_path = tcx.def_path(def_id.to_def_id()); |
| 164 | + // SAFETY: For soundness we need to ensure that the bodies have |
| 165 | + // the same lifetime (`'tcx`), which they had before they were |
| 166 | + // stored in the thread local. |
| 167 | + (def_path.to_string_no_crate_verbose(), body) |
| 168 | + }) |
| 169 | + .collect() |
| 170 | + }) |
| 171 | +} |
0 commit comments