Expose nested bodies in rustc_borrowck::consumers

4 files changed, 89 insertions, 54 deletions

diff --git a/compiler/rustc_borrowck/src/consumers.rs b/compiler/rustc_borrowck/src/consumers.rs
index 1f087b09234..1c4ff5a6779 100644
--- a/compiler/rustc_borrowck/src/consumers.rs
+++ b/compiler/rustc_borrowck/src/consumers.rs
@@ -1,7 +1,9 @@
 //! This file provides API for compiler consumers.
 
+use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::def_id::LocalDefId;
 use rustc_index::IndexVec;
+use rustc_middle::bug;
 use rustc_middle::mir::{Body, Promoted};
 use rustc_middle::ty::TyCtxt;
 
@@ -17,7 +19,39 @@ pub use super::polonius::legacy::{
 pub use super::region_infer::RegionInferenceContext;
 use crate::{BorrowCheckRootCtxt, do_mir_borrowck};
 
-/// Options determining the output behavior of [`get_body_with_borrowck_facts`].
+/// Struct used during mir borrowck to collect bodies with facts for a typeck root and all
+/// its nested bodies.
+pub(crate) struct BorrowckConsumer<'tcx> {
+    options: ConsumerOptions,
+    bodies: FxHashMap<LocalDefId, BodyWithBorrowckFacts<'tcx>>,
+}
+
+impl<'tcx> BorrowckConsumer<'tcx> {
+    pub(crate) fn new(options: ConsumerOptions) -> Self {
+        Self { options, bodies: Default::default() }
+    }
+
+    pub(crate) fn insert_body(&mut self, def_id: LocalDefId, body: BodyWithBorrowckFacts<'tcx>) {
+        if self.bodies.insert(def_id, body).is_some() {
+            bug!("unexpected previous body for {def_id:?}");
+        }
+    }
+
+    /// Should the Polonius input facts be computed?
+    pub(crate) fn polonius_input(&self) -> bool {
+        matches!(
+            self.options,
+            ConsumerOptions::PoloniusInputFacts | ConsumerOptions::PoloniusOutputFacts
+        )
+    }
+
+    /// Should we run Polonius and collect the output facts?
+    pub(crate) fn polonius_output(&self) -> bool {
+        matches!(self.options, ConsumerOptions::PoloniusOutputFacts)
+    }
+}
+
+/// Options determining the output behavior of [`get_bodies_with_borrowck_facts`].
 ///
 /// If executing under `-Z polonius` the choice here has no effect, and everything as if
 /// [`PoloniusOutputFacts`](ConsumerOptions::PoloniusOutputFacts) had been selected
