Rollup merge of #111840 - voidc:borrowck-consumers, r=oli-obk

Expose more information in `get_body_with_borrowck_facts`

Verification tools for Rust such as, for example, Creusot or Prusti would benefit from having access to more information computed by the borrow checker.
As a first step in that direction, #86977 added the `get_body_with_borrowck_facts` API, allowing compiler consumers to obtain a `mir::Body` with accompanying borrow checker information.
At RustVerify 2023, multiple people working on verification tools expressed their need for a more comprehensive API.
While eventually borrow information could be part of Stable MIR, in the meantime, this PR proposes a more limited approach, extending the existing `get_body_with_borrowck_facts` API.
In summary, we propose the following changes:

- Permit obtaining the borrow-checked body without necessarily running Polonius
- Return the `BorrowSet` and the `RegionInferenceContext` in `BodyWithBorrowckFacts`
- Provide a way to compute the `borrows_out_of_scope_at_location` map
- Make some helper methods public

This is similar to #108328 but smaller in scope.
`@smoelius` Do you think these changes would also be sufficient for your needs?

r? `@oli-obk`
cc `@JonasAlaif`

compiler/rustc_borrowck/src/borrow_set.rs

@@ -30,7 +30,7 @@ pub struct BorrowSet<'tcx> {
     /// Map from local to all the borrows on that local.
     pub local_map: FxIndexMap<mir::Local, FxIndexSet<BorrowIndex>>,
 
