Enable compiler consumers to obtain Body with Polonius facts.

6 files changed, 94 insertions, 17 deletions

diff --git a/compiler/rustc_mir/src/borrow_check/consumers.rs b/compiler/rustc_mir/src/borrow_check/consumers.rs
new file mode 100644
index 00000000000..28c2c750136
--- /dev/null
+++ b/compiler/rustc_mir/src/borrow_check/consumers.rs
@@ -0,0 +1,28 @@
+//! 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.
+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()
+    })
+}
diff --git a/compiler/rustc_mir/src/borrow_check/facts.rs b/compiler/rustc_mir/src/borrow_check/facts.rs
index 6d6b94ecf64..daea5e538ed 100644
--- a/compiler/rustc_mir/src/borrow_check/facts.rs
+++ b/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
diff --git a/compiler/rustc_mir/src/borrow_check/location.rs b/compiler/rustc_mir/src/borrow_check/location.rs
index 375ff72679f..d378a2cbea3 100644
--- a/compiler/rustc_mir/src/borrow_check/location.rs
+++ b/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
diff --git a/compiler/rustc_mir/src/borrow_check/mod.rs b/compiler/rustc_mir/src/borrow_check/mod.rs
index 36eb8a4baa8..bda1c48e765 100644
--- a/compiler/rustc_mir/src/borrow_check/mod.rs
+++ b/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().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> {
diff --git a/compiler/rustc_mir/src/borrow_check/nll.rs b/compiler/rustc_mir/src/borrow_check/nll.rs
index bfeafa33a91..7742b76d9a4 100644
--- a/compiler/rustc_mir/src/borrow_check/nll.rs
+++ b/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,
diff --git a/compiler/rustc_mir/src/lib.rs b/compiler/rustc_mir/src/lib.rs
index a58ded9cfd3..5fb37b1b372 100644
--- a/compiler/rustc_mir/src/lib.rs
+++ b/compiler/rustc_mir/src/lib.rs
@@ -46,6 +46,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) {