Auto merge of #98203 - kckeiks:gather-body-owners-in-hir-item-queries, r=cjgillot

gather body owners

Issue #96341

Author: bors

compiler/rustc_middle/src/hir/map/mod.rs

@@ -39,6 +39,7 @@ pub fn fn_sig<'hir>(node: Node<'hir>) -> Option<&'hir FnSig<'hir>> {
     }
 }
 
+#[inline]
 pub fn associated_body<'hir>(node: Node<'hir>) -> Option<BodyId> {
     match node {
         Node::Item(Item {
@@ -486,35 +487,13 @@ impl<'hir> Map<'hir> {
     /// crate. If you would prefer to iterate over the bodies
     /// themselves, you can do `self.hir().krate().body_ids.iter()`.
     pub fn body_owners(self) -> impl Iterator<Item = LocalDefId> + 'hir {
-        self.krate()
-            .owners
-            .iter_enumerated()
-            .flat_map(move |(owner, owner_info)| {
-                let bodies = &owner_info.as_owner()?.nodes.bodies;
-                Some(bodies.iter().map(move |&(local_id, _)| {
-                    let hir_id = HirId { owner, local_id };
-                    let body_id = BodyId { hir_id };
-                    self.body_owner_def_id(body_id)
-                }))
-            })
-            .flatten()
+        self.tcx.hir_crate_items(()).body_owners.iter().copied()
     }
 
     pub fn par_body_owners<F: Fn(LocalDefId) + Sync + Send>(self, f: F) {
         use rustc_data_structures::sync::{par_iter, ParallelIterator};
-        #[cfg(parallel_compiler)]
-        use rustc_rayon::iter::IndexedParallelIterator;
-
-        par_iter(&self.krate().owners.raw).enumerate().for_each(|(owner, owner_info)| {
-            let owner = LocalDefId::new(owner);
-            if let MaybeOwner::Owner(owner_info) = owner_info {
-                par_iter(owner_info.nodes.bodies.range(..)).for_each(|(local_id, _)| {
-                    let hir_id = HirId { owner, local_id: *local_id };
-                    let body_id = BodyId { hir_id };
-                    f(self.body_owner_def_id(body_id))
-                })
-            }
-        });
+
+        par_iter(&self.tcx.hir_crate_items(()).body_owners[..]).for_each(|&def_id| f(def_id));
     }
 
     pub fn ty_param_owner(self, def_id: LocalDefId) -> LocalDefId {
@@ -1283,133 +1262,145 @@ fn hir_id_to_string(map: Map<'_>, id: HirId) -> String {
 }
 
 pub(super) fn hir_module_items(tcx: TyCtxt<'_>, module_id: LocalDefId) -> ModuleItems {
-    let mut collector = ModuleCollector {
-        tcx,
-        submodules: Vec::default(),
-        items: Vec::default(),
-        trait_items: Vec::default(),
-        impl_items: Vec::default(),
-        foreign_items: Vec::default(),
-    };
+    let mut collector = ItemCollector::new(tcx, false);
 
     let (hir_mod, span, hir_id) = tcx.hir().get_module(module_id);
     collector.visit_mod(hir_mod, span, hir_id);
 
-    let ModuleCollector { submodules, items, trait_items, impl_items, foreign_items, .. } =
-        collector;
+    let ItemCollector {
+        submodules,
+        items,
+        trait_items,
+        impl_items,
+        foreign_items,
+        body_owners,
+        ..
+    } = collector;
     return ModuleItems {
         submodules: submodules.into_boxed_slice(),
         items: items.into_boxed_slice(),
         trait_items: trait_items.into_boxed_slice(),
         impl_items: impl_items.into_boxed_slice(),
         foreign_items: foreign_items.into_boxed_slice(),
+        body_owners: body_owners.into_boxed_slice(),
     };
-
-    struct ModuleCollector<'tcx> {
-        tcx: TyCtxt<'tcx>,
-        submodules: Vec<LocalDefId>,
-        items: Vec<ItemId>,
-        trait_items: Vec<TraitItemId>,
-        impl_items: Vec<ImplItemId>,
-        foreign_items: Vec<ForeignItemId>,
-    }
-
-    impl<'hir> Visitor<'hir> for ModuleCollector<'hir> {
-        type NestedFilter = nested_filter::All;
-
-        fn nested_visit_map(&mut self) -> Self::Map {
-            self.tcx.hir()
-        }
-
-        fn visit_item(&mut self, item: &'hir Item<'hir>) {
-            self.items.push(item.item_id());
-            if let ItemKind::Mod(..) = item.kind {
-                // If this declares another module, do not recurse inside it.
-                self.submodules.push(item.def_id);
-            } else {
-                intravisit::walk_item(self, item)
-            }
-        }
-
-        fn visit_trait_item(&mut self, item: &'hir TraitItem<'hir>) {
-            self.trait_items.push(item.trait_item_id());
-            intravisit::walk_trait_item(self, item)
-        }
-
-        fn visit_impl_item(&mut self, item: &'hir ImplItem<'hir>) {
-            self.impl_items.push(item.impl_item_id());
-            intravisit::walk_impl_item(self, item)
-        }
-
-        fn visit_foreign_item(&mut self, item: &'hir ForeignItem<'hir>) {
-            self.foreign_items.push(item.foreign_item_id());
-            intravisit::walk_foreign_item(self, item)
-        }
-    }
 }
 
 pub(crate) fn hir_crate_items(tcx: TyCtxt<'_>, _: ()) -> ModuleItems {
-    let mut collector = CrateCollector {
-        tcx,
-        submodules: Vec::default(),
-        items: Vec::default(),
-        trait_items: Vec::default(),
-        impl_items: Vec::default(),
-        foreign_items: Vec::default(),
-    };
+    let mut collector = ItemCollector::new(tcx, true);
 
+    // A "crate collector" and "module collector" start at a
+    // module item (the former starts at the crate root) but only
+    // the former needs to collect it. ItemCollector does not do this for us.
+    collector.submodules.push(CRATE_DEF_ID);
     tcx.hir().walk_toplevel_module(&mut collector);
 
-    let CrateCollector { submodules, items, trait_items, impl_items, foreign_items, .. } =
-        collector;
+    let ItemCollector {
+        submodules,
+        items,
+        trait_items,
+        impl_items,
+        foreign_items,
+        body_owners,
+        ..
+    } = collector;
 
     return ModuleItems {
         submodules: submodules.into_boxed_slice(),
         items: items.into_boxed_slice(),
         trait_items: trait_items.into_boxed_slice(),
         impl_items: impl_items.into_boxed_slice(),
         foreign_items: foreign_items.into_boxed_slice(),
+        body_owners: body_owners.into_boxed_slice(),
     };
+}
 
-    struct CrateCollector<'tcx> {
-        tcx: TyCtxt<'tcx>,
-        submodules: Vec<LocalDefId>,
-        items: Vec<ItemId>,
-        trait_items: Vec<TraitItemId>,
-        impl_items: Vec<ImplItemId>,
-        foreign_items: Vec<ForeignItemId>,
+struct ItemCollector<'tcx> {
+    // When true, it collects all items in the create,
+    // otherwise it collects items in some module.
+    crate_collector: bool,
+    tcx: TyCtxt<'tcx>,
+    submodules: Vec<LocalDefId>,
+    items: Vec<ItemId>,
+    trait_items: Vec<TraitItemId>,
+    impl_items: Vec<ImplItemId>,
+    foreign_items: Vec<ForeignItemId>,
+    body_owners: Vec<LocalDefId>,
+}
+
+impl<'tcx> ItemCollector<'tcx> {
+    fn new(tcx: TyCtxt<'tcx>, crate_collector: bool) -> ItemCollector<'tcx> {
+        ItemCollector {
+            crate_collector,
+            tcx,
+            submodules: Vec::default(),
+            items: Vec::default(),
+            trait_items: Vec::default(),
+            impl_items: Vec::default(),
+            foreign_items: Vec::default(),
+            body_owners: Vec::default(),
+        }
     }
+}
+
+impl<'hir> Visitor<'hir> for ItemCollector<'hir> {
+    type NestedFilter = nested_filter::All;
 
-    impl<'hir> Visitor<'hir> for CrateCollector<'hir> {
-        type NestedFilter = nested_filter::All;
+    fn nested_visit_map(&mut self) -> Self::Map {
+        self.tcx.hir()
+    }
 
-        fn nested_visit_map(&mut self) -> Self::Map {
-            self.tcx.hir()
+    fn visit_item(&mut self, item: &'hir Item<'hir>) {
+        if associated_body(Node::Item(item)).is_some() {
+            self.body_owners.push(item.def_id);
         }
 
-        fn visit_item(&mut self, item: &'hir Item<'hir>) {
-            self.items.push(item.item_id());
+        self.items.push(item.item_id());
+
+        // Items that are modules are handled here instead of in visit_mod.
+        if let ItemKind::Mod(module) = &item.kind {
+            self.submodules.push(item.def_id);
+            // A module collector does not recurse inside nested modules.
+            if self.crate_collector {
+                intravisit::walk_mod(self, module, item.hir_id());
+            }
+        } else {
             intravisit::walk_item(self, item)
         }
+    }
 
-        fn visit_mod(&mut self, m: &'hir Mod<'hir>, _s: Span, n: HirId) {
-            self.submodules.push(n.owner);
-            intravisit::walk_mod(self, m, n);
-        }
+    fn visit_foreign_item(&mut self, item: &'hir ForeignItem<'hir>) {
+        self.foreign_items.push(item.foreign_item_id());
+        intravisit::walk_foreign_item(self, item)
+    }
+
+    fn visit_anon_const(&mut self, c: &'hir AnonConst) {
+        self.body_owners.push(self.tcx.hir().local_def_id(c.hir_id));
+        intravisit::walk_anon_const(self, c)
+    }
 
-        fn visit_foreign_item(&mut self, item: &'hir ForeignItem<'hir>) {
-            self.foreign_items.push(item.foreign_item_id());
-            intravisit::walk_foreign_item(self, item)
+    fn visit_expr(&mut self, ex: &'hir Expr<'hir>) {
+        if matches!(ex.kind, ExprKind::Closure { .. }) {
+            self.body_owners.push(self.tcx.hir().local_def_id(ex.hir_id));
         }
+        intravisit::walk_expr(self, ex)
+    }
 
-        fn visit_trait_item(&mut self, item: &'hir TraitItem<'hir>) {
-            self.trait_items.push(item.trait_item_id());
-            intravisit::walk_trait_item(self, item)
+    fn visit_trait_item(&mut self, item: &'hir TraitItem<'hir>) {
+        if associated_body(Node::TraitItem(item)).is_some() {
+            self.body_owners.push(item.def_id);
         }
 
-        fn visit_impl_item(&mut self, item: &'hir ImplItem<'hir>) {
-            self.impl_items.push(item.impl_item_id());
-            intravisit::walk_impl_item(self, item)
+        self.trait_items.push(item.trait_item_id());
+        intravisit::walk_trait_item(self, item)
+    }
+
+    fn visit_impl_item(&mut self, item: &'hir ImplItem<'hir>) {
+        if associated_body(Node::ImplItem(item)).is_some() {
+            self.body_owners.push(item.def_id);
         }
+
+        self.impl_items.push(item.impl_item_id());
+        intravisit::walk_impl_item(self, item)
     }
 }

compiler/rustc_middle/src/hir/mod.rs

@@ -44,6 +44,7 @@ pub struct ModuleItems {
     trait_items: Box<[TraitItemId]>,
     impl_items: Box<[ImplItemId]>,
     foreign_items: Box<[ForeignItemId]>,
+    body_owners: Box<[LocalDefId]>,
 }
 
 impl ModuleItems {

src/test/ui/asm/type-check-1.stderr

@@ -33,33 +33,6 @@ LL |         asm!("{}", sym x);
    |
    = help: `sym` operands must refer to either a function or a static
 
-error[E0308]: mismatched types
-  --> $DIR/type-check-1.rs:58:26
-   |
-LL |         asm!("{}", const 0f32);
-   |                          ^^^^ expected integer, found `f32`
-
-error[E0308]: mismatched types
-  --> $DIR/type-check-1.rs:60:26
-   |
-LL |         asm!("{}", const 0 as *mut u8);
-   |                          ^^^^^^^^^^^^ expected integer, found *-ptr
-   |
-   = note:     expected type `{integer}`
-           found raw pointer `*mut u8`
-
-error[E0308]: mismatched types
-  --> $DIR/type-check-1.rs:62:26
-   |
-LL |         asm!("{}", const &0);
-   |                          ^^ expected integer, found `&{integer}`
-   |
-help: consider removing the borrow
-   |
-LL -         asm!("{}", const &0);
-LL +         asm!("{}", const 0);
-   |
-
 error: invalid asm output
   --> $DIR/type-check-1.rs:15:29
    |
@@ -123,6 +96,33 @@ LL |         asm!("{}", inout(reg) v[..]);
    |
    = note: only integers, floats, SIMD vectors, pointers and function pointers can be used as arguments for inline assembly
 
+error[E0308]: mismatched types
+  --> $DIR/type-check-1.rs:58:26
+   |
+LL |         asm!("{}", const 0f32);
+   |                          ^^^^ expected integer, found `f32`
+
+error[E0308]: mismatched types
+  --> $DIR/type-check-1.rs:60:26
+   |
+LL |         asm!("{}", const 0 as *mut u8);
+   |                          ^^^^^^^^^^^^ expected integer, found *-ptr
+   |
+   = note:     expected type `{integer}`
+           found raw pointer `*mut u8`
+
+error[E0308]: mismatched types
+  --> $DIR/type-check-1.rs:62:26
+   |
+LL |         asm!("{}", const &0);
+   |                          ^^ expected integer, found `&{integer}`
+   |
+help: consider removing the borrow
+   |
+LL -         asm!("{}", const &0);
+LL +         asm!("{}", const 0);
+   |
+
 error[E0308]: mismatched types
   --> $DIR/type-check-1.rs:76:25
    |