Split the node_id_to_def_id table into a per-owner table

compiler/rustc_ast_lowering/src/delegation.rs

@@ -122,8 +122,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let span = self.lower_span(delegation.path.segments.last().unwrap().ident.span);
 
         // Delegation can be unresolved in illegal places such as function bodies in extern blocks (see #151356)
-        let sig_id = if let Some(delegation_info) =
-            self.resolver.delegation_info(self.local_def_id(item_id))
+        let sig_id = if let Some(delegation_info) = self.resolver.delegation_info(self.owner.def_id)
         {
             self.get_sig_id(delegation_info.resolution_node, span)
         } else {
@@ -143,8 +142,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
                 let (param_count, c_variadic) = self.param_count(sig_id);
 
-                let mut generics =
-                    self.uplift_delegation_generics(delegation, sig_id, item_id, is_method);
+                let mut generics = self.uplift_delegation_generics(delegation, sig_id, is_method);
 
                 let body_id = self.lower_delegation_body(
                     delegation,

compiler/rustc_ast_lowering/src/delegation/generics.rs

@@ -231,11 +231,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
         &mut self,
         delegation: &Delegation,
         sig_id: DefId,
-        item_id: NodeId,
         is_method: bool,
     ) -> GenericsGenerationResults<'hir> {
-        let delegation_parent_kind =
-            self.tcx.def_kind(self.tcx.local_parent(self.local_def_id(item_id)));
+        let delegation_parent_kind = self.tcx.def_kind(self.tcx.local_parent(self.owner.def_id));
 
         let segments = &delegation.path.segments;
         let len = segments.len();

compiler/rustc_ast_lowering/src/item.rs

@@ -98,7 +98,7 @@ impl<'hir> ItemLowerer<'_, 'hir> {
             match node {
                 AstOwner::NonOwner => {}
                 AstOwner::Crate(c) => {
-                    assert_eq!(self.resolver.local_def_id(CRATE_NODE_ID), CRATE_DEF_ID);
+                    assert_eq!(self.resolver.owner_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?
@@ -445,7 +445,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                             ty,
                             ImplTraitContext::OpaqueTy {
                                 origin: hir::OpaqueTyOrigin::TyAlias {
-                                    parent: this.local_def_id(id),
+                                    parent: this.owner.def_id,
                                     in_assoc_ty: false,
                                 },
                             },
@@ -598,7 +598,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             ItemKind::MacroDef(ident, MacroDef { body, macro_rules, eii_declaration: _ }) => {
                 let ident = self.lower_ident(*ident);
                 let body = Box::new(self.lower_delim_args(body));
-                let def_id = self.local_def_id(id);
+                let def_id = self.owner.def_id;
                 let def_kind = self.tcx.def_kind(def_id);
                 let DefKind::Macro(macro_kinds) = def_kind else {
                     unreachable!(
@@ -1282,7 +1282,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                                     ty,
                                     ImplTraitContext::OpaqueTy {
                                         origin: hir::OpaqueTyOrigin::TyAlias {
-                                            parent: this.local_def_id(i.id),
+                                            parent: this.owner.def_id,
                                             in_assoc_ty: true,
                                         },
                                     },

compiler/rustc_ast_lowering/src/lib.rs

@@ -61,7 +61,7 @@ use rustc_index::{Idx, IndexSlice, IndexVec};
 use rustc_macros::extension;
 use rustc_middle::hir::{self as mid_hir};
 use rustc_middle::span_bug;
-use rustc_middle::ty::{DelegationInfo, ResolverAstLowering, TyCtxt};
+use rustc_middle::ty::{DelegationInfo, PerOwnerResolverData, ResolverAstLowering, TyCtxt};
 use rustc_session::errors::add_feature_diagnostics;
 use rustc_span::symbol::{Ident, Symbol, kw, sym};
 use rustc_span::{DUMMY_SP, DesugaringKind, Span};
@@ -126,6 +126,7 @@ struct LoweringContext<'a, 'hir> {
     is_in_dyn_type: bool,
 
     current_hir_id_owner: hir::OwnerId,
+    owner: &'a PerOwnerResolverData,
     item_local_id_counter: hir::ItemLocalId,
     trait_map: ItemLocalMap<&'hir [TraitCandidate<'hir>]>,
 
@@ -168,6 +169,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             tcx,
             resolver,
             current_disambiguator: Default::default(),
+            owner: &resolver.owners[&CRATE_NODE_ID],
             arena: tcx.hir_arena,
 
             // HirId handling.
@@ -310,12 +312,8 @@ impl<'tcx> ResolverAstLowering<'tcx> {
         self.delegation_infos.get(&id)
     }
 
-    fn opt_local_def_id(&self, id: NodeId) -> Option<LocalDefId> {
-        self.node_id_to_def_id.get(&id).copied()
-    }
-
-    fn local_def_id(&self, id: NodeId) -> LocalDefId {
-        self.opt_local_def_id(id).expect("must have def_id")
+    fn owner_id(&self, id: NodeId) -> LocalDefId {
+        self.owners[&id].def_id
     }
 
     fn lifetime_elision_allowed(&self, id: NodeId) -> bool {
@@ -483,20 +481,20 @@ fn index_crate<'a, 'b>(
         }
 
         fn visit_item(&mut self, item: &'a ast::Item) {
-            let def_id = self.resolver.local_def_id(item.id);
+            let def_id = self.resolver.owner_id(item.id);
             *self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner) = AstOwner::Item(item);
             visit::walk_item(self, item)
         }
 
         fn visit_assoc_item(&mut self, item: &'a ast::AssocItem, ctxt: visit::AssocCtxt) {
-            let def_id = self.resolver.local_def_id(item.id);
+            let def_id = self.resolver.owner_id(item.id);
             *self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner) =
                 AstOwner::AssocItem(item, ctxt);
             visit::walk_assoc_item(self, item, ctxt);
         }
 
         fn visit_foreign_item(&mut self, item: &'a ast::ForeignItem) {
-            let def_id = self.resolver.local_def_id(item.id);
+            let def_id = self.resolver.owner_id(item.id);
             *self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner) =
                 AstOwner::ForeignItem(item);
             visit::walk_item(self, item);
@@ -662,12 +660,25 @@ impl<'hir> LoweringContext<'_, 'hir> {
     fn opt_local_def_id(&self, node: NodeId) -> Option<LocalDefId> {
         self.node_id_to_def_id
             .get(&node)
-            .or_else(|| self.resolver.node_id_to_def_id.get(&node))
+            .or_else(|| self.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})",
+                            items.node_id_to_def_id.get(id),
+                        )
+                    }
+                    false
+                })
+            });
+            panic!("no entry for node id: `{node:?}`");
+        })
     }
 
     fn get_partial_res(&self, id: NodeId) -> Option<PartialRes> {
@@ -679,7 +690,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
     /// 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].def_id }
     }
 
     /// Freshen the `LoweringContext` and ready it to lower a nested item.
