Split the resolver tables into per-owner tables, starting with node_id_to_def_id

Author: oli-obk

compiler/rustc_ast_lowering/src/item.rs

@@ -87,7 +87,10 @@ impl<'a, 'hir> ItemLowerer<'a, 'hir> {
 
     #[instrument(level = "debug", skip(self, c))]
     fn lower_crate(&mut self, c: &Crate) {
-        debug_assert_eq!(self.resolver.node_id_to_def_id[&CRATE_NODE_ID], CRATE_DEF_ID);
+        debug_assert_eq!(
+            self.resolver.owners[&CRATE_NODE_ID].node_id_to_def_id[&CRATE_NODE_ID],
+            CRATE_DEF_ID
+        );
         self.with_lctx(CRATE_NODE_ID, |lctx| {
             let module = lctx.lower_mod(&c.items, &c.spans);
             // FIXME(jdonszelman): is dummy span ever a problem here?
@@ -101,6 +104,7 @@ impl<'a, 'hir> ItemLowerer<'a, 'hir> {
         self.with_lctx(item.id, |lctx| hir::OwnerNode::Item(lctx.lower_item(item)))
     }
 
+    #[instrument(level = "debug", skip(self))]
     fn lower_assoc_item(&mut self, item: &AssocItem, ctxt: AssocCtxt) {
         self.with_lctx(item.id, |lctx| lctx.lower_assoc_item(item, ctxt))
     }

compiler/rustc_ast_lowering/src/lib.rs

@@ -35,6 +35,7 @@
 #![doc(rust_logo)]
 #![feature(assert_matches)]
 #![feature(box_patterns)]
+#![feature(closure_track_caller)]
 #![feature(exact_size_is_empty)]
 #![feature(if_let_guard)]
 #![feature(let_chains)]
@@ -59,7 +60,7 @@ use rustc_hir::{
 };
 use rustc_index::{Idx, IndexSlice, IndexVec};
 use rustc_middle::span_bug;
-use rustc_middle::ty::{ResolverAstLowering, TyCtxt};
+use rustc_middle::ty::{PerOwnerResolverData, ResolverAstLowering, TyCtxt};
 use rustc_session::parse::{add_feature_diagnostics, feature_err};
 use rustc_span::symbol::{Ident, Symbol, kw, sym};
 use rustc_span::{DUMMY_SP, DesugaringKind, Span};
@@ -151,7 +152,10 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         Self {
             // Pseudo-globals.
             tcx,
-            resolver: PerOwnerResolver { general: resolver },
+            resolver: PerOwnerResolver {
+                general: resolver,
+                item: PerOwnerResolverData::new(DUMMY_NODE_ID),
+            },
             arena: tcx.hir_arena,
 
             // HirId handling.
@@ -201,6 +205,15 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
 pub(crate) struct PerOwnerResolver<'a> {
     pub general: &'a mut ResolverAstLowering,
+    pub item: PerOwnerResolverData,
+}
+
+impl<'a> std::ops::Deref for PerOwnerResolver<'a> {
+    type Target = PerOwnerResolverData;
+
+    fn deref(&self) -> &Self::Target {
+        &self.item
+    }
 }
 
 impl PerOwnerResolver<'_> {
@@ -368,17 +381,22 @@ enum AstOwner<'a> {
 }
 
 fn index_crate<'a>(
-    node_id_to_def_id: &NodeMap<LocalDefId>,
+    owners: &NodeMap<PerOwnerResolverData>,
     krate: &'a Crate,
 ) -> IndexVec<LocalDefId, AstOwner<'a>> {
-    let mut indexer = Indexer { node_id_to_def_id, index: IndexVec::new() };
+    let mut indexer = Indexer {
+        owners,
+        node_id_to_def_id: &owners[&CRATE_NODE_ID].node_id_to_def_id,
+        index: IndexVec::new(),
+    };
     *indexer.index.ensure_contains_elem(CRATE_DEF_ID, || AstOwner::NonOwner) =
         AstOwner::Crate(krate);
     visit::walk_crate(&mut indexer, krate);
     return indexer.index;
 
     struct Indexer<'s, 'a> {
         node_id_to_def_id: &'s NodeMap<LocalDefId>,
+        owners: &'s NodeMap<PerOwnerResolverData>,
         index: IndexVec<LocalDefId, AstOwner<'a>>,
     }
 
@@ -389,23 +407,38 @@ fn index_crate<'a>(
         }
 
         fn visit_item(&mut self, item: &'a ast::Item) {
+            let old = std::mem::replace(
+                &mut self.node_id_to_def_id,
+                &self.owners[&item.id].node_id_to_def_id,
+            );
             let def_id = self.node_id_to_def_id[&item.id];
             *self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner) = AstOwner::Item(item);
-            visit::walk_item(self, item)
+            visit::walk_item(self, item);
+            self.node_id_to_def_id = old;
         }
 
         fn visit_assoc_item(&mut self, item: &'a ast::AssocItem, ctxt: visit::AssocCtxt) {
+            let old = std::mem::replace(
+                &mut self.node_id_to_def_id,
+                &self.owners[&item.id].node_id_to_def_id,
+            );
             let def_id = self.node_id_to_def_id[&item.id];
             *self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner) =
                 AstOwner::AssocItem(item, ctxt);
             visit::walk_assoc_item(self, item, ctxt);
