Auto merge of #86977 - vakaras:body_with_borrowck_facts, r=nikomatsakis

Enable compiler consumers to obtain mir::Body with Polonius facts.

This PR adds a function (``get_body_with_borrowck_facts``) that can be used by compiler consumers to obtain ``mir::Body`` with accompanying borrow checker information.

The most important borrow checker information that [our verifier called Prusti](https://github.com/viperproject/prusti-dev) needs is lifetime constraints. I have not found a reasonable way to compute the lifetime constraints on the Prusti side. In the compiler, the constraints are computed during the borrow checking phase and then dropped. This PR adds an additional parameter to the `do_mir_borrowck` function that tells it to return the computed information instead of dropping it.

The additionally returned information by `do_mir_borrowck` contains a ``mir::Body`` with non-erased lifetime regions and Polonius facts. I have decided to reuse the Polonius facts because this way I needed fewer changes to the compiler and Polonius facts contains other useful information that we otherwise would need to recompute.

Just FYI: up to now, Prusti was obtaining this information by [parsing the compiler logs](b58ced8dfd/prusti-interface/src/environment/borrowck/regions.rs (L25-L39)). This is not only a hacky approach, but we also reached its limits.

r? `@nikomatsakis`

compiler/rustc_mir/src/borrow_check/consumers.rs

@@ -0,0 +1,39 @@
+//! This file provides API for compiler consumers.
+
+use rustc_hir::def_id::LocalDefId;
+use rustc_index::vec::IndexVec;
+use rustc_infer::infer::TyCtxtInferExt;
+use rustc_middle::mir::Body;
+use rustc_middle::ty::{self, TyCtxt};
+
+pub use super::{
+    facts::{AllFacts as PoloniusInput, RustcFacts},
+    location::{LocationTable, RichLocation},
+    nll::PoloniusOutput,
+    BodyWithBorrowckFacts,
+};
+
+/// This function computes Polonius facts for the given body. It makes a copy of
+/// the body because it needs to regenerate the region identifiers.
+///
+/// Note:
+/// *   This function will panic if the required body was already stolen. This
+///     can, for example, happen when requesting a body of a `const` function
+///     because they are evaluated during typechecking. The panic can be avoided
+///     by overriding the `mir_borrowck` query. You can find a complete example
+///     that shows how to do this at `src/test/run-make/obtain-borrowck/`.
+/// *   This function will also panic if computation of Polonius facts
+///     (`-Zpolonius` flag) is not enabled.
+///
+/// *   Polonius is highly unstable, so expect regular changes in its signature or other details.
+pub fn get_body_with_borrowck_facts<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def: ty::WithOptConstParam<LocalDefId>,
+) -> BodyWithBorrowckFacts<'tcx> {
+    let (input_body, promoted) = tcx.mir_promoted(def);
+    tcx.infer_ctxt().enter(|infcx| {
+        let input_body: &Body<'_> = &input_body.borrow();
+        let promoted: &IndexVec<_, _> = &promoted.borrow();
+        *super::do_mir_borrowck(&infcx, input_body, promoted, true).1.unwrap()
+    })
+}

compiler/rustc_mir/src/borrow_check/facts.rs

@@ -12,7 +12,7 @@ use std::io::{BufWriter, Write};
 use std::path::Path;
 
 #[derive(Copy, Clone, Debug)]
