Auto merge of #98203 - kckeiks:gather-body-owners-in-hir-item-queries…

…, r=cjgillot

gather body owners

Issue #96341

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
  |