+            self.node_id_to_def_id = old;
         }
 
         fn visit_foreign_item(&mut self, item: &'a ast::ForeignItem) {
+            let old = std::mem::replace(
+                &mut self.node_id_to_def_id,
+                &self.owners[&item.id].node_id_to_def_id,
+            );
             let def_id = self.node_id_to_def_id[&item.id];
             *self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner) =
                 AstOwner::ForeignItem(item);
             visit::walk_item(self, item);
+            self.node_id_to_def_id = old;
         }
     }
 }
@@ -442,7 +475,7 @@ pub fn lower_to_hir(tcx: TyCtxt<'_>, (): ()) -> hir::Crate<'_> {
     tcx.ensure_done().get_lang_items(());
     let (mut resolver, krate) = tcx.resolver_for_lowering().steal();
 
-    let ast_index = index_crate(&resolver.node_id_to_def_id, &krate);
+    let ast_index = index_crate(&resolver.owners, &krate);
     let mut owners = IndexVec::from_fn_n(
         |_| hir::MaybeOwner::Phantom,
         tcx.definitions_untracked().def_index_count(),
@@ -517,7 +550,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             self.tcx.at(span).create_def(self.current_hir_id_owner.def_id, name, def_kind).def_id();
 
         debug!("create_def: def_id_to_node_id[{:?}] <-> {:?}", def_id, node_id);
-        self.resolver.general.node_id_to_def_id.insert(node_id, def_id);
+        self.resolver.item.node_id_to_def_id.insert(node_id, def_id);
 
         def_id
     }
@@ -532,16 +565,29 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
     /// Given the id of some node in the AST, finds the `LocalDefId` associated with it by the name
     /// resolver (if any).
     fn opt_local_def_id(&self, node: NodeId) -> Option<LocalDefId> {
-        self.resolver.general.node_id_to_def_id.get(&node).copied()
+        self.resolver.node_id_to_def_id.get(&node).copied()
     }
 
     fn local_def_id(&self, node: NodeId) -> LocalDefId {
-        self.opt_local_def_id(node).unwrap_or_else(|| panic!("no entry for node id: `{node:?}`"))
+        self.opt_local_def_id(node).unwrap_or_else(|| {
+            self.resolver.general.owners.items().any(|(id, items)| {
+                items.node_id_to_def_id.items().any(|(node_id, def_id)| {
+                    if *node_id == node {
+                        panic!(
+                            "{def_id:?} ({node_id}) was found in {:?} ({id})",
+                            items.node_id_to_def_id.get(id),
+                        )
+                    }
+                    false
+                })
+            });
+            panic!("no entry for node id: `{node:?}`");
+        })
     }
 
     /// Given the id of an owner node in the AST, returns the corresponding `OwnerId`.
     fn owner_id(&self, node: NodeId) -> hir::OwnerId {
-        hir::OwnerId { def_id: self.local_def_id(node) }
+        hir::OwnerId { def_id: self.resolver.general.owners[&node].node_id_to_def_id[&node] }
     }
 
     /// Freshen the `LoweringContext` and ready it to lower a nested item.
@@ -560,6 +606,10 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         let current_attrs = std::mem::take(&mut self.attrs);
         let current_bodies = std::mem::take(&mut self.bodies);
         let current_define_opaque = std::mem::take(&mut self.define_opaque);
+        let current_ast_owner = std::mem::replace(
+            &mut self.resolver.item,
+            self.resolver.general.owners.remove(&owner).unwrap(),
+        );
         let current_ident_and_label_to_local_id =
             std::mem::take(&mut self.ident_and_label_to_local_id);
 
@@ -605,6 +655,10 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         self.impl_trait_defs = current_impl_trait_defs;
         self.impl_trait_bounds = current_impl_trait_bounds;
 
+        let _prev_owner_data = std::mem::replace(&mut self.resolver.item, current_ast_owner);
+        #[cfg(debug_assertions)]
+        self.resolver.general.owners.insert(owner, _prev_owner_data);
+
         debug_assert!(!self.children.iter().any(|(id, _)| id == &owner_id.def_id));
         self.children.push((owner_id.def_id, hir::MaybeOwner::Owner(info)));
     }

compiler/rustc_expand/src/base.rs

@@ -1116,6 +1116,13 @@ pub trait ResolverExpand {
         trait_def_id: DefId,
         impl_def_id: LocalDefId,
     ) -> Result<Vec<(Ident, Option<Ident>)>, Indeterminate>;
+
+    /// Set a new owner and return the old one. Use only in the implementation
+    /// of `with_owner` and always call [Self::reset_owner] afterwards
+    fn set_owner(&mut self, id: NodeId) -> NodeId;
+
+    /// Switch back to the original owner.
+    fn reset_owner(&mut self, id: NodeId);
 }
 
 pub trait LintStoreExpand {

compiler/rustc_expand/src/expand.rs

@@ -4,16 +4,15 @@ use std::rc::Rc;
 use std::sync::Arc;
 use std::{iter, mem};
 
-use rustc_ast as ast;
 use rustc_ast::mut_visit::*;
 use rustc_ast::ptr::P;
 use rustc_ast::token::{self, Delimiter};
 use rustc_ast::tokenstream::TokenStream;
 use rustc_ast::visit::{self, AssocCtxt, Visitor, VisitorResult, try_visit, walk_list};
 use rustc_ast::{
-    AssocItemKind, AstNodeWrapper, AttrArgs, AttrStyle, AttrVec, ExprKind, ForeignItemKind,
-    HasAttrs, HasNodeId, Inline, ItemKind, MacStmtStyle, MetaItemInner, MetaItemKind, ModKind,
-    NodeId, PatKind, StmtKind, TyKind,
+    self as ast, AssocItemKind, AstNodeWrapper, AttrArgs, AttrStyle, AttrVec, CRATE_NODE_ID,
+    DUMMY_NODE_ID, ExprKind, ForeignItemKind, HasAttrs, HasNodeId, Inline, ItemKind, MacStmtStyle,
+    MetaItemInner, MetaItemKind, ModKind, NodeId, PatKind, StmtKind, TyKind,
 };
 use rustc_ast_pretty::pprust;
 use rustc_data_structures::flat_map_in_place::FlatMapInPlace;
@@ -2178,34 +2177,45 @@ impl<'a, 'b> InvocationCollector<'a, 'b> {
             };
         }
     }
