Record item-likes in ItemLowerer.

cjgillot · cjgillot · commit 6b099db18cbf · 2022-03-31T13:48:16.000+02:00
diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
@@ -5,7 +5,7 @@ use crate::{Arena, FnDeclKind};
 use rustc_ast::ptr::P;
 use rustc_ast::visit::AssocCtxt;
 use rustc_ast::*;
-use rustc_data_structures::fx::FxHashSet;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::sorted_map::SortedMap;
 use rustc_errors::struct_span_err;
 use rustc_hir as hir;
@@ -53,18 +53,22 @@ fn add_ty_alias_where_clause(
 }
 
 impl<'a, 'hir> ItemLowerer<'a, 'hir> {
-    fn make_lctx(&mut self) -> LoweringContext<'_, 'hir> {
-        LoweringContext {
+    fn with_lctx(
+        &mut self,
+        owner: NodeId,
+        f: impl FnOnce(&mut LoweringContext<'_, 'hir>) -> hir::OwnerNode<'hir>,
+    ) {
+        let mut lctx = LoweringContext {
             // Pseudo-globals.
             sess: &self.sess,
             resolver: self.resolver,
             nt_to_tokenstream: self.nt_to_tokenstream,
             arena: self.arena,
-            owners: self.owners,
 
             // HirId handling.
             bodies: Vec::new(),
             attrs: SortedMap::default(),
+            children: FxHashMap::default(),
             current_hir_id_owner: CRATE_DEF_ID,
             item_local_id_counter: hir::ItemLocalId::new(0),
             node_id_to_local_id: Default::default(),
@@ -87,6 +91,13 @@ impl<'a, 'hir> ItemLowerer<'a, 'hir> {
             allow_try_trait: Some([sym::try_trait_v2][..].into()),
             allow_gen_future: Some([sym::gen_future][..].into()),
             allow_into_future: Some([sym::into_future][..].into()),
+        };
+        lctx.with_hir_id_owner(owner, |lctx| f(lctx));
+
+        for (def_id, info) in lctx.children {
+            self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
+            debug_assert!(matches!(self.owners[def_id], hir::MaybeOwner::Phantom));
+            self.owners[def_id] = info;
         }
     }
 
@@ -109,23 +120,21 @@ impl<'a, 'hir> ItemLowerer<'a, 'hir> {
         self.owners[def_id]
     }
 
-    fn lower_crate(&mut self, c: &'a Crate) {
+    fn lower_crate(&mut self, c: &Crate) {
         debug_assert_eq!(self.resolver.local_def_id(CRATE_NODE_ID), CRATE_DEF_ID);
 
-        let mut lctx = self.make_lctx();
-        lctx.with_hir_id_owner(CRATE_NODE_ID, |lctx| {
+        self.with_lctx(CRATE_NODE_ID, |lctx| {
             let module = lctx.lower_mod(&c.items, c.spans.inner_span);
             lctx.lower_attrs(hir::CRATE_HIR_ID, &c.attrs);
             hir::OwnerNode::Crate(lctx.arena.alloc(module))
         })
     }
 
-    fn lower_item(&mut self, item: &'a Item) {
-        let mut lctx = self.make_lctx();
-        lctx.with_hir_id_owner(item.id, |lctx| hir::OwnerNode::Item(lctx.lower_item(item)))
+    fn lower_item(&mut self, item: &Item) {
+        self.with_lctx(item.id, |lctx| hir::OwnerNode::Item(lctx.lower_item(item)))
     }
 
-    fn lower_assoc_item(&mut self, item: &'a AssocItem, ctxt: AssocCtxt) {
+    fn lower_assoc_item(&mut self, item: &AssocItem, ctxt: AssocCtxt) {
         let def_id = self.resolver.local_def_id(item.id);
 
         let parent_id = {
@@ -135,43 +144,44 @@ impl<'a, 'hir> ItemLowerer<'a, 'hir> {
         };
 
         let parent_hir = self.lower_node(parent_id).unwrap().node().expect_item();
-        let mut lctx = self.make_lctx();
-
-        // Evaluate with the lifetimes in `params` in-scope.
-        // This is used to track which lifetimes have already been defined,
-        // and which need to be replicated when lowering an async fn.
-        match parent_hir.kind {
-            hir::ItemKind::Impl(hir::Impl { ref of_trait, ref generics, .. }) => {
-                lctx.is_in_trait_impl = of_trait.is_some();
-                lctx.in_scope_lifetimes = generics
-                    .params
-                    .iter()
-                    .filter(|param| matches!(param.kind, hir::GenericParamKind::Lifetime { .. }))
-                    .map(|param| param.name)
-                    .collect();
-            }
-            hir::ItemKind::Trait(_, _, ref generics, ..) => {
-                lctx.in_scope_lifetimes = generics
-                    .params
-                    .iter()
-                    .filter(|param| matches!(param.kind, hir::GenericParamKind::Lifetime { .. }))
-                    .map(|param| param.name)
-                    .collect();
-            }
-            _ => {}
-        };
+        self.with_lctx(item.id, |lctx| {
+            // Evaluate with the lifetimes in `params` in-scope.
+            // This is used to track which lifetimes have already been defined,
+            // and which need to be replicated when lowering an async fn.
+            match parent_hir.kind {
+                hir::ItemKind::Impl(hir::Impl { ref of_trait, ref generics, .. }) => {
+                    lctx.is_in_trait_impl = of_trait.is_some();
+                    lctx.in_scope_lifetimes = generics
+                        .params
+                        .iter()
+                        .filter(|param| {
+                            matches!(param.kind, hir::GenericParamKind::Lifetime { .. })
+                        })
+                        .map(|param| param.name)
+                        .collect();
+                }
+                hir::ItemKind::Trait(_, _, ref generics, ..) => {
+                    lctx.in_scope_lifetimes = generics
+                        .params
+                        .iter()
+                        .filter(|param| {
+                            matches!(param.kind, hir::GenericParamKind::Lifetime { .. })
+                        })
+                        .map(|param| param.name)
+                        .collect();
+                }
+                _ => {}
+            };
 
-        lctx.with_hir_id_owner(item.id, |lctx| match ctxt {
-            AssocCtxt::Trait => hir::OwnerNode::TraitItem(lctx.lower_trait_item(item)),
-            AssocCtxt::Impl => hir::OwnerNode::ImplItem(lctx.lower_impl_item(item)),
+            match ctxt {
+                AssocCtxt::Trait => hir::OwnerNode::TraitItem(lctx.lower_trait_item(item)),
+                AssocCtxt::Impl => hir::OwnerNode::ImplItem(lctx.lower_impl_item(item)),
+            }
         })
     }
 
-    fn lower_foreign_item(&mut self, item: &'a ForeignItem) {
-        let mut lctx = self.make_lctx();
-        lctx.with_hir_id_owner(item.id, |lctx| {
-            hir::OwnerNode::ForeignItem(lctx.lower_foreign_item(item))
-        })
+    fn lower_foreign_item(&mut self, item: &ForeignItem) {
+        self.with_lctx(item.id, |lctx| hir::OwnerNode::ForeignItem(lctx.lower_foreign_item(item)))
     }
 }
 
@@ -555,12 +565,11 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     let new_id = self.resolver.local_def_id(new_node_id);
                     let Some(res) = resolutions.next() else {
                         // Associate an HirId to both ids even if there is no resolution.
-                        self.owners.ensure_contains_elem(new_id, || hir::MaybeOwner::Phantom);
-                        let _old = std::mem::replace(
-                            &mut self.owners[new_id],
+                        let _old = self.children.insert(
+                            new_id,
                             hir::MaybeOwner::NonOwner(hir::HirId::make_owner(new_id)),
                         );
-                        debug_assert!(matches!(_old, hir::MaybeOwner::Phantom));
+                        debug_assert!(_old.is_none());
                         continue;
                     };
                     let ident = *ident;
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
@@ -85,7 +85,7 @@ mod path;
 
 rustc_hir::arena_types!(rustc_arena::declare_arena);
 
-struct LoweringContext<'a, 'hir> {
+struct LoweringContext<'a, 'hir: 'a> {
     /// Used to assign IDs to HIR nodes that do not directly correspond to AST nodes.
     sess: &'a Session,
 
@@ -99,12 +99,12 @@ struct LoweringContext<'a, 'hir> {
     /// Used to allocate HIR nodes.
     arena: &'hir Arena<'hir>,
 
-    /// The items being lowered are collected here.
-    owners: &'a mut IndexVec<LocalDefId, hir::MaybeOwner<&'hir hir::OwnerInfo<'hir>>>,
     /// Bodies inside the owner being lowered.
     bodies: Vec<(hir::ItemLocalId, &'hir hir::Body<'hir>)>,
     /// Attributes inside the owner being lowered.
     attrs: SortedMap<hir::ItemLocalId, &'hir [Attribute]>,
+    /// Collect items that were created by lowering the current owner.
+    children: FxHashMap<LocalDefId, hir::MaybeOwner<&'hir hir::OwnerInfo<'hir>>>,
 
     generator_kind: Option<hir::GeneratorKind>,
 
@@ -536,13 +536,15 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         self.current_hir_id_owner = current_owner;
         self.item_local_id_counter = current_local_counter;
 
-        self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
-        self.owners[def_id] = hir::MaybeOwner::Owner(self.arena.alloc(info));
+        let _old = self.children.insert(def_id, hir::MaybeOwner::Owner(info));
+        debug_assert!(_old.is_none())
     }
 
-    fn make_owner_info(&mut self, node: hir::OwnerNode<'hir>) -> hir::OwnerInfo<'hir> {
+    fn make_owner_info(&mut self, node: hir::OwnerNode<'hir>) -> &'hir hir::OwnerInfo<'hir> {
         let attrs = std::mem::take(&mut self.attrs);
         let mut bodies = std::mem::take(&mut self.bodies);
+        let local_id_to_def_id = std::mem::take(&mut self.local_id_to_def_id);
+        let trait_map = std::mem::take(&mut self.trait_map);
 
         #[cfg(debug_assertions)]
         for (id, attrs) in attrs.iter() {
@@ -562,7 +564,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             hash_without_bodies,
             nodes,
             bodies,
-            local_id_to_def_id: std::mem::take(&mut self.local_id_to_def_id),
+            local_id_to_def_id,
         };
         let attrs = {
             let mut hcx = self.resolver.create_stable_hashing_context();
@@ -572,7 +574,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             hir::AttributeMap { map: attrs, hash }
         };
 
-        hir::OwnerInfo { nodes, parenting, attrs, trait_map: std::mem::take(&mut self.trait_map) }
+        self.arena.alloc(hir::OwnerInfo { nodes, parenting, attrs, trait_map })
     }
 
     /// Hash the HIR node twice, one deep and one shallow hash.  This allows to differentiate
@@ -620,11 +622,8 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
                 assert_ne!(local_id, hir::ItemLocalId::new(0));
                 if let Some(def_id) = self.resolver.opt_local_def_id(ast_node_id) {
-                    self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
-                    if let o @ hir::MaybeOwner::Phantom = &mut self.owners[def_id] {
-                        // Do not override a `MaybeOwner::Owner` that may already here.
-                        *o = hir::MaybeOwner::NonOwner(hir_id);
-                    }
+                    // Do not override a `MaybeOwner::Owner` that may already here.
+                    self.children.entry(def_id).or_insert(hir::MaybeOwner::NonOwner(hir_id));
                     self.local_id_to_def_id.insert(local_id, def_id);
                 }
 