@@ -704,6 +715,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             .unwrap_or_else(|| PerParentDisambiguatorState::new(def_id));
 
         let disambiguator = std::mem::replace(&mut self.current_disambiguator, new_disambig);
+        let current_ast_owner = std::mem::replace(&mut self.owner, &self.resolver.owners[&owner]);
         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);
@@ -739,6 +751,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let info = self.make_owner_info(item);
 
         self.current_disambiguator = disambiguator;
+        self.owner = current_ast_owner;
         self.attrs = current_attrs;
         self.bodies = current_bodies;
         self.define_opaque = current_define_opaque;
@@ -1787,27 +1800,27 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
         let output = match coro {
             Some(coro) => {
-                let fn_def_id = self.local_def_id(fn_node_id);
+                let fn_def_id = self.owner.def_id;
                 self.lower_coroutine_fn_ret_ty(&decl.output, fn_def_id, coro, kind)
             }
             None => match &decl.output {
                 FnRetTy::Ty(ty) => {
                     let itctx = match kind {
                         FnDeclKind::Fn | FnDeclKind::Inherent => ImplTraitContext::OpaqueTy {
                             origin: hir::OpaqueTyOrigin::FnReturn {
-                                parent: self.local_def_id(fn_node_id),
+                                parent: self.owner.def_id,
                                 in_trait_or_impl: None,
                             },
                         },
                         FnDeclKind::Trait => ImplTraitContext::OpaqueTy {
                             origin: hir::OpaqueTyOrigin::FnReturn {
-                                parent: self.local_def_id(fn_node_id),
+                                parent: self.owner.def_id,
                                 in_trait_or_impl: Some(hir::RpitContext::Trait),
                             },
                         },
                         FnDeclKind::Impl => ImplTraitContext::OpaqueTy {
                             origin: hir::OpaqueTyOrigin::FnReturn {
-                                parent: self.local_def_id(fn_node_id),
+                                parent: self.owner.def_id,
                                 in_trait_or_impl: Some(hir::RpitContext::TraitImpl),
                             },
                         },

compiler/rustc_data_structures/src/unord.rs

@@ -34,6 +34,11 @@ use crate::stable_hasher::{
 pub struct UnordItems<T, I: Iterator<Item = T>>(I);
 
 impl<T, I: Iterator<Item = T>> UnordItems<T, I> {
+    #[inline]
+    pub fn wrap(iter: I) -> Self {
+        Self(iter)
+    }
+
     #[inline]
     pub fn map<U, F: Fn(T) -> U>(self, f: F) -> UnordItems<U, impl Iterator<Item = U>> {
         UnordItems(self.0.map(f))
@@ -62,6 +67,14 @@ impl<T, I: Iterator<Item = T>> UnordItems<T, I> {
         UnordItems(self.0.filter_map(f))
     }
 
+    #[inline]
+    pub fn chain(
+        self,
+        other: UnordItems<T, impl Iterator<Item = T>>,
+    ) -> UnordItems<T, impl Iterator<Item = T>> {
+        UnordItems(self.0.chain(other.0))
+    }
+
     #[inline]
     pub fn max(self) -> Option<T>
     where
@@ -102,10 +115,10 @@ impl<T, I: Iterator<Item = T>> UnordItems<T, I> {
     #[inline]
     pub fn flat_map<U, F, O>(self, f: F) -> UnordItems<O, impl Iterator<Item = O>>
     where
-        U: IntoIterator<Item = O>,
-        F: Fn(T) -> U,
+        U: Iterator<Item = O>,
+        F: Fn(T) -> UnordItems<O, U>,
     {
-        UnordItems(self.0.flat_map(f))
+        UnordItems(self.0.flat_map(move |x| f(x).0))
     }
 
     pub fn collect<C: From<UnordItems<T, I>>>(self) -> C {

compiler/rustc_expand/src/base.rs

@@ -1180,6 +1180,13 @@ pub trait ResolverExpand {
     /// Mark the scope as having a compile error so that error for lookup in this scope
     /// should be suppressed
     fn mark_scope_with_compile_error(&mut self, parent_node: NodeId);
+
+    /// 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

@@ -8,9 +8,9 @@ use rustc_ast::tokenstream::TokenStream;
 use rustc_ast::visit::{self, AssocCtxt, Visitor, VisitorResult, try_visit, walk_list};
 use rustc_ast::{
     self as ast, AssocItemKind, AstNodeWrapper, AttrArgs, AttrItemKind, AttrStyle, AttrVec,
-    DUMMY_NODE_ID, DelegationSuffixes, EarlyParsedAttribute, ExprKind, ForeignItemKind, HasAttrs,
-    HasNodeId, Inline, ItemKind, MacStmtStyle, MetaItemInner, MetaItemKind, ModKind, NodeId,
-    PatKind, StmtKind, TyKind, token,
+    CRATE_NODE_ID, DUMMY_NODE_ID, DelegationSuffixes, EarlyParsedAttribute, ExprKind,
+    ForeignItemKind, HasAttrs, HasNodeId, Inline, ItemKind, MacStmtStyle, MetaItemInner,
+    MetaItemKind, ModKind, NodeId, PatKind, StmtKind, TyKind, token,
 };
 use rustc_ast_pretty::pprust;
 use rustc_attr_parsing::parser::AllowExprMetavar;
@@ -2474,34 +2474,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: Box<ast::Item>) -> SmallVec<[Box<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: Box<ast::AssocItem>,
         ctxt: AssocCtxt,
     ) -> SmallVec<[Box<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: Box<ast::ForeignItem>,
     ) -> SmallVec<[Box<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]> {
@@ -2572,7 +2583,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 ast::Ty) {

compiler/rustc_middle/src/ty/mod.rs

@@ -198,6 +198,21 @@ pub struct ResolverGlobalCtxt {
     pub stripped_cfg_items: Vec<StrippedCfgItem>,
 }
 
+#[derive(Debug)]
+pub struct PerOwnerResolverData {
+    pub node_id_to_def_id: NodeMap<LocalDefId>,
+    /// The id of the owner
+    pub id: ast::NodeId,
+    /// The `DefId` of the owner, can't be found in `node_id_to_def_id`.
+    pub def_id: LocalDefId,
+}
+
+impl PerOwnerResolverData {
+    pub fn new(id: ast::NodeId, def_id: LocalDefId) -> Self {
+        Self { node_id_to_def_id: Default::default(), id, def_id }
+    }
+}
+
 /// Resolutions that should only be used for lowering.
 /// This struct is meant to be consumed by lowering.
 #[derive(Debug)]
@@ -215,7 +230,7 @@ pub struct ResolverAstLowering<'tcx> {
 
     pub next_node_id: ast::NodeId,
 
-    pub node_id_to_def_id: NodeMap<LocalDefId>,
+    pub owners: NodeMap<PerOwnerResolverData>,
 
     pub trait_map: NodeMap<&'tcx [hir::TraitCandidate<'tcx>]>,
     /// List functions and methods for which lifetime elision was successful.