+
+    fn with_owner<T>(&mut self, id: NodeId, f: impl FnOnce(&mut Self) -> T) -> T {
+        if id == DUMMY_NODE_ID {
+            f(self)
+        } else {
+            let old = self.cx.resolver.set_owner(id);
+            let val = f(self);
+            self.cx.resolver.reset_owner(old);
+            val
+        }
+    }
 }
 
 impl<'a, 'b> MutVisitor for InvocationCollector<'a, 'b> {
     fn flat_map_item(&mut self, node: P<ast::Item>) -> SmallVec<[P<ast::Item>; 1]> {
-        self.flat_map_node(node)
+        self.with_owner(node.id, |this| this.flat_map_node(node))
     }
 
     fn flat_map_assoc_item(
         &mut self,
         node: P<ast::AssocItem>,
         ctxt: AssocCtxt,
     ) -> SmallVec<[P<ast::AssocItem>; 1]> {
-        match ctxt {
-            AssocCtxt::Trait => self.flat_map_node(AstNodeWrapper::new(node, TraitItemTag)),
+        self.with_owner(node.id, |this| match ctxt {
+            AssocCtxt::Trait => this.flat_map_node(AstNodeWrapper::new(node, TraitItemTag)),
             AssocCtxt::Impl { of_trait: false } => {
-                self.flat_map_node(AstNodeWrapper::new(node, ImplItemTag))
+                this.flat_map_node(AstNodeWrapper::new(node, ImplItemTag))
             }
             AssocCtxt::Impl { of_trait: true } => {
-                self.flat_map_node(AstNodeWrapper::new(node, TraitImplItemTag))
+                this.flat_map_node(AstNodeWrapper::new(node, TraitImplItemTag))
             }
-        }
+        })
     }
 
     fn flat_map_foreign_item(
         &mut self,
         node: P<ast::ForeignItem>,
     ) -> SmallVec<[P<ast::ForeignItem>; 1]> {
-        self.flat_map_node(node)
+        self.with_owner(node.id, |this| this.flat_map_node(node))
     }
 
     fn flat_map_variant(&mut self, node: ast::Variant) -> SmallVec<[ast::Variant; 1]> {
@@ -2276,7 +2286,7 @@ impl<'a, 'b> MutVisitor for InvocationCollector<'a, 'b> {
     }
 
     fn visit_crate(&mut self, node: &mut ast::Crate) {
-        self.visit_node(node)
+        self.with_owner(CRATE_NODE_ID, |this| this.visit_node(node))
     }
 
     fn visit_ty(&mut self, node: &mut P<ast::Ty>) {

compiler/rustc_middle/src/ty/mod.rs

@@ -185,6 +185,19 @@ pub struct ResolverGlobalCtxt {
     pub stripped_cfg_items: Steal<Vec<StrippedCfgItem>>,
 }
 
+#[derive(Debug)]
+pub struct PerOwnerResolverData {
+    pub node_id_to_def_id: NodeMap<LocalDefId>,
+    /// The id of the owner
+    pub id: ast::NodeId,
+}
+
+impl PerOwnerResolverData {
+    pub fn new(id: ast::NodeId) -> Self {
+        Self { node_id_to_def_id: Default::default(), id }
+    }
+}
+
 /// Resolutions that should only be used for lowering.
 /// This struct is meant to be consumed by lowering.
 #[derive(Debug)]
@@ -202,9 +215,9 @@ pub struct ResolverAstLowering {
     /// Lifetime parameters that lowering will have to introduce.
     pub extra_lifetime_params_map: NodeMap<Vec<(Ident, ast::NodeId, LifetimeRes)>>,
 
-    pub next_node_id: ast::NodeId,
+    pub owners: NodeMap<PerOwnerResolverData>,
 
-    pub node_id_to_def_id: NodeMap<LocalDefId>,
+    pub next_node_id: ast::NodeId,
 
     pub trait_map: NodeMap<Vec<hir::TraitCandidate>>,
     /// List functions and methods for which lifetime elision was successful.

compiler/rustc_passes/src/lang_items.rs

@@ -291,7 +291,7 @@ impl<'ast, 'tcx> visit::Visitor<'ast> for LanguageItemCollector<'ast, 'tcx> {
 
         self.check_for_lang(
             target,
-            self.resolver.node_id_to_def_id[&i.id],
+            self.resolver.owners[&i.id].node_id_to_def_id[&i.id],
             &i.attrs,
             i.span,
             i.opt_generics(),
@@ -306,7 +306,7 @@ impl<'ast, 'tcx> visit::Visitor<'ast> for LanguageItemCollector<'ast, 'tcx> {
         for variant in &enum_definition.variants {
             self.check_for_lang(
                 Target::Variant,
-                self.resolver.node_id_to_def_id[&variant.id],
+                self.resolver.owners[&self.parent_item.unwrap().id].node_id_to_def_id[&variant.id],
                 &variant.attrs,
                 variant.span,
                 None,
@@ -348,7 +348,7 @@ impl<'ast, 'tcx> visit::Visitor<'ast> for LanguageItemCollector<'ast, 'tcx> {
 
         self.check_for_lang(
             target,
-            self.resolver.node_id_to_def_id[&i.id],
+            self.resolver.owners[&i.id].node_id_to_def_id[&i.id],
             &i.attrs,
             i.span,
             generics,

compiler/rustc_resolve/src/def_collector.rs

@@ -8,6 +8,7 @@ use rustc_expand::expand::AstFragment;
 use rustc_hir as hir;
 use rustc_hir::def::{CtorKind, CtorOf, DefKind};
 use rustc_hir::def_id::LocalDefId;
+use rustc_middle::ty::PerOwnerResolverData;
 use rustc_span::hygiene::LocalExpnId;
 use rustc_span::{Span, Symbol, sym};
 use tracing::debug;
@@ -20,8 +21,13 @@ pub(crate) fn collect_definitions(
     expansion: LocalExpnId,
 ) {
     let invocation_parent = resolver.invocation_parents[&expansion];
+    assert_eq!(resolver.current_owner.id, DUMMY_NODE_ID);
+    resolver.current_owner = resolver.owners.remove(&invocation_parent.owner).unwrap();
     let mut visitor = DefCollector { resolver, expansion, invocation_parent };
     fragment.visit_with(&mut visitor);
+    let tables =
+        mem::replace(&mut resolver.current_owner, PerOwnerResolverData::new(DUMMY_NODE_ID));
+    assert!(resolver.owners.insert(invocation_parent.owner, tables).is_none());
 }
 
 /// Creates `DefId`s for nodes in the AST.
@@ -44,7 +50,8 @@ impl<'a, 'ra, 'tcx> DefCollector<'a, 'ra, 'tcx> {
             "create_def(node_id={:?}, def_kind={:?}, parent_def={:?})",
             node_id, def_kind, parent_def
         );
-        self.resolver
+        let def_id = self
+            .resolver
             .create_def(
                 parent_def,
                 node_id,
@@ -53,7 +60,9 @@ impl<'a, 'ra, 'tcx> DefCollector<'a, 'ra, 'tcx> {
                 self.expansion.to_expn_id(),
                 span.with_parent(None),
             )
-            .def_id()
+            .def_id();
+        assert!(self.resolver.node_id_to_def_id.insert(node_id, def_id).is_none());
+        def_id
     }
 
     fn with_parent<F: FnOnce(&mut Self)>(&mut self, parent_def: LocalDefId, f: F) {
@@ -62,6 +71,25 @@ impl<'a, 'ra, 'tcx> DefCollector<'a, 'ra, 'tcx> {
         self.invocation_parent.parent_def = orig_parent_def;
     }
 
+    fn with_owner<F: FnOnce(&mut Self)>(&mut self, owner: NodeId, f: F) {
+        assert_ne!(owner, DUMMY_NODE_ID);
+        if owner == CRATE_NODE_ID {
+            // Special case: we always have an invocation parent of at least the crate root,
+            // even for visiting the crate root, which would then remove the crate root from the tables
+            // list twice and try to insert it twice afterwards.
+            assert_eq!(self.resolver.current_owner.id, CRATE_NODE_ID);
+            return f(self);
+        }
+        let tables =
+            self.resolver.owners.remove(&owner).unwrap_or_else(|| PerOwnerResolverData::new(owner));
+        let orig_owner = mem::replace(&mut self.resolver.current_owner, tables);
+        let orig_invoc_owner = mem::replace(&mut self.invocation_parent.owner, owner);
+        f(self);
+        let tables = mem::replace(&mut self.resolver.current_owner, orig_owner);
+        assert!(self.resolver.owners.insert(owner, tables).is_none());
+        assert_eq!(mem::replace(&mut self.invocation_parent.owner, orig_invoc_owner), owner);
+    }
+
     fn with_impl_trait<F: FnOnce(&mut Self)>(
         &mut self,
         impl_trait_context: ImplTraitContext,
@@ -152,31 +180,34 @@ impl<'a, 'ra, 'tcx> visit::Visitor<'a> for DefCollector<'a, 'ra, 'tcx> {
                 return self.visit_macro_invoc(i.id);
             }
         };
-        let def_id =
-            self.create_def(i.id, i.kind.ident().map(|ident| ident.name), def_kind, i.span);
+        self.with_owner(i.id, |this| {
+            let def_id =
+                this.create_def(i.id, i.kind.ident().map(|ident| ident.name), def_kind, i.span);
 
-        if let Some(macro_data) = opt_macro_data {
-            self.resolver.macro_map.insert(def_id.to_def_id(), macro_data);
-        }
+            if let Some(macro_data) = opt_macro_data {
+                this.resolver.macro_map.insert(def_id.to_def_id(), macro_data);
+            }
 
-        self.with_parent(def_id, |this| {
-            this.with_impl_trait(ImplTraitContext::Existential, |this| {
-                match i.kind {
-                    ItemKind::Struct(_, ref struct_def, _)
-                    | ItemKind::Union(_, ref struct_def, _) => {
-                        // If this is a unit or tuple-like struct, register the constructor.
-                        if let Some((ctor_kind, ctor_node_id)) = CtorKind::from_ast(struct_def) {
-                            this.create_def(
-                                ctor_node_id,
-                                None,
-                                DefKind::Ctor(CtorOf::Struct, ctor_kind),
-                                i.span,
-                            );
+            this.with_parent(def_id, |this| {
+                this.with_impl_trait(ImplTraitContext::Existential, |this| {
+                    match i.kind {
+                        ItemKind::Struct(_, ref struct_def, _)
+                        | ItemKind::Union(_, ref struct_def, _) => {
+                            // If this is a unit or tuple-like struct, register the constructor.
+                            if let Some((ctor_kind, ctor_node_id)) = CtorKind::from_ast(struct_def)
+                            {
+                                this.create_def(
+                                    ctor_node_id,
+                                    None,
+                                    DefKind::Ctor(CtorOf::Struct, ctor_kind),
+                                    i.span,
+                                );
+                            }
                         }
+                        _ => {}
                     }
-                    _ => {}
-                }
-                visit::walk_item(this, i);
+                    visit::walk_item(this, i);
+                })
             })
         });
     }
@@ -233,8 +264,10 @@ impl<'a, 'ra, 'tcx> visit::Visitor<'a> for DefCollector<'a, 'ra, 'tcx> {
     }
 
     fn visit_use_tree(&mut self, use_tree: &'a UseTree, id: NodeId, _nested: bool) {
-        self.create_def(id, None, DefKind::Use, use_tree.span);
-        visit::walk_use_tree(self, use_tree, id);
+        self.with_owner(id, |this| {
+            this.create_def(id, None, DefKind::Use, use_tree.span);
+            visit::walk_use_tree(this, use_tree, id);
+        })
     }
 
     fn visit_foreign_item(&mut self, fi: &'a ForeignItem) {
@@ -259,9 +292,11 @@ impl<'a, 'ra, 'tcx> visit::Visitor<'a> for DefCollector<'a, 'ra, 'tcx> {
             ForeignItemKind::MacCall(_) => return self.visit_macro_invoc(fi.id),
         };
 
-        let def = self.create_def(fi.id, Some(ident.name), def_kind, fi.span);
+        self.with_owner(fi.id, |this| {
+            let def = this.create_def(fi.id, Some(ident.name), def_kind, fi.span);
 
-        self.with_parent(def, |this| visit::walk_item(this, fi));
+            this.with_parent(def, |this| visit::walk_item(this, fi));
+        })
     }
 
     fn visit_variant(&mut self, v: &'a Variant) {
@@ -333,8 +368,10 @@ impl<'a, 'ra, 'tcx> visit::Visitor<'a> for DefCollector<'a, 'ra, 'tcx> {
             }
         };
 
-        let def = self.create_def(i.id, Some(ident.name), def_kind, i.span);
-        self.with_parent(def, |this| visit::walk_assoc_item(this, i, ctxt));
+        self.with_owner(i.id, |this| {
+            let def = this.create_def(i.id, Some(ident.name), def_kind, i.span);
+            this.with_parent(def, |this| visit::walk_assoc_item(this, i, ctxt));
+        })
     }
 
     fn visit_pat(&mut self, pat: &'a Pat) {
@@ -451,7 +488,7 @@ impl<'a, 'ra, 'tcx> visit::Visitor<'a> for DefCollector<'a, 'ra, 'tcx> {
         if krate.is_placeholder {
             self.visit_macro_invoc(krate.id)
         } else {
-            visit::walk_crate(self, krate)
+            self.with_owner(CRATE_NODE_ID, |this| visit::walk_crate(this, krate))
         }
     }
 

compiler/rustc_resolve/src/late.rs

@@ -765,7 +765,9 @@ impl<'ra: 'ast, 'ast, 'tcx> Visitor<'ast> for LateResolutionVisitor<'_, 'ast, 'r
         let prev = replace(&mut self.diag_metadata.current_item, Some(item));
         // Always report errors in items we just entered.
         let old_ignore = replace(&mut self.in_func_body, false);
-        self.with_lifetime_rib(LifetimeRibKind::Item, |this| this.resolve_item(item));
+        self.with_owner(item.id, |this| {
+            this.with_lifetime_rib(LifetimeRibKind::Item, |this| this.resolve_item(item))
+        });
         self.in_func_body = old_ignore;
         self.diag_metadata.current_item = prev;
     }
@@ -1680,6 +1682,16 @@ impl<'a, 'ast, 'ra: 'ast, 'tcx> LateResolutionVisitor<'a, 'ast, 'ra, 'tcx> {
         })
     }
 
+    #[instrument(level = "debug", skip(self, work))]
+    fn with_owner<T>(&mut self, owner: NodeId, work: impl FnOnce(&mut Self) -> T) -> T {
+        let old_owner =
+            std::mem::replace(&mut self.r.current_owner, self.r.owners.remove(&owner).unwrap());
+        let ret = work(self);
+        let prev = std::mem::replace(&mut self.r.current_owner, old_owner);
+        self.r.owners.insert(prev.id, prev);
+        ret
+    }
+
     #[instrument(level = "debug", skip(self, work))]
     fn with_lifetime_rib<T>(
         &mut self,
@@ -3098,7 +3110,9 @@ impl<'a, 'ast, 'ra: 'ast, 'tcx> LateResolutionVisitor<'a, 'ast, 'ra, 'tcx> {
             replace(&mut self.diag_metadata.current_trait_assoc_items, Some(trait_items));
 
         for item in trait_items {
-            self.resolve_trait_item(item);
+            self.with_owner(item.id, |this| {
+                this.resolve_trait_item(item);
+            })
         }
 
         self.diag_metadata.current_trait_assoc_items = trait_assoc_items;
@@ -3296,7 +3310,9 @@ impl<'a, 'ast, 'ra: 'ast, 'tcx> LateResolutionVisitor<'a, 'ast, 'ra, 'tcx> {
                                                 debug!("resolve_implementation with_self_rib_ns(ValueNS, ...)");
                                                 let mut seen_trait_items = Default::default();
                                                 for item in impl_items {
-                                                    this.resolve_impl_item(&**item, &mut seen_trait_items, trait_id);
+                                                    this.with_owner(item.id, |this| {
+                                                        this.resolve_impl_item(&**item, &mut seen_trait_items, trait_id);
+                                                    })
                                                 }
                                             });
                                         });
@@ -5227,18 +5243,21 @@ impl<'ast> Visitor<'ast> for ItemInfoCollector<'_, '_, '_> {
 
 impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
     pub(crate) fn late_resolve_crate(&mut self, krate: &Crate) {
-        visit::walk_crate(&mut ItemInfoCollector { r: self }, krate);
-        let mut late_resolution_visitor = LateResolutionVisitor::new(self);
-        late_resolution_visitor.resolve_doc_links(&krate.attrs, MaybeExported::Ok(CRATE_NODE_ID));
-        visit::walk_crate(&mut late_resolution_visitor, krate);
-        for (id, span) in late_resolution_visitor.diag_metadata.unused_labels.iter() {
-            self.lint_buffer.buffer_lint(
-                lint::builtin::UNUSED_LABELS,
-                *id,
-                *span,
-                BuiltinLintDiag::UnusedLabel,
-            );
-        }
+        self.with_owner(CRATE_NODE_ID, |this| {
+            visit::walk_crate(&mut ItemInfoCollector { r: this }, krate);
+            let mut late_resolution_visitor = LateResolutionVisitor::new(this);
+            late_resolution_visitor
+                .resolve_doc_links(&krate.attrs, MaybeExported::Ok(CRATE_NODE_ID));
+            visit::walk_crate(&mut late_resolution_visitor, krate);
+            for (id, span) in late_resolution_visitor.diag_metadata.unused_labels.iter() {
+                this.lint_buffer.buffer_lint(
+                    lint::builtin::UNUSED_LABELS,
+                    *id,
+                    *span,
+                    BuiltinLintDiag::UnusedLabel,
+                );
+            }
+        })
     }
 }
 

compiler/rustc_resolve/src/lib.rs

@@ -39,8 +39,8 @@ use rustc_arena::{DroplessArena, TypedArena};
 use rustc_ast::expand::StrippedCfgItem;
 use rustc_ast::node_id::NodeMap;
 use rustc_ast::{
-    self as ast, AngleBracketedArg, CRATE_NODE_ID, Crate, Expr, ExprKind, GenericArg, GenericArgs,
-    LitKind, NodeId, Path, attr,
+    self as ast, AngleBracketedArg, CRATE_NODE_ID, Crate, DUMMY_NODE_ID, Expr, ExprKind,
+    GenericArg, GenericArgs, LitKind, NodeId, Path, attr,
 };
 use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexMap, FxIndexSet};
 use rustc_data_structures::intern::Interned;
