Avoid indexing the owners table all the time

oli-obk · oli-obk · commit a16a6f1a6139 · 2025-03-28T12:16:58.000Z
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
@@ -125,7 +125,7 @@ struct LoweringContext<'a, 'hir> {
     is_in_dyn_type: bool,
 
     current_hir_id_owner: hir::OwnerId,
-    current_ast_id_owner: NodeId,
+    current_ast_id_owner: PerOwnerResolverData,
     item_local_id_counter: hir::ItemLocalId,
     trait_map: ItemLocalMap<Box<[TraitCandidate]>>,
 
@@ -164,7 +164,10 @@ 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,
+            current_ast_id_owner: PerOwnerResolverData {
+                node_id_to_def_id: Default::default(),
+                id: DUMMY_NODE_ID,
+            },
             item_local_id_counter: hir::ItemLocalId::ZERO,
             ident_and_label_to_local_id: Default::default(),
             #[cfg(debug_assertions)]
@@ -535,12 +538,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
-            .owners
-            .get_mut(&self.current_ast_id_owner)
-            .unwrap()
-            .node_id_to_def_id
-            .insert(node_id, def_id);
+        self.current_ast_id_owner.node_id_to_def_id.insert(node_id, def_id);
 
         def_id
     }
@@ -555,7 +553,7 @@ 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.owners[&self.current_ast_id_owner].node_id_to_def_id.get(&node).copied()
+        self.current_ast_id_owner.node_id_to_def_id.get(&node).copied()
     }
 
     fn local_def_id(&self, node: NodeId) -> LocalDefId {
@@ -564,10 +562,8 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
                 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 {:?} ({})",
+                            "{def_id:?} ({node_id}) was found in {:?} ({id})",
                             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
@@ -598,7 +594,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.current_ast_id_owner, owner);
+        let current_ast_owner = std::mem::replace(
+            &mut self.current_ast_id_owner,
+            self.resolver.owners.remove(&owner).unwrap(),
+        );
         let current_ident_and_label_to_local_id =
             std::mem::take(&mut self.ident_and_label_to_local_id);
 
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
@@ -187,9 +187,17 @@ pub struct ResolverGlobalCtxt {
     pub stripped_cfg_items: Steal<Vec<StrippedCfgItem>>,
 }
 
-#[derive(Debug, Default)]
+#[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.
diff --git a/compiler/rustc_resolve/src/def_collector.rs b/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,10 +21,13 @@ pub(crate) fn collect_definitions(
     expansion: LocalExpnId,
 ) {
     let invocation_parent = resolver.invocation_parents[&expansion];
-    resolver.current_owner = invocation_parent.owner;
+    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);
-    resolver.current_owner = DUMMY_NODE_ID;
+    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.
@@ -68,12 +72,21 @@ impl<'a, 'ra, 'tcx> DefCollector<'a, 'ra, 'tcx> {
     }
 
     fn with_owner<F: FnOnce(&mut Self)>(&mut self, owner: NodeId, f: F) {
-        let orig_owner = mem::replace(&mut self.resolver.current_owner, owner);
-        let orig_invoc_owner = mem::replace(&mut self.invocation_parent.owner, owner);
         assert_ne!(owner, DUMMY_NODE_ID);
-        self.resolver.owners.entry(owner).or_default();
+        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);
-        assert_eq!(mem::replace(&mut self.resolver.current_owner, orig_owner), owner);
+        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);
     }
 
diff --git a/compiler/rustc_resolve/src/late.rs b/compiler/rustc_resolve/src/late.rs
@@ -1688,9 +1688,11 @@ 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, owner);
+        let old_owner =
+            std::mem::replace(&mut self.r.current_owner, self.r.owners.remove(&owner).unwrap());
         let ret = work(self);
-        self.r.current_owner = old_owner;
+        let prev = std::mem::replace(&mut self.r.current_owner, old_owner);
+        self.r.owners.insert(prev.id, prev);
         ret
     }
 