@@ -43,17 +77,6 @@ pub enum ConsumerOptions {
     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.
 ///
@@ -82,25 +105,35 @@ pub struct BodyWithBorrowckFacts<'tcx> {
     pub output_facts: Option<Box<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`].
+/// This function computes borrowck facts for the given def id and all its nested bodies.
+/// It must be called with a typeck root which will then borrowck all nested bodies as well.
+/// The [`ConsumerOptions`] determine which facts are returned. This function makes a copy
+/// of the bodies 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
+/// *   This function will panic if the required bodies were 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 `tests/run-make/obtain-borrowck/`.
+///     that shows how to do this at `tests/ui-fulldeps/obtain-borrowck.rs`.
 ///
 /// *   Polonius is highly unstable, so expect regular changes in its signature or other details.
-pub fn get_body_with_borrowck_facts(
+pub fn get_bodies_with_borrowck_facts(
     tcx: TyCtxt<'_>,
-    def_id: LocalDefId,
+    root_def_id: LocalDefId,
     options: ConsumerOptions,
-) -> BodyWithBorrowckFacts<'_> {
-    let mut root_cx = BorrowCheckRootCtxt::new(tcx, def_id);
-    *do_mir_borrowck(&mut root_cx, def_id, Some(options)).1.unwrap()
+) -> FxHashMap<LocalDefId, BodyWithBorrowckFacts<'_>> {
+    let mut root_cx =
+        BorrowCheckRootCtxt::new(tcx, root_def_id, Some(BorrowckConsumer::new(options)));
+
+    // See comment in `rustc_borrowck::mir_borrowck`
+    let nested_bodies = tcx.nested_bodies_within(root_def_id);
+    for def_id in nested_bodies {
+        root_cx.get_or_insert_nested(def_id);
+    }
+
+    do_mir_borrowck(&mut root_cx, root_def_id);
+    root_cx.consumer.unwrap().bodies
 }
diff --git a/compiler/rustc_borrowck/src/lib.rs b/compiler/rustc_borrowck/src/lib.rs
index 82b300dcb17..321b18c9b78 100644
--- a/compiler/rustc_borrowck/src/lib.rs
+++ b/compiler/rustc_borrowck/src/lib.rs
@@ -51,7 +51,7 @@ use smallvec::SmallVec;
 use tracing::{debug, instrument};
 
 use crate::borrow_set::{BorrowData, BorrowSet};
-use crate::consumers::{BodyWithBorrowckFacts, ConsumerOptions};
+use crate::consumers::BodyWithBorrowckFacts;
 use crate::dataflow::{BorrowIndex, Borrowck, BorrowckDomain, Borrows};
 use crate::diagnostics::{
     AccessKind, BorrowckDiagnosticsBuffer, IllegalMoveOriginKind, MoveError, RegionName,
@@ -124,7 +124,7 @@ fn mir_borrowck(
         let opaque_types = ConcreteOpaqueTypes(Default::default());
         Ok(tcx.arena.alloc(opaque_types))
     } else {
-        let mut root_cx = BorrowCheckRootCtxt::new(tcx, def);
+        let mut root_cx = BorrowCheckRootCtxt::new(tcx, def, None);
         // We need to manually borrowck all nested bodies from the HIR as
         // we do not generate MIR for dead code. Not doing so causes us to
         // never check closures in dead code.
@@ -134,7 +134,7 @@ fn mir_borrowck(
         }
 
         let PropagatedBorrowCheckResults { closure_requirements, used_mut_upvars } =
-            do_mir_borrowck(&mut root_cx, def, None).0;
+            do_mir_borrowck(&mut root_cx, def);
         debug_assert!(closure_requirements.is_none());
         debug_assert!(used_mut_upvars.is_empty());
         root_cx.finalize()
@@ -289,17 +289,12 @@ impl<'tcx> ClosureOutlivesSubjectTy<'tcx> {
 
 /// Perform the actual borrow checking.
 ///
-/// Use `consumer_options: None` for the default behavior of returning
-/// [`PropagatedBorrowCheckResults`] only. Otherwise, return [`BodyWithBorrowckFacts`]
-/// according to the given [`ConsumerOptions`].
-///
 /// For nested bodies this should only be called through `root_cx.get_or_insert_nested`.
 #[instrument(skip(root_cx), level = "debug")]
 fn do_mir_borrowck<'tcx>(
     root_cx: &mut BorrowCheckRootCtxt<'tcx>,
     def: LocalDefId,
-    consumer_options: Option<ConsumerOptions>,
-) -> (PropagatedBorrowCheckResults<'tcx>, Option<Box<BodyWithBorrowckFacts<'tcx>>>) {
+) -> PropagatedBorrowCheckResults<'tcx> {
     let tcx = root_cx.tcx;
     let infcx = BorrowckInferCtxt::new(tcx, def);
     let (input_body, promoted) = tcx.mir_promoted(def);
@@ -343,7 +338,6 @@ fn do_mir_borrowck<'tcx>(
         &location_table,
         &move_data,
         &borrow_set,
-        consumer_options,
     );
 
     // Dump MIR results into a file, if that is enabled. This lets us
@@ -483,23 +477,24 @@ fn do_mir_borrowck<'tcx>(
         used_mut_upvars: mbcx.used_mut_upvars,
     };
 
-    let body_with_facts = if consumer_options.is_some() {
-        Some(Box::new(BodyWithBorrowckFacts {
-            body: body_owned,
-            promoted,
-            borrow_set,
-            region_inference_context: regioncx,
-            location_table: polonius_input.as_ref().map(|_| location_table),
-            input_facts: polonius_input,
-            output_facts: polonius_output,
-        }))
-    } else {
-        None
-    };
+    if let Some(consumer) = &mut root_cx.consumer {
+        consumer.insert_body(
+            def,
+            BodyWithBorrowckFacts {
+                body: body_owned,
+                promoted,
+                borrow_set,
+                region_inference_context: regioncx,
+                location_table: polonius_input.as_ref().map(|_| location_table),
+                input_facts: polonius_input,
+                output_facts: polonius_output,
+            },
+        );
+    }
 
     debug!("do_mir_borrowck: result = {:#?}", result);
 
-    (result, body_with_facts)
+    result
 }
 
 fn get_flow_results<'a, 'tcx>(
diff --git a/compiler/rustc_borrowck/src/nll.rs b/compiler/rustc_borrowck/src/nll.rs
index af450507296..41f67e78930 100644
--- a/compiler/rustc_borrowck/src/nll.rs
+++ b/compiler/rustc_borrowck/src/nll.rs
@@ -18,7 +18,6 @@ use rustc_span::sym;
 use tracing::{debug, instrument};
 
 use crate::borrow_set::BorrowSet;
-use crate::consumers::ConsumerOptions;
 use crate::diagnostics::RegionErrors;
 use crate::handle_placeholders::compute_sccs_applying_placeholder_outlives_constraints;
 use crate::polonius::PoloniusDiagnosticsContext;
@@ -83,12 +82,11 @@ pub(crate) fn compute_regions<'tcx>(
     location_table: &PoloniusLocationTable,
     move_data: &MoveData<'tcx>,
     borrow_set: &BorrowSet<'tcx>,
-    consumer_options: Option<ConsumerOptions>,
 ) -> NllOutput<'tcx> {
     let is_polonius_legacy_enabled = infcx.tcx.sess.opts.unstable_opts.polonius.is_legacy_enabled();
-    let polonius_input = consumer_options.map(|c| c.polonius_input()).unwrap_or_default()
+    let polonius_input = root_cx.consumer.as_ref().map_or(false, |c| c.polonius_input())
         || is_polonius_legacy_enabled;
-    let polonius_output = consumer_options.map(|c| c.polonius_output()).unwrap_or_default()
+    let polonius_output = root_cx.consumer.as_ref().map_or(false, |c| c.polonius_output())
         || is_polonius_legacy_enabled;
     let mut polonius_facts =
         (polonius_input || PoloniusFacts::enabled(infcx.tcx)).then_some(PoloniusFacts::default());
diff --git a/compiler/rustc_borrowck/src/root_cx.rs b/compiler/rustc_borrowck/src/root_cx.rs
index 66b526fa02a..9b1d12aede5 100644
--- a/compiler/rustc_borrowck/src/root_cx.rs
+++ b/compiler/rustc_borrowck/src/root_cx.rs
@@ -6,6 +6,7 @@ use rustc_middle::ty::{OpaqueHiddenType, Ty, TyCtxt, TypeVisitableExt};
 use rustc_span::ErrorGuaranteed;
 use smallvec::SmallVec;
 
+use crate::consumers::BorrowckConsumer;
 use crate::{ClosureRegionRequirements, ConcreteOpaqueTypes, PropagatedBorrowCheckResults};
 
 /// The shared context used by both the root as well as all its nested
@@ -16,16 +17,24 @@ pub(super) struct BorrowCheckRootCtxt<'tcx> {
     concrete_opaque_types: ConcreteOpaqueTypes<'tcx>,
     nested_bodies: FxHashMap<LocalDefId, PropagatedBorrowCheckResults<'tcx>>,
     tainted_by_errors: Option<ErrorGuaranteed>,
+    /// This should be `None` during normal compilation. See [`crate::consumers`] for more
+    /// information on how this is used.
+    pub(crate) consumer: Option<BorrowckConsumer<'tcx>>,
 }
 
 impl<'tcx> BorrowCheckRootCtxt<'tcx> {
-    pub(super) fn new(tcx: TyCtxt<'tcx>, root_def_id: LocalDefId) -> BorrowCheckRootCtxt<'tcx> {
+    pub(super) fn new(
+        tcx: TyCtxt<'tcx>,
+        root_def_id: LocalDefId,
+        consumer: Option<BorrowckConsumer<'tcx>>,
+    ) -> BorrowCheckRootCtxt<'tcx> {
         BorrowCheckRootCtxt {
             tcx,
             root_def_id,
             concrete_opaque_types: Default::default(),
             nested_bodies: Default::default(),
             tainted_by_errors: None,
+            consumer,
         }
     }
 
@@ -71,7 +80,7 @@ impl<'tcx> BorrowCheckRootCtxt<'tcx> {
             self.root_def_id.to_def_id()
         );
         if !self.nested_bodies.contains_key(&def_id) {
-            let result = super::do_mir_borrowck(self, def_id, None).0;
+            let result = super::do_mir_borrowck(self, def_id);
             if let Some(prev) = self.nested_bodies.insert(def_id, result) {
                 bug!("unexpected previous nested body: {prev:?}");
             }