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

@@ -88,7 +88,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?
@@ -102,6 +105,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_macros::extension;
 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};
@@ -124,6 +125,7 @@ struct LoweringContext<'a, 'hir> {
     is_in_dyn_type: bool,
 
     current_hir_id_owner: hir::OwnerId,
+    current_ast_id_owner: NodeId,
     item_local_id_counter: hir::ItemLocalId,
     trait_map: ItemLocalMap<Box<[TraitCandidate]>>,
 
@@ -162,6 +164,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             children: Vec::default(),
             contract_ensures: None,
             current_hir_id_owner: hir::CRATE_OWNER_ID,
+            current_ast_id_owner: CRATE_NODE_ID,
             item_local_id_counter: hir::ItemLocalId::ZERO,
             ident_and_label_to_local_id: Default::default(),
             #[cfg(debug_assertions)]
@@ -364,17 +367,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>>,
     }
 
@@ -385,23 +393,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;
         }
     }
 }
@@ -438,7 +461,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(),
@@ -512,7 +535,12 @@ 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.node_id_to_def_id.insert(node_id, def_id);
+        self.resolver
+            .owners
+            .get_mut(&self.current_ast_id_owner)
+            .unwrap()
+            .node_id_to_def_id
+            .insert(node_id, def_id);
 
         def_id
     }
@@ -527,16 +555,31 @@ 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.node_id_to_def_id.get(&node).copied()
+        self.resolver.owners[&self.current_ast_id_owner].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.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}) instead of in {:?} ({})",
+                            items.node_id_to_def_id.get(id),
+                            self.resolver.owners[&self.current_ast_id_owner].node_id_to_def_id[&self.current_ast_id_owner],
+                            self.current_ast_id_owner,
+                        )
+                    }
+                    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.owners[&node].node_id_to_def_id[&node] }
     }
 
     /// Freshen the `LoweringContext` and ready it to lower a nested item.
@@ -555,6 +598,7 @@ 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.current_ast_id_owner, owner);
         let current_ident_and_label_to_local_id =
             std::mem::take(&mut self.ident_and_label_to_local_id);
 
@@ -599,6 +643,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         self.item_local_id_counter = current_local_counter;
         self.impl_trait_defs = current_impl_trait_defs;
         self.impl_trait_bounds = current_impl_trait_bounds;
+        self.current_ast_id_owner = current_ast_owner;
 
         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,9 @@ 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
+    fn set_owner(&mut self, id: NodeId) -> 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,43 @@ impl<'a, 'b> InvocationCollector<'a, 'b> {
             };
         }
     }
+
+    fn with_owner<T>(&mut self, id: NodeId, f: impl FnOnce(&mut Self) -> T) -> T {
+        let old = if id != DUMMY_NODE_ID { self.cx.resolver.set_owner(id) } else { DUMMY_NODE_ID };
+        let val = f(self);
+        if id != DUMMY_NODE_ID {
+            assert_eq!(self.cx.resolver.set_owner(old), id);
+        }
+        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 +2284,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,11 @@ pub struct ResolverGlobalCtxt {
     pub stripped_cfg_items: Steal<Vec<StrippedCfgItem>>,
 }
 
+#[derive(Debug, Default)]
+pub struct PerOwnerResolverData {
+    pub node_id_to_def_id: NodeMap<LocalDefId>,
+}
+
 /// Resolutions that should only be used for lowering.
 /// This struct is meant to be consumed by lowering.
 #[derive(Debug)]
@@ -204,9 +209,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,