Auto merge of rust-lang#141057 - compiler-errors:fast-body, r=<try>

Don't do expensive HIR walk to collect nested bodies

An attempt to claw back some perf hit from the HIR-based impl of `nested_bodies_within`.

r? `@ghost`

Author: bors

compiler/rustc_hir/src/def.rs

@@ -339,6 +339,10 @@ impl DefKind {
             | DefKind::ExternCrate => false,
         }
     }
+
+    pub fn is_typeck_child(self) -> bool {
+        matches!(self, DefKind::Closure | DefKind::InlineConst | DefKind::SyntheticCoroutineBody)
+    }
 }
 
 /// The resolution of a path or export.

compiler/rustc_middle/src/ty/mod.rs

@@ -182,6 +182,7 @@ pub struct ResolverGlobalCtxt {
     pub extern_crate_map: UnordMap<LocalDefId, CrateNum>,
     pub maybe_unused_trait_imports: FxIndexSet<LocalDefId>,
     pub module_children: LocalDefIdMap<Vec<ModChild>>,
+    pub typeck_children: LocalDefIdMap<Vec<LocalDefId>>,
     pub glob_map: FxIndexMap<LocalDefId, FxIndexSet<Symbol>>,
     pub main_def: Option<MainDefinition>,
     pub trait_impls: FxIndexMap<DefId, Vec<LocalDefId>>,

compiler/rustc_middle/src/ty/util.rs

@@ -608,10 +608,7 @@ impl<'tcx> TyCtxt<'tcx> {
     /// Returns `true` if `def_id` refers to a definition that does not have its own
     /// type-checking context, i.e. closure, coroutine or inline const.
     pub fn is_typeck_child(self, def_id: DefId) -> bool {
-        matches!(
-            self.def_kind(def_id),
-            DefKind::Closure | DefKind::InlineConst | DefKind::SyntheticCoroutineBody
-        )
+        self.def_kind(def_id).is_typeck_child()
     }
 
     /// Returns `true` if `def_id` refers to a trait (i.e., `trait Foo { ... }`).

compiler/rustc_resolve/src/lib.rs

@@ -1073,6 +1073,8 @@ pub struct Resolver<'ra, 'tcx> {
     module_children: LocalDefIdMap<Vec<ModChild>>,
     trait_map: NodeMap<Vec<TraitCandidate>>,
 
+    typeck_children: LocalDefIdMap<Vec<LocalDefId>>,
+
     /// A map from nodes to anonymous modules.
     /// Anonymous modules are pseudo-modules that are implicitly created around items
     /// contained within blocks.
@@ -1347,6 +1349,11 @@ impl<'tcx> Resolver<'_, 'tcx> {
         let feed = self.tcx.create_def(parent, name, def_kind, None, &mut self.disambiguator);
         let def_id = feed.def_id();
 
+        if def_kind.is_typeck_child() {
+            let typeck_root_def_id = self.tcx.typeck_root_def_id(def_id.to_def_id()).expect_local();
+            self.typeck_children.entry(typeck_root_def_id).or_default().push(def_id);
+        }
+
         // Create the definition.
         if expn_id != ExpnId::root() {
             self.expn_that_defined.insert(def_id, expn_id);
@@ -1481,6 +1488,7 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
             extra_lifetime_params_map: Default::default(),
             extern_crate_map: Default::default(),
             module_children: Default::default(),
+            typeck_children: Default::default(),
             trait_map: NodeMap::default(),
             underscore_disambiguator: 0,
             empty_module,
@@ -1664,6 +1672,7 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
             effective_visibilities,
             extern_crate_map,
             module_children: self.module_children,
+            typeck_children: self.typeck_children,
             glob_map,
             maybe_unused_trait_imports,
             main_def,

compiler/rustc_ty_utils/src/nested_bodies.rs

@@ -1,15 +1,56 @@
+use rustc_data_structures::fx::FxIndexSet;
 use rustc_hir as hir;
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::intravisit::Visitor;
 use rustc_middle::query::Providers;
+use rustc_middle::span_bug;
 use rustc_middle::ty::{self, TyCtxt};
 
 fn nested_bodies_within<'tcx>(tcx: TyCtxt<'tcx>, item: LocalDefId) -> &'tcx ty::List<LocalDefId> {
-    let body = tcx.hir_body_owned_by(item);
-    let mut collector =
-        NestedBodiesVisitor { tcx, root_def_id: item.to_def_id(), nested_bodies: vec![] };
-    collector.visit_body(body);
-    tcx.mk_local_def_ids(&collector.nested_bodies)
+    let children = if let Some(children) = tcx.resolutions(()).typeck_children.get(&item) {
+        // We re-intern this as a `List` since `List` comparisons are cheap pointer equality.
+        tcx.mk_local_def_ids_from_iter(
+            children
+                .iter()
+                .copied()
+                // We need to filter out the `Node::Err` bodies which were unsuccessfully
+                // lowered in HIR lowering, like items that are contained in never pattern
+                // match arms, which don't get lowered.
+                .filter(|def_id| !matches!(tcx.hir_node_by_def_id(*def_id), hir::Node::Err(_))),
+        )
+    } else {
+        ty::List::empty()
+    };
+
+    // Double check that the list of children we're collecting here is consistent
+    // with what we see in the HIR.
+    if cfg!(debug_assertions) {
+        let body = tcx.hir_body_owned_by(item);
+        let mut collector =
+            NestedBodiesVisitor { tcx, root_def_id: item.to_def_id(), nested_bodies: vec![] };
+        collector.visit_body(body);
+
+        let mut expected = FxIndexSet::from_iter(children);
+        for found in collector.nested_bodies {
+            if !expected.shift_remove(&found) {
+                span_bug!(
+                    tcx.def_span(found),
+                    "did not expect {found:#?} to be collected as a nested \
+                    body child of {item:#?}, but it appeared in the HIR",
+                );
+            }
+        }
+
+        for expected in expected {
+            span_bug!(
+                tcx.def_span(expected),
+                "expected {expected:#?} to be collected as a \
+                nested body child of {item:#?}",
+            );
+        }
+    }
+
+    children
 }
 
 struct NestedBodiesVisitor<'tcx> {