@@ -62,16 +62,16 @@ use rustc_middle::middle::privacy::EffectiveVisibilities;
 use rustc_middle::query::Providers;
 use rustc_middle::span_bug;
 use rustc_middle::ty::{
-    self, DelegationFnSig, Feed, MainDefinition, RegisteredTools, ResolverGlobalCtxt,
-    ResolverOutputs, TyCtxt, TyCtxtFeed,
+    self, DelegationFnSig, Feed, MainDefinition, PerOwnerResolverData, RegisteredTools,
+    ResolverGlobalCtxt, ResolverOutputs, TyCtxt, TyCtxtFeed,
 };
 use rustc_query_system::ich::StableHashingContext;
 use rustc_session::lint::builtin::PRIVATE_MACRO_USE;
 use rustc_session::lint::{BuiltinLintDiag, LintBuffer};
 use rustc_span::hygiene::{ExpnId, LocalExpnId, MacroKind, SyntaxContext, Transparency};
 use rustc_span::{DUMMY_SP, Ident, Span, Symbol, kw, sym};
 use smallvec::{SmallVec, smallvec};
-use tracing::debug;
+use tracing::{debug, instrument};
 
 type Res = def::Res<NodeId>;
 
@@ -176,13 +176,15 @@ struct InvocationParent {
     parent_def: LocalDefId,
     impl_trait_context: ImplTraitContext,
     in_attr: bool,
+    owner: NodeId,
 }
 
 impl InvocationParent {
     const ROOT: Self = Self {
         parent_def: CRATE_DEF_ID,
         impl_trait_context: ImplTraitContext::Existential,
         in_attr: false,
+        owner: CRATE_NODE_ID,
     };
 }
 
