Auto merge of #138995 - oli-obk:split-resolver, r=petrochenkov

Split the node_id_to_def_id table into a per-owner table

*[View all comments](https://triagebot.infra.rust-lang.org/gh-comments/rust-lang/rust/pull/138995)*

My goal is to split all the resolver tables that get passed to act lowering into per-owner tables, so that all information that ast lowering needs from the resolver is separated by owners. This should allow us to fully split ast lowering to have one query invocation per owner that steal the individual resolver results for each owner.

part of rust-lang/rust-project-goals#620

Author: bors

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_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?
@@ -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>]>,
 
@@ -173,6 +174,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.
@@ -316,12 +318,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_def_id(&self, id: NodeId) -> LocalDefId {
+        self.owners[&id].def_id
     }
 
     fn lifetime_elision_allowed(&self, id: NodeId) -> bool {
@@ -489,20 +487,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_def_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_def_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_def_id(item.id);
             *self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner) =
                 AstOwner::ForeignItem(item);
             visit::walk_item(self, item);
@@ -668,12 +666,23 @@ 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 {
+                        let actual_owner = items.node_id_to_def_id.get(id);
+                        panic!("{def_id:?} ({node_id}) was found in {actual_owner:?} ({id})",)
+                    }
+                    false
+                })
+            });
+            panic!("no entry for node id: `{node:?}`");
+        })
     }
 
     fn get_partial_res(&self, id: NodeId) -> Option<PartialRes> {
@@ -685,7 +694,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.
@@ -710,6 +719,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);
@@ -745,6 +755,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;
@@ -1793,27 +1804,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_hash::{
 pub struct UnordItems<T, I: Iterator<Item = T>>(I);
 
 impl<T, I: Iterator<Item = T>> UnordItems<T, I> {
+    #[inline]
+    pub fn new(iter: I) -> UnordItems<T, I> {
+        UnordItems(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_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) -> PerOwnerResolverData {
+        PerOwnerResolverData { 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.

compiler/rustc_passes/src/lang_items.rs

@@ -295,7 +295,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].def_id,
             &i.attrs,
             i.span,
             i.opt_generics(),
@@ -309,7 +309,7 @@ impl<'ast, 'tcx> visit::Visitor<'ast> for LanguageItemCollector<'ast, 'tcx> {
     fn visit_variant(&mut self, variant: &'ast ast::Variant) {
         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,
@@ -346,13 +346,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],
-            &i.attrs,
-            i.span,
-            generics,
-        );
+        self.check_for_lang(target, self.resolver.owners[&i.id].def_id, &i.attrs, i.span, generics);
 
         visit::walk_assoc_item(self, i, ctxt);
     }

compiler/rustc_resolve/src/build_reduced_graph.rs

@@ -743,12 +743,13 @@ impl<'a, 'ra, 'tcx> DefCollector<'a, 'ra, 'tcx> {
             }
             ast::UseTreeKind::Nested { ref items, .. } => {
                 for &(ref tree, id) in items {
-                    let feed = self.create_def(id, None, DefKind::Use, use_tree.span());
-                    self.build_reduced_graph_for_use_tree(
-                        // This particular use tree
-                        tree, id, &prefix, true, false, // The whole `use` item
-                        item, vis, root_span, feed,
-                    );
+                    self.with_owner(id, None, DefKind::Use, use_tree.span(), |this, feed| {
+                        this.build_reduced_graph_for_use_tree(
+                            // This particular use tree
+                            tree, id, &prefix, true, false, // The whole `use` item
+                            item, vis, root_span, feed,
+                        )
+                    });
                 }
 
                 // Empty groups `a::b::{}` are turned into synthetic `self` imports

compiler/rustc_resolve/src/check_unused.rs

@@ -31,6 +31,7 @@ use rustc_data_structures::fx::{FxHashMap, FxIndexMap, FxIndexSet};
 use rustc_data_structures::unord::UnordSet;
 use rustc_errors::{DiagArgValue, Diagnostic, MultiSpan};
 use rustc_hir::def::{DefKind, Res};
+use rustc_hir::def_id::LocalDefId;
 use rustc_session::lint::builtin::{
     MACRO_USE_EXTERN_CRATE, UNUSED_EXTERN_CRATES, UNUSED_IMPORTS, UNUSED_QUALIFICATIONS,
 };
@@ -81,9 +82,8 @@ struct ExternCrateToLint {
 impl<'a, 'ra, 'tcx> UnusedImportCheckVisitor<'a, 'ra, 'tcx> {
     // We have information about whether `use` (import) items are actually
     // used now. If an import is not used at all, we signal a lint error.
-    fn check_import(&mut self, id: ast::NodeId) {
+    fn check_import(&mut self, id: ast::NodeId, def_id: LocalDefId) {
         let used = self.r.used_imports.contains(&id);
-        let def_id = self.r.local_def_id(id);
         if !used {
             if self.r.maybe_unused_trait_imports.contains(&def_id) {
                 // Check later.
@@ -102,7 +102,8 @@ impl<'a, 'ra, 'tcx> UnusedImportCheckVisitor<'a, 'ra, 'tcx> {
     }
 
     fn check_use_tree(&mut self, use_tree: &'a ast::UseTree, id: ast::NodeId) {
-        if self.r.effective_visibilities.is_exported(self.r.local_def_id(id)) {
+        let def_id = self.r.owner_def_id(id);
+        if self.r.effective_visibilities.is_exported(def_id) {
             self.check_import_as_underscore(use_tree, id);
             return;
         }
@@ -112,7 +113,7 @@ impl<'a, 'ra, 'tcx> UnusedImportCheckVisitor<'a, 'ra, 'tcx> {
                 self.unused_import(self.base_id).add(id);
             }
         } else {
-            self.check_import(id);
+            self.check_import(id, def_id);
         }
     }
 
@@ -212,7 +213,7 @@ impl<'a, 'ra, 'tcx> UnusedImportCheckVisitor<'a, 'ra, 'tcx> {
 
             let module = self
                 .r
-                .get_nearest_non_block_module(self.r.local_def_id(extern_crate.id).to_def_id());
+                .get_nearest_non_block_module(self.r.owner_def_id(extern_crate.id).to_def_id());
             if module.no_implicit_prelude {
                 // If the module has `no_implicit_prelude`, then we don't suggest
                 // replacing the extern crate with a use, as it would not be
@@ -494,7 +495,7 @@ impl Resolver<'_, '_> {
                 None
             } else {
                 let parent_module = visitor.r.get_nearest_non_block_module(
-                    visitor.r.local_def_id(unused.use_tree_id).to_def_id(),
+                    visitor.r.owner_def_id(unused.use_tree_id).to_def_id(),
                 );
                 match module_to_string(parent_module) {
                     Some(module)