Auto merge of #83723 - cjgillot:ownernode, r=petrochenkov

Store all HIR owners in the same container

This replaces the previous storage in a BTreeMap for each of Item/ImplItem/TraitItem/ForeignItem.
This should allow for a more compact storage.

Based on #83114

Author: bors

compiler/rustc_ast_lowering/src/item.rs

@@ -82,15 +82,11 @@ impl<'a> Visitor<'a> for ItemLowerer<'a, '_, '_> {
         self.lctx.with_hir_id_owner(item.id, |lctx| match ctxt {
             AssocCtxt::Trait => {
                 let hir_item = lctx.lower_trait_item(item);
-                let id = hir_item.trait_item_id();
-                lctx.trait_items.insert(id, hir_item);
-                lctx.modules.entry(lctx.current_module).or_default().trait_items.insert(id);
+                lctx.insert_trait_item(hir_item);
             }
             AssocCtxt::Impl => {
                 let hir_item = lctx.lower_impl_item(item);
-                let id = hir_item.impl_item_id();
-                lctx.impl_items.insert(id, hir_item);
-                lctx.modules.entry(lctx.current_module).or_default().impl_items.insert(id);
+                lctx.insert_impl_item(hir_item);
             }
         });
 
@@ -101,9 +97,7 @@ impl<'a> Visitor<'a> for ItemLowerer<'a, '_, '_> {
         self.lctx.allocate_hir_id_counter(item.id);
         self.lctx.with_hir_id_owner(item.id, |lctx| {
             let hir_item = lctx.lower_foreign_item(item);
-            let id = hir_item.foreign_item_id();
-            lctx.foreign_items.insert(id, hir_item);
-            lctx.modules.entry(lctx.current_module).or_default().foreign_items.insert(id);
+            lctx.insert_foreign_item(hir_item);
         });
 
         visit::walk_foreign_item(self, item);