@@ -1060,6 +1062,12 @@ pub struct Resolver<'ra, 'tcx> {
     // Used for suggestions during error reporting.
     pat_span_map: NodeMap<Span>,
 
+    /// Preserves per owner data once the owner is finished resolving.
+    owners: NodeMap<PerOwnerResolverData>,
+
+    /// An entry of `owners` that gets taken out and reinserted whenever an owner is handled.
+    current_owner: PerOwnerResolverData,
+
     /// Resolutions for nodes that have a single resolution.
     partial_res_map: NodeMap<PartialRes>,
     /// Resolutions for import nodes, which have multiple resolutions in different namespaces.
@@ -1182,7 +1190,7 @@ pub struct Resolver<'ra, 'tcx> {
 
     next_node_id: NodeId,
 
-    node_id_to_def_id: NodeMap<Feed<'tcx, LocalDefId>>,
+    feed_for_node_id: NodeMap<Feed<'tcx, LocalDefId>>,
 
     /// Indices of unnamed struct or variant fields with unresolved attributes.
     placeholder_field_indices: FxHashMap<NodeId, usize>,
@@ -1229,6 +1237,22 @@ pub struct Resolver<'ra, 'tcx> {
     mods_with_parse_errors: FxHashSet<DefId>,
 }
 
+impl<'ra, 'tcx> std::ops::DerefMut for Resolver<'ra, 'tcx> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        assert_ne!(self.current_owner.id, DUMMY_NODE_ID);
+        &mut self.current_owner
+    }
+}
+
+impl<'ra, 'tcx> std::ops::Deref for Resolver<'ra, 'tcx> {
+    type Target = PerOwnerResolverData;
+
+    fn deref(&self) -> &Self::Target {
+        assert_ne!(self.current_owner.id, DUMMY_NODE_ID);
+        &self.current_owner
+    }
+}
+
 /// This provides memory for the rest of the crate. The `'ra` lifetime that is
 /// used by many types in this crate is an abbreviation of `ResolverArenas`.
 #[derive(Default)]
