use crate::{block_timer, mir::utils::SimplifyMir};
use rustc_borrowck::consumers::BodyWithBorrowckFacts;
use rustc_data_structures::fx::FxHashMap as HashMap;
use rustc_hir::def_id::LocalDefId;
use rustc_middle::ty::{
self,
query::{query_values::mir_borrowck, ExternProviders, Providers},
TyCtxt,
};
use rustc_mir_transform::MirPass;
use std::{cell::RefCell, pin::Pin};
pub fn override_queries(
_session: &rustc_session::Session,
local: &mut Providers,
_external: &mut ExternProviders,
) {
local.mir_borrowck = mir_borrowck;
}
thread_local! {
static MIR_BODIES: RefCell<HashMap<LocalDefId, Pin<Box<BodyWithBorrowckFacts<'static>>>>> =
RefCell::new(HashMap::default());
}
fn mir_borrowck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> mir_borrowck<'tcx> {
block_timer!(&format!(
"get_body_with_borrowck_facts for {}",
tcx.def_path_debug_str(def_id.to_def_id())
));
let mut body_with_facts = rustc_borrowck::consumers::get_body_with_borrowck_facts(
tcx,
ty::WithOptConstParam::unknown(def_id),
);
let body = &mut body_with_facts.body;
SimplifyMir.run_pass(tcx, body);
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, Pin::new(Box::new(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)
}
pub fn get_body_with_borrowck_facts(
tcx: TyCtxt<'tcx>,
def_id: LocalDefId,
) -> &'tcx BodyWithBorrowckFacts<'tcx> {
let _ = tcx.mir_borrowck(def_id);
MIR_BODIES.with(|state| {
let state = state.borrow();
let body = &*state.get(&def_id).unwrap();
unsafe {
std::mem::transmute::<&BodyWithBorrowckFacts<'static>, &'tcx BodyWithBorrowckFacts<'tcx>>(
body,
)
}
})
}