Merge the BTreeMap in hir::Crate.

2 files changed, 46 insertions, 29 deletions

diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index 9f9d41c3f3d..852f696ef14 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/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(),
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index d4caba92416..540abe6e48e 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/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(),
@@ -558,12 +549,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 +563,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);