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

@@ -36,6 +36,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)]
@@ -60,7 +61,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};
@@ -152,7 +153,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.
@@ -202,6 +206,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<'_> {
@@ -369,17 +382,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>>,
     }
 
@@ -390,23 +408,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;
         }
     }
 }
@@ -443,7 +476,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(),
@@ -518,7 +551,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
     }
@@ -533,16 +566,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.
@@ -561,6 +607,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);
 
@@ -606,6 +656,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

@@ -187,6 +187,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)]
@@ -204,9 +217,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,