Try using MIR promoted.

Author: vakaras

prusti-interface/src/environment/body.rs

@@ -143,8 +143,10 @@ impl<'tcx> EnvBody<'tcx> {
     /// Get local MIR body of spec or pure functions. Retrieves the body from
     /// the compiler (relatively cheap).
     fn load_local_mir(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> MirBody<'tcx> {
-        // Throw away the borrowck facts
-        Self::load_local_mir_with_facts(tcx, def_id).body
+        // SAFETY: This is safe because we are feeding in the same `tcx`
+        // that was used to store the data.
+        let body = unsafe { mir_storage::retrieve_promoted_mir_body(tcx, def_id) };
+        MirBody(Rc::new(body))
     }
 
     fn get_monomorphised(

prusti-interface/src/environment/mir_storage.rs

@@ -7,15 +7,20 @@
 //! `'tcx`.
 
 use prusti_rustc_interface::{
-    borrowck::consumers::BodyWithBorrowckFacts, data_structures::fx::FxHashMap,
-    hir::def_id::LocalDefId, middle::ty::TyCtxt,
+    borrowck::consumers::BodyWithBorrowckFacts,
+    data_structures::fx::FxHashMap,
+    hir::def_id::LocalDefId,
+    middle::{mir, ty::TyCtxt},
 };
 use std::{cell::RefCell, thread_local};
 
 thread_local! {
-    pub static SHARED_STATE:
+    pub static SHARED_STATE_WITH_FACTS:
         RefCell<FxHashMap<LocalDefId, BodyWithBorrowckFacts<'static>>> =
         RefCell::new(FxHashMap::default());
+    pub static SHARED_STATE_WITHOUT_FACTS:
+        RefCell<FxHashMap<LocalDefId, mir::Body<'static>>> =
+        RefCell::new(FxHashMap::default());
 }
 
 /// # Safety
@@ -29,7 +34,7 @@ pub unsafe fn store_mir_body<'tcx>(
     // SAFETY: See the module level comment.
     let body_with_facts: BodyWithBorrowckFacts<'static> =
         unsafe { std::mem::transmute(body_with_facts) };
-    SHARED_STATE.with(|state| {
+    SHARED_STATE_WITH_FACTS.with(|state| {
         let mut map = state.borrow_mut();
         assert!(map.insert(def_id, body_with_facts).is_none());
     });
@@ -46,10 +51,39 @@ pub(super) unsafe fn retrieve_mir_body<'tcx>(
     _tcx: TyCtxt<'tcx>,
     def_id: LocalDefId,
 ) -> BodyWithBorrowckFacts<'tcx> {
-    let body_with_facts: BodyWithBorrowckFacts<'static> = SHARED_STATE.with(|state| {
+    let body_with_facts: BodyWithBorrowckFacts<'static> = SHARED_STATE_WITH_FACTS.with(|state| {
         let mut map = state.borrow_mut();
         map.remove(&def_id).unwrap()
     });
     // SAFETY: See the module level comment.
     unsafe { std::mem::transmute(body_with_facts) }
 }
+
+/// # Safety
+///
+/// See the module level comment.
+pub unsafe fn store_promoted_mir_body<'tcx>(
+    _tcx: TyCtxt<'tcx>,
+    def_id: LocalDefId,
+    body: mir::Body<'tcx>,
+) {
+    // SAFETY: See the module level comment.
+    let body: mir::Body<'static> = unsafe { std::mem::transmute(body) };
+    SHARED_STATE_WITHOUT_FACTS.with(|state| {
+        let mut map = state.borrow_mut();
+        assert!(map.insert(def_id, body).is_none());
+    });
+}
+
+#[allow(clippy::needless_lifetimes)] // We want to be very explicit about lifetimes here.
+pub(super) unsafe fn retrieve_promoted_mir_body<'tcx>(
+    _tcx: TyCtxt<'tcx>,
+    def_id: LocalDefId,
+) -> mir::Body<'tcx> {
+    let body_without_facts: mir::Body<'static> = SHARED_STATE_WITHOUT_FACTS.with(|state| {
+        let mut map = state.borrow_mut();
+        map.remove(&def_id).unwrap()
+    });
+    // SAFETY: See the module level comment.
+    unsafe { std::mem::transmute(body_without_facts) }
+}

prusti/src/callbacks.rs

@@ -6,11 +6,13 @@ use prusti_interface::{
 };
 use prusti_rustc_interface::{
     borrowck::consumers,
+    data_structures::steal::Steal,
     driver::Compilation,
     hir::{def::DefKind, def_id::LocalDefId},
+    index::IndexVec,
     interface::{interface::Compiler, Config, Queries},
     middle::{
-        mir::BorrowCheckResult,
+        mir::{self, BorrowCheckResult},
         query::{ExternProviders, Providers},
         ty::TyCtxt,
     },
@@ -50,12 +52,33 @@ fn mir_borrowck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> &BorrowCheckResu
     original_mir_borrowck(tcx, def_id)
 }
 
+#[allow(clippy::needless_lifetimes)]
+#[tracing::instrument(level = "debug", skip(tcx))]
+fn mir_promoted<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: LocalDefId,
+) -> (
+    &'tcx Steal<mir::Body<'tcx>>,
+    &'tcx Steal<IndexVec<mir::Promoted, mir::Body<'tcx>>>,
+) {
+    let original_mir_promoted =
+        prusti_rustc_interface::interface::DEFAULT_QUERY_PROVIDERS.mir_promoted;
+    let result = original_mir_promoted(tcx, def_id);
+    // SAFETY: This is safe because we are feeding in the same `tcx` that is
+    // going to be used as a witness when pulling out the data.
+    unsafe {
+        mir_storage::store_promoted_mir_body(tcx, def_id, result.0.borrow().clone());
+    }
+    result
+}
+
 impl prusti_rustc_interface::driver::Callbacks for PrustiCompilerCalls {
     fn config(&mut self, config: &mut Config) {
         assert!(config.override_queries.is_none());
         config.override_queries = Some(
             |_session: &Session, providers: &mut Providers, _external: &mut ExternProviders| {
                 providers.mir_borrowck = mir_borrowck;
+                providers.mir_promoted = mir_promoted;
             },
         );
     }