@@ -1316,7 +1340,7 @@ impl<'tcx> Resolver<'_, 'tcx> {
     }
 
     fn opt_feed(&self, node: NodeId) -> Option<Feed<'tcx, LocalDefId>> {
-        self.node_id_to_def_id.get(&node).copied()
+        self.feed_for_node_id.get(&node).copied()
     }
 
     fn feed(&self, node: NodeId) -> Feed<'tcx, LocalDefId> {
@@ -1338,12 +1362,12 @@ impl<'tcx> Resolver<'_, 'tcx> {
         span: Span,
     ) -> TyCtxtFeed<'tcx, LocalDefId> {
         assert!(
-            !self.node_id_to_def_id.contains_key(&node_id),
+            !self.feed_for_node_id.contains_key(&node_id),
             "adding a def for node-id {:?}, name {:?}, data {:?} but a previous def exists: {:?}",
             node_id,
             name,
             def_kind,
-            self.tcx.definitions_untracked().def_key(self.node_id_to_def_id[&node_id].key()),
+            self.tcx.definitions_untracked().def_key(self.feed_for_node_id[&node_id].key()),
         );
 
         // FIXME: remove `def_span` body, pass in the right spans here and call `tcx.at().create_def()`
@@ -1365,7 +1389,7 @@ impl<'tcx> Resolver<'_, 'tcx> {
         // we don't need a mapping from `NodeId` to `LocalDefId`.
         if node_id != ast::DUMMY_NODE_ID {
             debug!("create_def: def_id_to_node_id[{:?}] <-> {:?}", def_id, node_id);
-            self.node_id_to_def_id.insert(node_id, feed.downgrade());
+            self.feed_for_node_id.insert(node_id, feed.downgrade());
         }
 
         feed