@@ -123,7 +117,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
     ) -> T {
         let old_len = self.in_scope_lifetimes.len();
 
-        let parent_generics = match self.items.get(&parent_hir_id).unwrap().kind {
+        let parent_generics = match self.owners[parent_hir_id.def_id].unwrap().expect_item().kind {
             hir::ItemKind::Impl(hir::Impl { ref generics, .. })
             | hir::ItemKind::Trait(_, _, ref generics, ..) => generics.params,
             _ => &[],
@@ -224,7 +218,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 let hir_id = self.lower_node_id(i.id);
                 self.lower_attrs(hir_id, &i.attrs);
                 let body = P(self.lower_mac_args(body));
-                self.exported_macros.push(hir::MacroDef {
+                self.insert_macro_def(hir::MacroDef {
                     ident,
                     vis,
                     def_id: hir_id.expect_owner(),

compiler/rustc_ast_lowering/src/lib.rs

@@ -101,13 +101,8 @@ struct LoweringContext<'a, 'hir: 'a> {
     arena: &'hir Arena<'hir>,
 
     /// The items being lowered are collected here.
-    items: BTreeMap<hir::ItemId, hir::Item<'hir>>,
-
-    trait_items: BTreeMap<hir::TraitItemId, hir::TraitItem<'hir>>,
-    impl_items: BTreeMap<hir::ImplItemId, hir::ImplItem<'hir>>,
-    foreign_items: BTreeMap<hir::ForeignItemId, hir::ForeignItem<'hir>>,
+    owners: IndexVec<LocalDefId, Option<hir::OwnerNode<'hir>>>,
     bodies: BTreeMap<hir::BodyId, hir::Body<'hir>>,
-    exported_macros: Vec<hir::MacroDef<'hir>>,
     non_exported_macro_attrs: Vec<ast::Attribute>,
 
     trait_impls: BTreeMap<DefId, Vec<LocalDefId>>,
@@ -330,15 +325,11 @@ pub fn lower_crate<'a, 'hir>(
         resolver,
         nt_to_tokenstream,
         arena,
-        items: BTreeMap::new(),
-        trait_items: BTreeMap::new(),
-        impl_items: BTreeMap::new(),
-        foreign_items: BTreeMap::new(),
+        owners: IndexVec::default(),
         bodies: BTreeMap::new(),
         trait_impls: BTreeMap::new(),
         modules: BTreeMap::new(),
         attrs: BTreeMap::default(),
-        exported_macros: Vec::new(),
         non_exported_macro_attrs: Vec::new(),
         catch_scopes: Vec::new(),
         loop_scopes: Vec::new(),
@@ -521,8 +512,11 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         visit::walk_crate(&mut MiscCollector { lctx: &mut self }, c);
         visit::walk_crate(&mut item::ItemLowerer { lctx: &mut self }, c);
 
-        let module = self.lower_mod(&c.items, c.span);
+        let module = self.arena.alloc(self.lower_mod(&c.items, c.span));
         self.lower_attrs(hir::CRATE_HIR_ID, &c.attrs);
+        self.owners.ensure_contains_elem(CRATE_DEF_ID, || None);
+        self.owners[CRATE_DEF_ID] = Some(hir::OwnerNode::Crate(module));
+
         let body_ids = body_ids(&self.bodies);
         let proc_macros =
             c.proc_macros.iter().map(|id| self.node_id_to_hir_id[*id].unwrap()).collect();
@@ -557,13 +551,8 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         }
 
         let krate = hir::Crate {
-            item: module,
-            exported_macros: self.arena.alloc_from_iter(self.exported_macros),
             non_exported_macro_attrs: self.arena.alloc_from_iter(self.non_exported_macro_attrs),
-            items: self.items,
-            trait_items: self.trait_items,
-            impl_items: self.impl_items,
-            foreign_items: self.foreign_items,
+            owners: self.owners,
             bodies: self.bodies,
             body_ids,
             trait_impls: self.trait_impls,
@@ -576,12 +565,48 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
     }
 
     fn insert_item(&mut self, item: hir::Item<'hir>) -> hir::ItemId {
-        let id = hir::ItemId { def_id: item.def_id };
-        self.items.insert(id, item);
+        let id = item.item_id();
+        let item = self.arena.alloc(item);
+        self.owners.ensure_contains_elem(id.def_id, || None);
+        self.owners[id.def_id] = Some(hir::OwnerNode::Item(item));
         self.modules.entry(self.current_module).or_default().items.insert(id);
         id
     }
 
+    fn insert_foreign_item(&mut self, item: hir::ForeignItem<'hir>) -> hir::ForeignItemId {
+        let id = item.foreign_item_id();
+        let item = self.arena.alloc(item);
+        self.owners.ensure_contains_elem(id.def_id, || None);
+        self.owners[id.def_id] = Some(hir::OwnerNode::ForeignItem(item));
+        self.modules.entry(self.current_module).or_default().foreign_items.insert(id);
+        id
+    }
+
+    fn insert_impl_item(&mut self, item: hir::ImplItem<'hir>) -> hir::ImplItemId {
+        let id = item.impl_item_id();
+        let item = self.arena.alloc(item);
+        self.owners.ensure_contains_elem(id.def_id, || None);
+        self.owners[id.def_id] = Some(hir::OwnerNode::ImplItem(item));
+        self.modules.entry(self.current_module).or_default().impl_items.insert(id);
+        id
+    }
+
+    fn insert_trait_item(&mut self, item: hir::TraitItem<'hir>) -> hir::TraitItemId {
+        let id = item.trait_item_id();
+        let item = self.arena.alloc(item);
+        self.owners.ensure_contains_elem(id.def_id, || None);
+        self.owners[id.def_id] = Some(hir::OwnerNode::TraitItem(item));
+        self.modules.entry(self.current_module).or_default().trait_items.insert(id);
+        id
+    }
+
+    fn insert_macro_def(&mut self, item: hir::MacroDef<'hir>) {
+        let def_id = item.def_id;
+        let item = self.arena.alloc(item);
+        self.owners.ensure_contains_elem(def_id, || None);
+        self.owners[def_id] = Some(hir::OwnerNode::MacroDef(item));
+    }
+
     fn allocate_hir_id_counter(&mut self, owner: NodeId) -> hir::HirId {
         // Set up the counter if needed.
         self.item_local_id_counters.entry(owner).or_insert(0);

compiler/rustc_hir/src/arena.rs

@@ -29,11 +29,14 @@ macro_rules! arena_types {
             [] fn_decl: rustc_hir::FnDecl<$tcx>,
             [] foreign_item: rustc_hir::ForeignItem<$tcx>,
             [few] foreign_item_ref: rustc_hir::ForeignItemRef<$tcx>,
+            [] impl_item: rustc_hir::ImplItem<$tcx>,
             [] impl_item_ref: rustc_hir::ImplItemRef<$tcx>,
+            [] item: rustc_hir::Item<$tcx>,
             [few] inline_asm: rustc_hir::InlineAsm<$tcx>,
             [few] llvm_inline_asm: rustc_hir::LlvmInlineAsm<$tcx>,
             [] local: rustc_hir::Local<$tcx>,
             [few] macro_def: rustc_hir::MacroDef<$tcx>,
+            [few] mod_: rustc_hir::Mod<$tcx>,
             [] param: rustc_hir::Param<$tcx>,
             [] pat: rustc_hir::Pat<$tcx>,
             [] path: rustc_hir::Path<$tcx>,
@@ -42,6 +45,7 @@ macro_rules! arena_types {
             [] qpath: rustc_hir::QPath<$tcx>,
             [] stmt: rustc_hir::Stmt<$tcx>,
             [] field_def: rustc_hir::FieldDef<$tcx>,
+            [] trait_item: rustc_hir::TraitItem<$tcx>,
             [] trait_item_ref: rustc_hir::TraitItemRef,
             [] ty: rustc_hir::Ty<$tcx>,
             [] type_binding: rustc_hir::TypeBinding<$tcx>,