Keep TyCtxtFeed around longer in the resolver

Author: oli-obk

compiler/rustc_middle/src/ty/context.rs

@@ -76,6 +76,7 @@ use std::cmp::Ordering;
 use std::fmt;
 use std::hash::{Hash, Hasher};
 use std::iter;
+use std::marker::PhantomData;
 use std::mem;
 use std::ops::{Bound, Deref};
 
@@ -515,6 +516,23 @@ pub struct TyCtxtFeed<'tcx, KEY: Copy> {
     key: KEY,
 }
 
+/// The same as `TyCtxtFeed`, but does not contain a `TyCtxt`.
+/// Use this to pass around when you have a `TyCtxt` elsewhere.
+/// Just an optimization to save space and not store hundreds of
+/// `TyCtxtFeed` in the resolver.
+#[derive(Copy, Clone)]
+pub struct Feed<'tcx, KEY: Copy> {
+    _tcx: PhantomData<TyCtxt<'tcx>>,
+    // Do not allow direct access, as downstream code must not mutate this field.
+    key: KEY,
+}
+
+impl<T: fmt::Debug + Copy> fmt::Debug for Feed<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.key.fmt(f)
+    }
+}
+
 impl<'tcx> TyCtxt<'tcx> {
     pub fn feed_unit_query(self) -> TyCtxtFeed<'tcx, ()> {
         TyCtxtFeed { tcx: self, key: () }
@@ -540,6 +558,23 @@ impl<'tcx, KEY: Copy> TyCtxtFeed<'tcx, KEY> {
     pub fn key(&self) -> KEY {
         self.key
     }
+
+    #[inline(always)]
+    pub fn downgrade(self) -> Feed<'tcx, KEY> {
+        Feed { _tcx: PhantomData, key: self.key }
+    }
+}
+
+impl<'tcx, KEY: Copy> Feed<'tcx, KEY> {
+    #[inline(always)]
+    pub fn key(&self) -> KEY {
+        self.key
+    }
+
+    #[inline(always)]
+    pub fn upgrade(self, tcx: TyCtxt<'tcx>) -> TyCtxtFeed<'tcx, KEY> {
+        TyCtxtFeed { tcx, key: self.key }
+    }
 }
 
 impl<'tcx> TyCtxtFeed<'tcx, LocalDefId> {

compiler/rustc_middle/src/ty/mod.rs

@@ -84,7 +84,8 @@ pub use self::consts::{
     Const, ConstData, ConstInt, ConstKind, Expr, ScalarInt, UnevaluatedConst, ValTree,
 };
 pub use self::context::{
-    tls, CtxtInterners, DeducedParamAttrs, FreeRegionInfo, GlobalCtxt, Lift, TyCtxt, TyCtxtFeed,
+    tls, CtxtInterners, DeducedParamAttrs, Feed, FreeRegionInfo, GlobalCtxt, Lift, TyCtxt,
+    TyCtxtFeed,
 };
 pub use self::instance::{Instance, InstanceDef, ShortInstance, UnusedGenericParams};
 pub use self::list::List;

compiler/rustc_resolve/src/def_collector.rs

@@ -4,6 +4,7 @@ use rustc_ast::*;
 use rustc_expand::expand::AstFragment;
 use rustc_hir::def::{CtorKind, CtorOf, DefKind};
 use rustc_hir::def_id::LocalDefId;
+use rustc_middle::ty::Feed;
 use rustc_span::hygiene::LocalExpnId;
 use rustc_span::symbol::{kw, sym, Symbol};
 use rustc_span::Span;
@@ -20,7 +21,7 @@ pub(crate) fn collect_definitions(
 /// Creates `DefId`s for nodes in the AST.
 struct DefCollector<'a, 'b, 'tcx> {
     resolver: &'a mut Resolver<'b, 'tcx>,
-    parent_def: LocalDefId,
+    parent_def: Feed<'tcx, LocalDefId>,
     impl_trait_context: ImplTraitContext,
     expansion: LocalExpnId,
 }
@@ -32,8 +33,8 @@ impl<'a, 'b, 'tcx> DefCollector<'a, 'b, 'tcx> {
         name: Symbol,
         def_kind: DefKind,
         span: Span,
-    ) -> LocalDefId {
-        let parent_def = self.parent_def;
+    ) -> Feed<'tcx, LocalDefId> {
+        let parent_def = self.parent_def.key();
         debug!(
             "create_def(node_id={:?}, def_kind={:?}, parent_def={:?})",
             node_id, def_kind, parent_def
@@ -47,10 +48,10 @@ impl<'a, 'b, 'tcx> DefCollector<'a, 'b, 'tcx> {
                 self.expansion.to_expn_id(),
                 span.with_parent(None),
             )
-            .def_id()
+            .downgrade()
     }
 
-    fn with_parent<F: FnOnce(&mut Self)>(&mut self, parent_def: LocalDefId, f: F) {
+    fn with_parent<F: FnOnce(&mut Self)>(&mut self, parent_def: Feed<'tcx, LocalDefId>, f: F) {
         let orig_parent_def = std::mem::replace(&mut self.parent_def, parent_def);
         f(self);
         self.parent_def = orig_parent_def;
@@ -148,7 +149,7 @@ impl<'a, 'b, 'tcx> visit::Visitor<'a> for DefCollector<'a, 'b, 'tcx> {
         let def_id = self.create_def(i.id, i.ident.name, def_kind, i.span);
 
         if let Some(macro_data) = opt_macro_data {
-            self.resolver.macro_map.insert(def_id.to_def_id(), macro_data);
+            self.resolver.macro_map.insert(def_id.key().to_def_id(), macro_data);
         }
 
         self.with_parent(def_id, |this| {

compiler/rustc_resolve/src/lib.rs

@@ -52,7 +52,7 @@ use rustc_middle::metadata::ModChild;
 use rustc_middle::middle::privacy::EffectiveVisibilities;
 use rustc_middle::query::Providers;
 use rustc_middle::span_bug;
-use rustc_middle::ty::{self, MainDefinition, RegisteredTools, TyCtxt, TyCtxtFeed};
+use rustc_middle::ty::{self, Feed, MainDefinition, RegisteredTools, TyCtxt, TyCtxtFeed};
 use rustc_middle::ty::{ResolverGlobalCtxt, ResolverOutputs};
 use rustc_query_system::ich::StableHashingContext;
 use rustc_session::lint::builtin::PRIVATE_MACRO_USE;
@@ -1106,7 +1106,7 @@ pub struct Resolver<'a, 'tcx> {
     /// When collecting definitions from an AST fragment produced by a macro invocation `ExpnId`
     /// we know what parent node that fragment should be attached to thanks to this table,
     /// and how the `impl Trait` fragments were introduced.
-    invocation_parents: FxHashMap<LocalExpnId, (LocalDefId, ImplTraitContext)>,
+    invocation_parents: FxHashMap<LocalExpnId, (Feed<'tcx, LocalDefId>, ImplTraitContext)>,
 
     /// Some way to know that we are in a *trait* impl in `visit_assoc_item`.
     /// FIXME: Replace with a more general AST map (together with some other fields).
@@ -1318,7 +1318,10 @@ impl<'a, 'tcx> Resolver<'a, 'tcx> {
         node_id_to_def_id.insert(CRATE_NODE_ID, CRATE_DEF_ID);
 
         let mut invocation_parents = FxHashMap::default();
-        invocation_parents.insert(LocalExpnId::ROOT, (CRATE_DEF_ID, ImplTraitContext::Existential));
+        invocation_parents.insert(
+            LocalExpnId::ROOT,
+            (tcx.feed_local_def_id(CRATE_DEF_ID).downgrade(), ImplTraitContext::Existential),
+        );
 
         let mut extern_prelude: FxHashMap<Ident, ExternPreludeEntry<'_>> = tcx
             .sess

compiler/rustc_resolve/src/macros.rs

@@ -174,7 +174,7 @@ impl<'a, 'tcx> ResolverExpand for Resolver<'a, 'tcx> {
     }
 
     fn invocation_parent(&self, id: LocalExpnId) -> LocalDefId {
-        self.invocation_parents[&id].0
+        self.invocation_parents[&id].0.key()
     }
 
     fn resolve_dollar_crates(&mut self) {
@@ -804,7 +804,7 @@ impl<'a, 'tcx> Resolver<'a, 'tcx> {
                         let node_id = self
                             .invocation_parents
                             .get(&parent_scope.expansion)
-                            .map_or(ast::CRATE_NODE_ID, |id| self.def_id_to_node_id[id.0]);
+                            .map_or(ast::CRATE_NODE_ID, |id| self.def_id_to_node_id[id.0.key()]);
                         self.lint_buffer.buffer_lint_with_diagnostic(
                             LEGACY_DERIVE_HELPERS,
                             node_id,