-    pub(crate) locals_state_at_exit: LocalsStateAtExit,
+    pub locals_state_at_exit: LocalsStateAtExit,
 }
 
 impl<'tcx> Index<BorrowIndex> for BorrowSet<'tcx> {
@@ -153,7 +153,7 @@ impl<'tcx> BorrowSet<'tcx> {
         self.activation_map.get(&location).map_or(&[], |activations| &activations[..])
     }
 
-    pub(crate) fn len(&self) -> usize {
+    pub fn len(&self) -> usize {
         self.location_map.len()
     }
 

compiler/rustc_borrowck/src/consumers.rs

@@ -3,22 +3,95 @@
 //! This file provides API for compiler consumers.
 
 use rustc_hir::def_id::LocalDefId;
-use rustc_index::IndexSlice;
+use rustc_index::{IndexSlice, IndexVec};
 use rustc_infer::infer::{DefiningAnchor, TyCtxtInferExt};
-use rustc_middle::mir::Body;
+use rustc_middle::mir::{Body, Promoted};
 use rustc_middle::ty::TyCtxt;
+use std::rc::Rc;
+
+use crate::borrow_set::BorrowSet;
 
 pub use super::{
+    constraints::OutlivesConstraint,
+    dataflow::{calculate_borrows_out_of_scope_at_location, BorrowIndex, Borrows},
     facts::{AllFacts as PoloniusInput, RustcFacts},
     location::{LocationTable, RichLocation},
     nll::PoloniusOutput,
-    BodyWithBorrowckFacts,
+    place_ext::PlaceExt,
+    places_conflict::{places_conflict, PlaceConflictBias},
+    region_infer::RegionInferenceContext,
 };
 
-/// This function computes Polonius facts for the given body. It makes a copy of
-/// the body because it needs to regenerate the region identifiers. This function
-/// should never be invoked during a typical compilation session due to performance
-/// issues with Polonius.
+/// Options determining the output behavior of [`get_body_with_borrowck_facts`].
+///
+/// If executing under `-Z polonius` the choice here has no effect, and everything as if
+/// [`PoloniusOutputFacts`](ConsumerOptions::PoloniusOutputFacts) had been selected
+/// will be retrieved.
+#[derive(Debug, Copy, Clone)]
+pub enum ConsumerOptions {
+    /// Retrieve the [`Body`] along with the [`BorrowSet`](super::borrow_set::BorrowSet)
+    /// and [`RegionInferenceContext`]. If you would like the body only, use
+    /// [`TyCtxt::mir_promoted`].
+    ///
+    /// These can be used in conjunction with [`calculate_borrows_out_of_scope_at_location`].
+    RegionInferenceContext,
+    /// The recommended option. Retrieves the maximal amount of information
+    /// without significant slowdowns.
+    ///
+    /// Implies [`RegionInferenceContext`](ConsumerOptions::RegionInferenceContext),
+    /// and additionally retrieve the [`LocationTable`] and [`PoloniusInput`] that
+    /// would be given to Polonius. Critically, this does not run Polonius, which
+    /// one may want to avoid due to performance issues on large bodies.
+    PoloniusInputFacts,
+    /// Implies [`PoloniusInputFacts`](ConsumerOptions::PoloniusInputFacts),
+    /// and additionally runs Polonius to calculate the [`PoloniusOutput`].
+    PoloniusOutputFacts,
+}
+
+impl ConsumerOptions {
+    /// Should the Polonius input facts be computed?
+    pub(crate) fn polonius_input(&self) -> bool {
+        matches!(self, Self::PoloniusInputFacts | Self::PoloniusOutputFacts)
+    }
+    /// Should we run Polonius and collect the output facts?
+    pub(crate) fn polonius_output(&self) -> bool {
+        matches!(self, Self::PoloniusOutputFacts)
+    }
+}
+
+/// 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>,
+    /// The mir bodies of promoteds.
+    pub promoted: IndexVec<Promoted, Body<'tcx>>,
+    /// The set of borrows occurring in `body` with data about them.
+    pub borrow_set: Rc<BorrowSet<'tcx>>,
+    /// Context generated during borrowck, intended to be passed to
+    /// [`calculate_borrows_out_of_scope_at_location`].
+    pub region_inference_context: Rc<RegionInferenceContext<'tcx>>,
+    /// The table that maps Polonius points to locations in the table.
+    /// Populated when using [`ConsumerOptions::PoloniusInputFacts`]
+    /// or [`ConsumerOptions::PoloniusOutputFacts`].
+    pub location_table: Option<LocationTable>,
+    /// Polonius input facts.
+    /// Populated when using [`ConsumerOptions::PoloniusInputFacts`]
+    /// or [`ConsumerOptions::PoloniusOutputFacts`].
+    pub input_facts: Option<Box<PoloniusInput>>,
+    /// Polonius output facts. Populated when using
+    /// [`ConsumerOptions::PoloniusOutputFacts`].
+    pub output_facts: Option<Rc<PoloniusOutput>>,
+}
+
+/// This function computes borrowck facts for the given body. The [`ConsumerOptions`]
+/// determine which facts are returned. This function makes a copy of the body because
+/// it needs to regenerate the region identifiers. It should never be invoked during a
+/// typical compilation session due to the unnecessary overhead of returning
+/// [`BodyWithBorrowckFacts`].
 ///
 /// Note:
 /// *   This function will panic if the required body was already stolen. This
@@ -28,10 +101,14 @@ pub use super::{
 ///     that shows how to do this at `tests/run-make/obtain-borrowck/`.
 ///
 /// *   Polonius is highly unstable, so expect regular changes in its signature or other details.
-pub fn get_body_with_borrowck_facts(tcx: TyCtxt<'_>, def: LocalDefId) -> BodyWithBorrowckFacts<'_> {
+pub fn get_body_with_borrowck_facts(
+    tcx: TyCtxt<'_>,
+    def: LocalDefId,
+    options: ConsumerOptions,
+) -> BodyWithBorrowckFacts<'_> {
     let (input_body, promoted) = tcx.mir_promoted(def);
     let infcx = tcx.infer_ctxt().with_opaque_type_inference(DefiningAnchor::Bind(def)).build();
     let input_body: &Body<'_> = &input_body.borrow();
     let promoted: &IndexSlice<_, _> = &promoted.borrow();
-    *super::do_mir_borrowck(&infcx, input_body, promoted, true).1.unwrap()
+    *super::do_mir_borrowck(&infcx, input_body, promoted, Some(options)).1.unwrap()
 }

compiler/rustc_borrowck/src/dataflow.rs

@@ -231,27 +231,32 @@ impl<'tcx> OutOfScopePrecomputer<'_, 'tcx> {
     }
 }
 
+pub fn calculate_borrows_out_of_scope_at_location<'tcx>(
+    body: &Body<'tcx>,
+    regioncx: &RegionInferenceContext<'tcx>,
+    borrow_set: &BorrowSet<'tcx>,
+) -> FxIndexMap<Location, Vec<BorrowIndex>> {
+    let mut prec = OutOfScopePrecomputer::new(body, regioncx);
+    for (borrow_index, borrow_data) in borrow_set.iter_enumerated() {
+        let borrow_region = borrow_data.region;
+        let location = borrow_data.reserve_location;
+
+        prec.precompute_borrows_out_of_scope(borrow_index, borrow_region, location);
+    }
+
+    prec.borrows_out_of_scope_at_location
+}
+
 impl<'a, 'tcx> Borrows<'a, 'tcx> {
-    pub(crate) fn new(
+    pub fn new(
         tcx: TyCtxt<'tcx>,
         body: &'a Body<'tcx>,
         nonlexical_regioncx: &'a RegionInferenceContext<'tcx>,
         borrow_set: &'a BorrowSet<'tcx>,
     ) -> Self {
-        let mut prec = OutOfScopePrecomputer::new(body, nonlexical_regioncx);
-        for (borrow_index, borrow_data) in borrow_set.iter_enumerated() {
-            let borrow_region = borrow_data.region;
-            let location = borrow_data.reserve_location;
-
-            prec.precompute_borrows_out_of_scope(borrow_index, borrow_region, location);
-        }
-
-        Borrows {
-            tcx,
-            body,
-            borrow_set,
-            borrows_out_of_scope_at_location: prec.borrows_out_of_scope_at_location,
-        }
+        let borrows_out_of_scope_at_location =
+            calculate_borrows_out_of_scope_at_location(body, nonlexical_regioncx, borrow_set);
+        Borrows { tcx, body, borrow_set, borrows_out_of_scope_at_location }
     }
 
     pub fn location(&self, idx: BorrowIndex) -> &Location {

compiler/rustc_borrowck/src/lib.rs

@@ -61,7 +61,7 @@ use crate::session_diagnostics::VarNeedNotMut;
 use self::diagnostics::{AccessKind, RegionName};
 use self::location::LocationTable;
 use self::prefixes::PrefixSet;
-use facts::AllFacts;
+use consumers::{BodyWithBorrowckFacts, ConsumerOptions};
 
 use self::path_utils::*;
 
@@ -143,23 +143,23 @@ fn mir_borrowck(tcx: TyCtxt<'_>, def: LocalDefId) -> &BorrowCheckResult<'_> {
         tcx.infer_ctxt().with_opaque_type_inference(DefiningAnchor::Bind(hir_owner.def_id)).build();
     let input_body: &Body<'_> = &input_body.borrow();
     let promoted: &IndexSlice<_, _> = &promoted.borrow();
-    let opt_closure_req = do_mir_borrowck(&infcx, input_body, promoted, false).0;
+    let opt_closure_req = do_mir_borrowck(&infcx, input_body, promoted, None).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.
+/// Use `consumer_options: None` for the default behavior of returning
+/// [`BorrowCheckResult`] only. Otherwise, return [`BodyWithBorrowckFacts`] according
+/// to the given [`ConsumerOptions`].
 #[instrument(skip(infcx, input_body, input_promoted), fields(id=?input_body.source.def_id()), level = "debug")]
 fn do_mir_borrowck<'tcx>(
     infcx: &InferCtxt<'tcx>,
     input_body: &Body<'tcx>,
     input_promoted: &IndexSlice<Promoted, Body<'tcx>>,
-    return_body_with_facts: bool,
+    consumer_options: Option<ConsumerOptions>,
 ) -> (BorrowCheckResult<'tcx>, Option<Box<BodyWithBorrowckFacts<'tcx>>>) {
     let def = input_body.source.def_id().expect_local();
     debug!(?def);
@@ -240,8 +240,6 @@ fn do_mir_borrowck<'tcx>(
     let borrow_set =
         Rc::new(BorrowSet::build(tcx, body, locals_are_invalidated_at_exit, &mdpe.move_data));
 
-    let use_polonius = return_body_with_facts || infcx.tcx.sess.opts.unstable_opts.polonius;
-
     // Compute non-lexical lifetimes.
     let nll::NllOutput {
         regioncx,
@@ -261,7 +259,7 @@ fn do_mir_borrowck<'tcx>(
         &mdpe.move_data,
         &borrow_set,
         &upvars,
-        use_polonius,
+        consumer_options,
     );
 
     // Dump MIR results into a file, if that is enabled. This let us
@@ -441,13 +439,16 @@ fn do_mir_borrowck<'tcx>(
         tainted_by_errors,
     };
 
-    let body_with_facts = if return_body_with_facts {
-        let output_facts = mbcx.polonius_output.expect("Polonius output was not computed");
+    let body_with_facts = if consumer_options.is_some() {
+        let output_facts = mbcx.polonius_output;
         Some(Box::new(BodyWithBorrowckFacts {
             body: body_owned,
-            input_facts: *polonius_input.expect("Polonius input facts were not generated"),
+            promoted,
+            borrow_set,
+            region_inference_context: regioncx,
+            location_table: polonius_input.as_ref().map(|_| location_table_owned),
+            input_facts: polonius_input,
             output_facts,
-            location_table: location_table_owned,
         }))
     } else {
         None
@@ -458,22 +459,6 @@ fn do_mir_borrowck<'tcx>(
     (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,
-}
-
 pub struct BorrowckInferCtxt<'cx, 'tcx> {
     pub(crate) infcx: &'cx InferCtxt<'tcx>,
     pub(crate) reg_var_to_origin: RefCell<FxIndexMap<ty::RegionVid, RegionCtxt>>,

compiler/rustc_borrowck/src/nll.rs

@@ -27,6 +27,7 @@ use rustc_mir_dataflow::ResultsCursor;
 use crate::{
     borrow_set::BorrowSet,
     constraint_generation,
+    consumers::ConsumerOptions,
     diagnostics::RegionErrors,
     facts::{AllFacts, AllFactsExt, RustcFacts},
     invalidation,
@@ -165,10 +166,14 @@ pub(crate) fn compute_regions<'cx, 'tcx>(
     move_data: &MoveData<'tcx>,
     borrow_set: &BorrowSet<'tcx>,
     upvars: &[Upvar<'tcx>],
-    use_polonius: bool,
+    consumer_options: Option<ConsumerOptions>,
 ) -> NllOutput<'tcx> {
+    let polonius_input = consumer_options.map(|c| c.polonius_input()).unwrap_or_default()
+        || infcx.tcx.sess.opts.unstable_opts.polonius;
+    let polonius_output = consumer_options.map(|c| c.polonius_output()).unwrap_or_default()
+        || infcx.tcx.sess.opts.unstable_opts.polonius;
     let mut all_facts =
-        (use_polonius || AllFacts::enabled(infcx.tcx)).then_some(AllFacts::default());
+        (polonius_input || AllFacts::enabled(infcx.tcx)).then_some(AllFacts::default());
 
     let universal_regions = Rc::new(universal_regions);
 
@@ -189,7 +194,7 @@ pub(crate) fn compute_regions<'cx, 'tcx>(
             move_data,
             elements,
             upvars,
-            use_polonius,
+            polonius_input,
         );
 
     if let Some(all_facts) = &mut all_facts {
@@ -284,7 +289,7 @@ pub(crate) fn compute_regions<'cx, 'tcx>(
             all_facts.write_to_dir(dir_path, location_table).unwrap();
         }
 
-        if use_polonius {
+        if polonius_output {
             let algorithm =
                 env::var("POLONIUS_ALGORITHM").unwrap_or_else(|_| String::from("Hybrid"));
             let algorithm = Algorithm::from_str(&algorithm).unwrap();

compiler/rustc_borrowck/src/place_ext.rs

@@ -7,7 +7,7 @@ use rustc_middle::mir::{Body, Mutability, Place};
 use rustc_middle::ty::{self, TyCtxt};
 
 /// Extension methods for the `Place` type.
-pub(crate) trait PlaceExt<'tcx> {
+pub trait PlaceExt<'tcx> {
     /// Returns `true` if we can safely ignore borrows of this place.
     /// This is true whenever there is no action that the user can do
     /// to the place `self` that would invalidate the borrow. This is true

compiler/rustc_borrowck/src/places_conflict.rs

@@ -16,15 +16,15 @@ use std::iter;
 /// being run in the calling context, the conservative choice is to assume the compared indices
 /// are disjoint (and therefore, do not overlap).
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
-pub(crate) enum PlaceConflictBias {
+pub enum PlaceConflictBias {
     Overlap,
     NoOverlap,
 }
 
 /// Helper function for checking if places conflict with a mutable borrow and deep access depth.
 /// This is used to check for places conflicting outside of the borrow checking code (such as in
 /// dataflow).
-pub(crate) fn places_conflict<'tcx>(
+pub fn places_conflict<'tcx>(
     tcx: TyCtxt<'tcx>,
     body: &Body<'tcx>,
     borrow_place: Place<'tcx>,

compiler/rustc_borrowck/src/region_infer/mod.rs

@@ -585,6 +585,11 @@ impl<'tcx> RegionInferenceContext<'tcx> {
         self.universal_regions.to_region_vid(r)
     }
 
+    /// Returns an iterator over all the outlives constraints.
+    pub fn outlives_constraints(&self) -> impl Iterator<Item = OutlivesConstraint<'tcx>> + '_ {
+        self.constraints.outlives().iter().copied()
+    }
+
     /// Adds annotations for `#[rustc_regions]`; see `UniversalRegions::annotate`.
     pub(crate) fn annotate(&self, tcx: TyCtxt<'tcx>, err: &mut Diagnostic) {
         self.universal_regions.annotate(tcx, err)
@@ -698,7 +703,7 @@ impl<'tcx> RegionInferenceContext<'tcx> {
     #[instrument(skip(self, _body), level = "debug")]
     fn propagate_constraints(&mut self, _body: &Body<'tcx>) {
         debug!("constraints={:#?}", {
-            let mut constraints: Vec<_> = self.constraints.outlives().iter().collect();
+            let mut constraints: Vec<_> = self.outlives_constraints().collect();
             constraints.sort_by_key(|c| (c.sup, c.sub));
             constraints
                 .into_iter()