tests/run-make-fulldeps/obtain-borrowck/driver.rs

   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_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;
  20
  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;
  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     // 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);
  55     }
  56
  57     // In this callback we trigger borrow checking of all functions and obtain
  58     // the result.
  59     fn after_analysis<'tcx>(
  60         &mut self,
  61         compiler: &Compiler,
  62         queries: &'tcx Queries<'tcx>,
  63     ) -> Compilation {
  64         compiler.session().abort_if_errors();
  65         queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
  66             // Collect definition ids of MIR bodies.
  67             let hir = tcx.hir();
  68             let mut bodies = Vec::new();
  69
  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);
  74                 }
  75             }
  76
  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);
  83                         }
  84                     }
  85                 }
  86             }
  87
  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);
  91                 }
  92             }
  93
  94             // Trigger borrow checking of all bodies.
  95             for def_id in bodies {
  96                 let _ = tcx.optimized_mir(def_id);
  97             }
  98
  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);
 106             }
 107         });
 108
 109         Compilation::Continue
 110     }
 111 }
 112
 113 fn override_queries(_session: &Session, local: &mut Providers, _external: &mut ExternProviders) {
 114     local.mir_borrowck = mir_borrowck;
 115 }
 116
 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.
 119 //
 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.
 123 thread_local! {
 124     pub static MIR_BODIES:
 125         RefCell<HashMap<LocalDefId, BodyWithBorrowckFacts<'static>>> =
 126         RefCell::new(HashMap::new());
 127 }
 128
 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(
 131         tcx,
 132         ty::WithOptConstParam::unknown(def_id),
 133     );
 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());
 140     });
 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)
 145 }
 146
 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();
 151         map.drain()
 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) })
 158             })
 159             .collect()
 160     })
 161 }