Don't do expensive HIR walk to collect nested bodies

Author: compiler-errors

compiler/rustc_ty_utils/src/nested_bodies.rs

@@ -1,32 +1,22 @@
-use rustc_hir as hir;
-use rustc_hir::def_id::{DefId, LocalDefId};
-use rustc_hir::intravisit::Visitor;
+use rustc_hir::def_id::LocalDefId;
 use rustc_middle::query::Providers;
 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)
-}
-
-struct NestedBodiesVisitor<'tcx> {
-    tcx: TyCtxt<'tcx>,
-    root_def_id: DefId,
-    nested_bodies: Vec<LocalDefId>,
-}
-
-impl<'tcx> Visitor<'tcx> for NestedBodiesVisitor<'tcx> {
-    fn visit_nested_body(&mut self, id: hir::BodyId) {
-        let body_def_id = self.tcx.hir_body_owner_def_id(id);
-        if self.tcx.typeck_root_def_id(body_def_id.to_def_id()) == self.root_def_id {
-            self.nested_bodies.push(body_def_id);
-            let body = self.tcx.hir_body(id);
-            self.visit_body(body);
-        }
-    }
+    let owner = tcx.local_def_id_to_hir_id(item).owner;
+    let children = tcx.mk_local_def_ids_from_iter(
+        tcx.hir_owner_nodes(owner)
+            .bodies
+            .iter()
+            .map(|&(_, body)| tcx.hir_body_owner_def_id(body.id()))
+            .filter(|&child_item| {
+                // Anon consts are not owner IDs, but they may have (e.g.) closures in them.
+                // Filter this just down to bodies that share the typeck root.
+                child_item != item
+                    && tcx.typeck_root_def_id(child_item.to_def_id()).expect_local() == item
+            }),
+    );
+    children
 }
 
 pub(super) fn provide(providers: &mut Providers) {