@@ -1384,11 +1408,11 @@ impl<'tcx> Resolver<'_, 'tcx> {
     }
 
     /// This function is very slow, as it iterates over the entire
-    /// [Resolver::node_id_to_def_id] map just to find the [NodeId]
+    /// [Resolver::feed_for_node_id] map just to find the [NodeId]
     /// that corresponds to the given [LocalDefId]. Only use this in
     /// diagnostics code paths.
     fn def_id_to_node_id(&self, def_id: LocalDefId) -> NodeId {
-        self.node_id_to_def_id
+        self.feed_for_node_id
             .items()
             .filter(|(_, v)| v.key() == def_id)
             .map(|(k, _)| *k)
@@ -1427,12 +1451,16 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
             &mut Default::default(),
         );
 
-        let mut node_id_to_def_id = NodeMap::default();
+        let mut feed_for_node_id = NodeMap::default();
+        let mut owner_data = PerOwnerResolverData::new(CRATE_NODE_ID);
         let crate_feed = tcx.create_local_crate_def_id(crate_span);
 
         crate_feed.def_kind(DefKind::Mod);
         let crate_feed = crate_feed.downgrade();
-        node_id_to_def_id.insert(CRATE_NODE_ID, crate_feed);
+        owner_data.node_id_to_def_id.insert(CRATE_NODE_ID, crate_feed.key());
+        let mut owners = NodeMap::default();
+        owners.insert(CRATE_NODE_ID, owner_data);
+        feed_for_node_id.insert(CRATE_NODE_ID, crate_feed);
 
         let mut invocation_parents = FxHashMap::default();
         invocation_parents.insert(LocalExpnId::ROOT, InvocationParent::ROOT);
@@ -1476,6 +1504,8 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
             indeterminate_imports: Vec::new(),
 
             pat_span_map: Default::default(),