@@ -5239,20 +5241,21 @@ impl<'ast> Visitor<'ast> for ItemInfoCollector<'_, '_, '_> {
 
 impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
     pub(crate) fn late_resolve_crate(&mut self, krate: &Crate) {
-        self.current_owner = CRATE_NODE_ID;
-        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.current_owner = DUMMY_NODE_ID;
+        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,
+                );
+            }
+        })
     }
 }
 
diff --git a/compiler/rustc_resolve/src/lib.rs b/compiler/rustc_resolve/src/lib.rs
@@ -1064,8 +1064,8 @@ pub struct Resolver<'ra, 'tcx> {
     /// Preserves per owner data once the owner is finished resolving.
     owners: NodeMap<PerOwnerResolverData>,
 
-    /// Used to index into [Self::owners]. Refers to the owner of the nodes that are currently processed
-    current_owner: NodeId,
+    /// 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>,
@@ -1240,19 +1240,17 @@ pub struct Resolver<'ra, 'tcx> {
 
 impl<'ra, 'tcx> std::ops::DerefMut for Resolver<'ra, 'tcx> {
     fn deref_mut(&mut self) -> &mut Self::Target {
-        assert_ne!(self.current_owner, DUMMY_NODE_ID);
-        self.owners
-            .get_mut(&self.current_owner)
-            .unwrap_or_else(|| panic!("no entry for {}", self.current_owner))
+        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, DUMMY_NODE_ID);
-        &self.owners[&self.current_owner]
+        assert_ne!(self.current_owner.id, DUMMY_NODE_ID);
+        &self.current_owner
     }
 }
 
@@ -1445,7 +1443,7 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
         let mut def_id_to_node_id = IndexVec::default();
         assert_eq!(def_id_to_node_id.push(CRATE_NODE_ID), CRATE_DEF_ID);
         let mut feed_for_node_id = NodeMap::default();
-        let mut owner_data = PerOwnerResolverData::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);
@@ -1498,7 +1496,7 @@ impl<'ra, 'tcx> Resolver<'ra, 'tcx> {
 
             pat_span_map: Default::default(),
             owners,
-            current_owner: DUMMY_NODE_ID,
+            current_owner: PerOwnerResolverData::new(DUMMY_NODE_ID),
             partial_res_map: Default::default(),
             import_res_map: Default::default(),
             import_use_map: Default::default(),
@@ -2291,9 +2289,14 @@ impl<'ra, 'tcx> Resolver<'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.current_owner, owner);
+        let old_owner =
+            std::mem::replace(&mut self.current_owner, self.owners.remove(&owner).unwrap());
         let ret = work(self);
-        assert_eq!(std::mem::replace(&mut self.current_owner, old_owner), owner);
+        assert!(
+            self.owners
+                .insert(owner, std::mem::replace(&mut self.current_owner, old_owner))
+                .is_none()
+        );
         ret
     }
 }
diff --git a/compiler/rustc_resolve/src/macros.rs b/compiler/rustc_resolve/src/macros.rs
@@ -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,19 +511,34 @@ impl<'ra, 'tcx> ResolverExpand for Resolver<'ra, 'tcx> {
 
     fn set_owner(&mut self, id: NodeId) -> NodeId {
         assert_ne!(id, DUMMY_NODE_ID);
-        let old = std::mem::replace(&mut self.current_owner, id);
-        self.owners.entry(id).or_default();
-        old
+        let old = std::mem::replace(
+            &mut self.current_owner,
+            self.owners.remove(&id).unwrap_or_else(|| PerOwnerResolverData {
+                node_id_to_def_id: Default::default(),
+                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 old = std::mem::replace(&mut self.current_owner, id);
-        if cfg!(debug_assertions) {
-            if id != DUMMY_NODE_ID {
-                self.owners.get(&id).unwrap();
-            }
-            self.owners.get(&old).unwrap();
-        }
+        let new = if id == DUMMY_NODE_ID {
+            PerOwnerResolverData { node_id_to_def_id: Default::default(), id: 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());
     }
 }
 