-crate struct RustcFacts;
+pub struct RustcFacts;
 
 impl polonius_engine::FactTypes for RustcFacts {
     type Origin = RegionVid;
@@ -22,7 +22,7 @@ impl polonius_engine::FactTypes for RustcFacts {
     type Path = MovePathIndex;
 }
 
-crate type AllFacts = PoloniusFacts<RustcFacts>;
+pub type AllFacts = PoloniusFacts<RustcFacts>;
 
 crate trait AllFactsExt {
     /// Returns `true` if there is a need to gather `AllFacts` given the

compiler/rustc_mir/src/borrow_check/location.rs

@@ -12,7 +12,7 @@ use rustc_middle::mir::{BasicBlock, Body, Location};
 /// granularity through outlives relations; however, the rich location
 /// table serves another purpose: it compresses locations from
 /// multiple words into a single u32.
-crate struct LocationTable {
+pub struct LocationTable {
     num_points: usize,
     statements_before_block: IndexVec<BasicBlock, usize>,
 }
@@ -24,7 +24,7 @@ rustc_index::newtype_index! {
 }
 
 #[derive(Copy, Clone, Debug)]
-crate enum RichLocation {
+pub enum RichLocation {
     Start(Location),
     Mid(Location),
 }
@@ -48,23 +48,23 @@ impl LocationTable {
         Self { num_points, statements_before_block }
     }
 
-    crate fn all_points(&self) -> impl Iterator<Item = LocationIndex> {
+    pub fn all_points(&self) -> impl Iterator<Item = LocationIndex> {
         (0..self.num_points).map(LocationIndex::new)
     }
 
-    crate fn start_index(&self, location: Location) -> LocationIndex {
+    pub fn start_index(&self, location: Location) -> LocationIndex {
         let Location { block, statement_index } = location;
         let start_index = self.statements_before_block[block];
         LocationIndex::new(start_index + statement_index * 2)
     }
 
-    crate fn mid_index(&self, location: Location) -> LocationIndex {
+    pub fn mid_index(&self, location: Location) -> LocationIndex {
         let Location { block, statement_index } = location;
         let start_index = self.statements_before_block[block];
         LocationIndex::new(start_index + statement_index * 2 + 1)
     }
 
-    crate fn to_location(&self, index: LocationIndex) -> RichLocation {
+    pub fn to_location(&self, index: LocationIndex) -> RichLocation {
         let point_index = index.index();
 
         // Find the basic block. We have a vector with the

compiler/rustc_mir/src/borrow_check/mod.rs

@@ -42,12 +42,14 @@ use self::diagnostics::{AccessKind, RegionName};
 use self::location::LocationTable;
 use self::prefixes::PrefixSet;
 use self::MutateMode::{JustWrite, WriteAndRead};
+use facts::AllFacts;
 
 use self::path_utils::*;
 
 mod borrow_set;
 mod constraint_generation;
 mod constraints;
+pub mod consumers;
 mod def_use;
 mod diagnostics;
 mod facts;
@@ -108,22 +110,33 @@ fn mir_borrowck<'tcx>(
     let opt_closure_req = tcx.infer_ctxt().with_opaque_type_inference(def.did).enter(|infcx| {
         let input_body: &Body<'_> = &input_body.borrow();
         let promoted: &IndexVec<_, _> = &promoted.borrow();
-        do_mir_borrowck(&infcx, input_body, promoted)
+        do_mir_borrowck(&infcx, input_body, promoted, false).0
     });
     debug!("mir_borrowck done");
 
     tcx.arena.alloc(opt_closure_req)
 }
 
+/// Perform the actual borrow checking.
+///
+/// If `return_body_with_facts` is true, then return the body with non-erased
+/// region ids on which the borrow checking was performed together with Polonius
+/// facts.
 fn do_mir_borrowck<'a, 'tcx>(
     infcx: &InferCtxt<'a, 'tcx>,
     input_body: &Body<'tcx>,
     input_promoted: &IndexVec<Promoted, Body<'tcx>>,
-) -> BorrowCheckResult<'tcx> {
+    return_body_with_facts: bool,
+) -> (BorrowCheckResult<'tcx>, Option<Box<BodyWithBorrowckFacts<'tcx>>>) {
     let def = input_body.source.with_opt_param().as_local().unwrap();
 
     debug!("do_mir_borrowck(def = {:?})", def);
 
+    assert!(
+        !return_body_with_facts || infcx.tcx.sess.opts.debugging_opts.polonius,
+        "borrowck facts can be requested only when Polonius is enabled"
+    );
+
     let tcx = infcx.tcx;
     let param_env = tcx.param_env(def.did);
     let id = tcx.hir().local_def_id_to_hir_id(def.did);
@@ -169,12 +182,14 @@ fn do_mir_borrowck<'a, 'tcx>(
     // requires first making our own copy of the MIR. This copy will
     // be modified (in place) to contain non-lexical lifetimes. It
     // will have a lifetime tied to the inference context.
-    let mut body = input_body.clone();
+    let mut body_owned = input_body.clone();
     let mut promoted = input_promoted.clone();
-    let free_regions = nll::replace_regions_in_mir(infcx, param_env, &mut body, &mut promoted);
-    let body = &body; // no further changes
+    let free_regions =
+        nll::replace_regions_in_mir(infcx, param_env, &mut body_owned, &mut promoted);
+    let body = &body_owned; // no further changes
 
-    let location_table = &LocationTable::new(&body);
+    let location_table_owned = LocationTable::new(body);
+    let location_table = &location_table_owned;
 
     let mut errors_buffer = Vec::new();
     let (move_data, move_errors): (MoveData<'tcx>, Vec<(Place<'tcx>, MoveError<'tcx>)>) =
@@ -202,6 +217,7 @@ fn do_mir_borrowck<'a, 'tcx>(
     let nll::NllOutput {
         regioncx,
         opaque_type_values,
+        polonius_input,
         polonius_output,
         opt_closure_req,
         nll_errors,
@@ -446,9 +462,37 @@ fn do_mir_borrowck<'a, 'tcx>(
         used_mut_upvars: mbcx.used_mut_upvars,
     };
 
+    let body_with_facts = if return_body_with_facts {
+        let output_facts = mbcx.polonius_output.expect("Polonius output was not computed");
+        Some(box BodyWithBorrowckFacts {
+            body: body_owned,
+            input_facts: *polonius_input.expect("Polonius input facts were not generated"),
+            output_facts,
+            location_table: location_table_owned,
+        })
+    } else {
+        None
+    };
+
     debug!("do_mir_borrowck: result = {:#?}", result);
 
-    result
+    (result, body_with_facts)
+}
+
+/// A `Body` with information computed by the borrow checker. This struct is
+/// intended to be consumed by compiler consumers.
+///
+/// We need to include the MIR body here because the region identifiers must
+/// match the ones in the Polonius facts.
+pub struct BodyWithBorrowckFacts<'tcx> {
+    /// A mir body that contains region identifiers.
+    pub body: Body<'tcx>,
+    /// Polonius input facts.
+    pub input_facts: AllFacts,
+    /// Polonius output facts.
+    pub output_facts: Rc<self::nll::PoloniusOutput>,
+    /// The table that maps Polonius points to locations in the table.
+    pub location_table: LocationTable,
 }
 
 crate struct MirBorrowckCtxt<'cx, 'tcx> {

compiler/rustc_mir/src/borrow_check/nll.rs

@@ -40,13 +40,14 @@ use crate::borrow_check::{
     Upvar,
 };
 
-crate type PoloniusOutput = Output<RustcFacts>;
+pub type PoloniusOutput = Output<RustcFacts>;
 
 /// The output of `nll::compute_regions`. This includes the computed `RegionInferenceContext`, any
 /// closure requirements to propagate, and any generated errors.
 crate struct NllOutput<'tcx> {
     pub regioncx: RegionInferenceContext<'tcx>,
     pub opaque_type_values: VecMap<OpaqueTypeKey<'tcx>, Ty<'tcx>>,
+    pub polonius_input: Option<Box<AllFacts>>,
     pub polonius_output: Option<Rc<PoloniusOutput>>,
     pub opt_closure_req: Option<ClosureRegionRequirements<'tcx>>,
     pub nll_errors: RegionErrors<'tcx>,
@@ -271,7 +272,7 @@ pub(in crate::borrow_check) fn compute_regions<'cx, 'tcx>(
     let def_id = body.source.def_id();
 
     // Dump facts if requested.
-    let polonius_output = all_facts.and_then(|all_facts| {
+    let polonius_output = all_facts.as_ref().and_then(|all_facts| {
         if infcx.tcx.sess.opts.debugging_opts.nll_facts {
             let def_path = infcx.tcx.def_path(def_id);
             let dir_path = PathBuf::from(&infcx.tcx.sess.opts.debugging_opts.nll_facts_dir)
@@ -305,6 +306,7 @@ pub(in crate::borrow_check) fn compute_regions<'cx, 'tcx>(
     NllOutput {
         regioncx,
         opaque_type_values: remapped_opaque_tys,
+        polonius_input: all_facts.map(Box::new),
         polonius_output,
         opt_closure_req: closure_region_requirements,
         nll_errors,

compiler/rustc_mir/src/lib.rs

@@ -47,6 +47,9 @@ mod shim;
 pub mod transform;
 pub mod util;
 
+// A public API provided for the Rust compiler consumers.
+pub use self::borrow_check::consumers;
+
 use rustc_middle::ty::query::Providers;
 
 pub fn provide(providers: &mut Providers) {

src/test/run-make-fulldeps/obtain-borrowck/Makefile

@@ -0,0 +1,26 @@
+include ../tools.mk
+
+# This example shows how to implement a rustc driver that retrieves MIR bodies
+# together with the borrow checker information.
+
+# How to run this
+# $ ./x.py test src/test/run-make-fulldeps/obtain-borrowck
+
+DRIVER_BINARY := "$(TMPDIR)"/driver
+SYSROOT := $(shell $(RUSTC) --print sysroot)
+
+ifdef IS_WINDOWS
+LIBSTD := -L "$(SYSROOT)\\lib\\rustlib\\$(TARGET)\\lib"
+else
+LIBSTD :=
+endif
+
+all:
+	$(RUSTC) driver.rs -o "$(DRIVER_BINARY)"
+	$(TARGET_RPATH_ENV) "$(DRIVER_BINARY)" --sysroot $(SYSROOT) $(LIBSTD) test.rs -o "$(TMPDIR)/driver_test" > "$(TMPDIR)"/output.stdout
+
+ifdef RUSTC_BLESS_TEST
+	cp "$(TMPDIR)"/output.stdout output.stdout
+else
+	$(DIFF) output.stdout "$(TMPDIR)"/output.stdout
+endif

src/test/run-make-fulldeps/obtain-borrowck/driver.rs

@@ -0,0 +1,171 @@
+#![feature(rustc_private)]
+
+//! This program implements a rustc driver that retrieves MIR bodies with
+//! borrowck information. This cannot be done in a straightforward way because
+//! `get_body_with_borrowck_facts`–the function for retrieving a MIR body with
+//! borrowck facts–can panic if the body is stolen before it is invoked.
+//! Therefore, the driver overrides `mir_borrowck` query (this is done in the
+//! `config` callback), which retrieves the body that is about to be borrow
+//! checked and stores it in a thread local `MIR_BODIES`. Then, `after_analysis`
+//! callback triggers borrow checking of all MIR bodies by retrieving
+//! `optimized_mir` and pulls out the MIR bodies with the borrowck information
+//! from the thread local storage.
+
+extern crate rustc_driver;
+extern crate rustc_hir;
+extern crate rustc_interface;
+extern crate rustc_middle;
+extern crate rustc_mir;
+extern crate rustc_session;
+
+use rustc_driver::Compilation;
+use rustc_hir::def_id::LocalDefId;
+use rustc_hir::itemlikevisit::ItemLikeVisitor;
+use rustc_interface::interface::Compiler;
+use rustc_interface::{Config, Queries};
+use rustc_middle::ty::query::query_values::mir_borrowck;
+use rustc_middle::ty::query::Providers;
+use rustc_middle::ty::{self, TyCtxt};
+use rustc_mir::consumers::BodyWithBorrowckFacts;
+use rustc_session::Session;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::thread_local;
+
+fn main() {
+    let exit_code = rustc_driver::catch_with_exit_code(move || {
+        let mut rustc_args: Vec<_> = std::env::args().collect();
+        // We must pass -Zpolonius so that the borrowck information is computed.
+        rustc_args.push("-Zpolonius".to_owned());
+        let mut callbacks = CompilerCalls::default();
+        // Call the Rust compiler with our callbacks.
+        rustc_driver::RunCompiler::new(&rustc_args, &mut callbacks).run()
+    });
+    std::process::exit(exit_code);
+}
+
+#[derive(Default)]
+pub struct CompilerCalls;
+
+impl rustc_driver::Callbacks for CompilerCalls {
+
+    // In this callback we override the mir_borrowck query.
+    fn config(&mut self, config: &mut Config) {
+        assert!(config.override_queries.is_none());
+        config.override_queries = Some(override_queries);
+    }
+
+    // In this callback we trigger borrow checking of all functions and obtain
+    // the result.
+    fn after_analysis<'tcx>(
+        &mut self,
+        compiler: &Compiler,
+        queries: &'tcx Queries<'tcx>,
+    ) -> Compilation {
+        compiler.session().abort_if_errors();
+        queries.global_ctxt().unwrap().peek_mut().enter(|tcx| {
+
+            // Collect definition ids of MIR bodies.
+            let hir = tcx.hir();
+            let krate = hir.krate();
+            let mut visitor = HirVisitor { bodies: Vec::new() };
+            krate.visit_all_item_likes(&mut visitor);
+
+            // Trigger borrow checking of all bodies.
+            for def_id in visitor.bodies {
+                let _ = tcx.optimized_mir(def_id);
+            }
+
+            // See what bodies were borrow checked.
+            let mut bodies = get_bodies(tcx);
+            bodies.sort_by(|(def_id1, _), (def_id2, _)| def_id1.cmp(def_id2));
+            println!("Bodies retrieved for:");
+            for (def_id, body) in bodies {
+                println!("{}", def_id);
+                assert!(body.input_facts.cfg_edge.len() > 0);
+            }
+        });
+
+        Compilation::Continue
+    }
+}
+
+fn override_queries(_session: &Session, local: &mut Providers, external: &mut Providers) {
+    local.mir_borrowck = mir_borrowck;
+    external.mir_borrowck = mir_borrowck;
+}
+
+// Since mir_borrowck does not have access to any other state, we need to use a
+// thread-local for storing the obtained MIR bodies.
+//
+// Note: We are using 'static lifetime here, which is in general unsound.
+// Unfortunately, that is the only lifetime allowed here. Our use is safe
+// because we cast it back to `'tcx` before using.
+thread_local! {
+    pub static MIR_BODIES:
+        RefCell<HashMap<LocalDefId, BodyWithBorrowckFacts<'static>>> =
+        RefCell::new(HashMap::new());
+}
+
+fn mir_borrowck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> mir_borrowck<'tcx> {
+    let body_with_facts = rustc_mir::consumers::get_body_with_borrowck_facts(
+        tcx,
+        ty::WithOptConstParam::unknown(def_id),
+    );
+    // SAFETY: The reader casts the 'static lifetime to 'tcx before using it.
+    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, body_with_facts).is_none());
+    });
+    let mut providers = Providers::default();
+    rustc_mir::provide(&mut providers);
+    let original_mir_borrowck = providers.mir_borrowck;
+    original_mir_borrowck(tcx, def_id)
+}
+
+/// Visitor that collects all body definition ids mentioned in the program.
+struct HirVisitor {
+    bodies: Vec<LocalDefId>,
+}
+
+impl<'tcx> ItemLikeVisitor<'tcx> for HirVisitor {
+    fn visit_item(&mut self, item: &rustc_hir::Item) {
+        if let rustc_hir::ItemKind::Fn(..) = item.kind {
+            self.bodies.push(item.def_id);
+        }
+    }
+
+    fn visit_trait_item(&mut self, trait_item: &rustc_hir::TraitItem) {
+        if let rustc_hir::TraitItemKind::Fn(_, trait_fn) = &trait_item.kind {
+            if let rustc_hir::TraitFn::Provided(_) = trait_fn {
+                self.bodies.push(trait_item.def_id);
+            }
+        }
+    }
+
+    fn visit_impl_item(&mut self, impl_item: &rustc_hir::ImplItem) {
+        if let rustc_hir::ImplItemKind::Fn(..) = impl_item.kind {
+            self.bodies.push(impl_item.def_id);
+        }
+    }
+
+    fn visit_foreign_item(&mut self, _foreign_item: &rustc_hir::ForeignItem) {}
+}
+
+/// Pull MIR bodies stored in the thread-local.
+fn get_bodies<'tcx>(tcx: TyCtxt<'tcx>) -> Vec<(String, BodyWithBorrowckFacts<'tcx>)> {
+    MIR_BODIES.with(|state| {
+        let mut map = state.borrow_mut();
+        map.drain()
+            .map(|(def_id, body)| {
+                let def_path = tcx.def_path(def_id.to_def_id());
+                // SAFETY: For soundness we need to ensure that the bodies have
+                // the same lifetime (`'tcx`), which they had before they were
+                // stored in the thread local.
+                (def_path.to_string_no_crate_verbose(), body)
+            })
+            .collect()
+    })
+}

src/test/run-make-fulldeps/obtain-borrowck/output.stdout

@@ -0,0 +1,8 @@
+Bodies retrieved for:
+::X::provided
+::foo
+::main
+::main::{constant#0}
+::{impl#0}::new
+::{impl#1}::provided
+::{impl#1}::required

src/test/run-make-fulldeps/obtain-borrowck/test.rs

@@ -0,0 +1,32 @@
+trait X {
+    fn provided(&self) -> usize {
+        5
+    }
+    fn required(&self) -> u32;
+}
+
+struct Bar;
+
+impl Bar {
+    fn new() -> Self {
+        Self
+    }
+}
+
+impl X for Bar {
+    fn provided(&self) -> usize {
+        1
+    }
+    fn required(&self) -> u32 {
+        7
+    }
+}
+
+const fn foo() -> usize {
+    1
+}
+
+fn main() {
+    let bar: [Bar; foo()] = [Bar::new()];
+    assert_eq!(bar[0].provided(), foo());
+}