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](https://github.com/viperproject/prusti-dev/blob/b58ced8dfd14ef30582b503d517167ccd771eaff/prusti-interface/src/environment/borrowck/regions.rs#L25-L39). This is not only a hacky approach, but we also reached its limits.

r? `@nikomatsakis`

Author: bors

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