1 files changed, 57 insertions, 24 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
 }