src/test/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::itemlikevisit::ItemLikeVisitor;
  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::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 visitor = HirVisitor { bodies: Vec::new() };
  69             hir.visit_all_item_likes(&mut visitor);
  70
  71             // Trigger borrow checking of all bodies.
  72             for def_id in visitor.bodies {
  73                 let _ = tcx.optimized_mir(def_id);
  74             }
  75
  76             // See what bodies were borrow checked.
  77             let mut bodies = get_bodies(tcx);
  78             bodies.sort_by(|(def_id1, _), (def_id2, _)| def_id1.cmp(def_id2));
  79             println!("Bodies retrieved for:");
  80             for (def_id, body) in bodies {
  81                 println!("{}", def_id);
  82                 assert!(body.input_facts.cfg_edge.len() > 0);
  83             }
  84         });
  85
  86         Compilation::Continue
  87     }
  88 }
  89
  90 fn override_queries(_session: &Session, local: &mut Providers, external: &mut Providers) {
  91     local.mir_borrowck = mir_borrowck;
  92     external.mir_borrowck = mir_borrowck;
  93 }
  94
  95 // Since mir_borrowck does not have access to any other state, we need to use a
  96 // thread-local for storing the obtained MIR bodies.
  97 //
  98 // Note: We are using 'static lifetime here, which is in general unsound.
  99 // Unfortunately, that is the only lifetime allowed here. Our use is safe
 100 // because we cast it back to `'tcx` before using.
 101 thread_local! {
 102     pub static MIR_BODIES:
 103         RefCell<HashMap<LocalDefId, BodyWithBorrowckFacts<'static>>> =
 104         RefCell::new(HashMap::new());
 105 }
 106
 107 fn mir_borrowck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> mir_borrowck<'tcx> {
 108     let body_with_facts = rustc_borrowck::consumers::get_body_with_borrowck_facts(
 109         tcx,
 110         ty::WithOptConstParam::unknown(def_id),
 111     );
 112     // SAFETY: The reader casts the 'static lifetime to 'tcx before using it.
 113     let body_with_facts: BodyWithBorrowckFacts<'static> =
 114         unsafe { std::mem::transmute(body_with_facts) };
 115     MIR_BODIES.with(|state| {
 116         let mut map = state.borrow_mut();
 117         assert!(map.insert(def_id, body_with_facts).is_none());
 118     });
 119     let mut providers = Providers::default();
 120     rustc_borrowck::provide(&mut providers);
 121     let original_mir_borrowck = providers.mir_borrowck;
 122     original_mir_borrowck(tcx, def_id)
 123 }
 124
 125 /// Visitor that collects all body definition ids mentioned in the program.
 126 struct HirVisitor {
 127     bodies: Vec<LocalDefId>,
 128 }
 129
 130 impl<'tcx> ItemLikeVisitor<'tcx> for HirVisitor {
 131     fn visit_item(&mut self, item: &rustc_hir::Item) {
 132         if let rustc_hir::ItemKind::Fn(..) = item.kind {
 133             self.bodies.push(item.def_id);
 134         }
 135     }
 136
 137     fn visit_trait_item(&mut self, trait_item: &rustc_hir::TraitItem) {
 138         if let rustc_hir::TraitItemKind::Fn(_, trait_fn) = &trait_item.kind {
 139             if let rustc_hir::TraitFn::Provided(_) = trait_fn {
 140                 self.bodies.push(trait_item.def_id);
 141             }
 142         }
 143     }
 144
 145     fn visit_impl_item(&mut self, impl_item: &rustc_hir::ImplItem) {
 146         if let rustc_hir::ImplItemKind::Fn(..) = impl_item.kind {
 147             self.bodies.push(impl_item.def_id);
 148         }
 149     }
 150
 151     fn visit_foreign_item(&mut self, _foreign_item: &rustc_hir::ForeignItem) {}
 152 }
 153
 154 /// Pull MIR bodies stored in the thread-local.
 155 fn get_bodies<'tcx>(tcx: TyCtxt<'tcx>) -> Vec<(String, BodyWithBorrowckFacts<'tcx>)> {
 156     MIR_BODIES.with(|state| {
 157         let mut map = state.borrow_mut();
 158         map.drain()
 159             .map(|(def_id, body)| {
 160                 let def_path = tcx.def_path(def_id.to_def_id());
 161                 // SAFETY: For soundness we need to ensure that the bodies have
 162                 // the same lifetime (`'tcx`), which they had before they were
 163                 // stored in the thread local.
 164                 (def_path.to_string_no_crate_verbose(), body)
 165             })
 166             .collect()
 167     })
 168 }