+            owners,
+            current_owner: PerOwnerResolverData::new(DUMMY_NODE_ID),
             partial_res_map: Default::default(),
             import_res_map: Default::default(),
             import_use_map: Default::default(),
@@ -1560,7 +1590,7 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
             containers_deriving_copy: Default::default(),
             lint_buffer: LintBuffer::default(),
             next_node_id: CRATE_NODE_ID,
-            node_id_to_def_id,
+            feed_for_node_id,
             placeholder_field_indices: Default::default(),
             invocation_parents,
             legacy_const_generic_args: Default::default(),
@@ -1653,8 +1683,12 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
             self.stripped_cfg_items
                 .into_iter()
                 .filter_map(|item| {
-                    let parent_module =
-                        self.node_id_to_def_id.get(&item.parent_module)?.key().to_def_id();
+                    let parent_module = self
+                        .owners
+                        .get(&item.parent_module)?
+                        .node_id_to_def_id
+                        .get(&item.parent_module)?
+                        .to_def_id();
                     Some(StrippedCfgItem { parent_module, ident: item.ident, cfg: item.cfg })
                 })
                 .collect(),
@@ -1678,18 +1712,14 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
             stripped_cfg_items,
         };
         let ast_lowering = ty::ResolverAstLowering {
+            owners: self.owners,
             legacy_const_generic_args: self.legacy_const_generic_args,
             partial_res_map: self.partial_res_map,
             import_res_map: self.import_res_map,
             label_res_map: self.label_res_map,
             lifetimes_res_map: self.lifetimes_res_map,
             extra_lifetime_params_map: self.extra_lifetime_params_map,
             next_node_id: self.next_node_id,
-            node_id_to_def_id: self
-                .node_id_to_def_id
-                .into_items()
-                .map(|(k, f)| (k, f.key()))
-                .collect(),
             trait_map: self.trait_map,
             lifetime_elision_allowed: self.lifetime_elision_allowed,
             lint_buffer: Steal::new(self.lint_buffer),
@@ -2264,6 +2294,19 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
         }
         self.main_def = Some(MainDefinition { res, is_import, span });
     }
+
+    #[instrument(level = "debug", skip(self, work))]
+    fn with_owner<T>(&mut self, owner: NodeId, work: impl FnOnce(&mut Self) -> T) -> T {
+        let old_owner =
+            std::mem::replace(&mut self.current_owner, self.owners.remove(&owner).unwrap());
+        let ret = work(self);
+        assert!(
+            self.owners
+                .insert(owner, std::mem::replace(&mut self.current_owner, old_owner))
+                .is_none()
+        );
+        ret
+    }
 }
 
 fn names_to_string(names: impl Iterator<Item = Symbol>) -> String {

compiler/rustc_resolve/src/macros.rs

@@ -6,7 +6,7 @@ use std::mem;
 use std::sync::Arc;
 
 use rustc_ast::expand::StrippedCfgItem;
-use rustc_ast::{self as ast, Crate, NodeId, attr};
+use rustc_ast::{self as ast, Crate, DUMMY_NODE_ID, NodeId, attr};
 use rustc_ast_pretty::pprust;
 use rustc_attr_parsing::{AttributeKind, StabilityLevel, find_attr};
 use rustc_data_structures::intern::Interned;
@@ -21,7 +21,7 @@ use rustc_expand::expand::{
 use rustc_hir::def::{self, DefKind, Namespace, NonMacroAttrKind};
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId};
 use rustc_middle::middle::stability;
-use rustc_middle::ty::{RegisteredTools, TyCtxt, Visibility};
+use rustc_middle::ty::{PerOwnerResolverData, RegisteredTools, TyCtxt, Visibility};
 use rustc_session::lint::BuiltinLintDiag;
 use rustc_session::lint::builtin::{
     LEGACY_DERIVE_HELPERS, OUT_OF_SCOPE_MACRO_CALLS, UNKNOWN_OR_MALFORMED_DIAGNOSTIC_ATTRIBUTES,
@@ -511,6 +511,35 @@ impl<'ra, 'tcx> ResolverExpand for Resolver<'ra, 'tcx> {
         });
         Ok(idents)
     }
+
+    fn set_owner(&mut self, id: NodeId) -> NodeId {
+        assert_ne!(id, DUMMY_NODE_ID);
+        let old = std::mem::replace(
+            &mut self.current_owner,
+            self.owners.remove(&id).unwrap_or_else(|| PerOwnerResolverData::new(id)),
+        );
+        let old_id = old.id;
+        if old.id == DUMMY_NODE_ID {
+            if cfg!(debug_assertions) {
+                let PerOwnerResolverData { node_id_to_def_id, id: _ } = old;
+                assert!(node_id_to_def_id.is_empty());
+            }
+        } else {
+            assert!(self.owners.insert(old.id, old).is_none());
+        }
+        old_id
+    }
+
+    fn reset_owner(&mut self, id: NodeId) {
+        let new = if id == DUMMY_NODE_ID {
+            PerOwnerResolverData::new(DUMMY_NODE_ID)
+        } else {
+            self.owners.remove(&id).unwrap()
+        };
+        let old = std::mem::replace(&mut self.current_owner, new);
+        assert_ne!(old.id, DUMMY_NODE_ID);
+        assert!(self.owners.insert(old.id, old).is_none());
+    }
 }
 
 impl<'ra, 'tcx> Resolver<'ra, 'tcx> {