Index the crate before creating the TyCtxt.

cjgillot · cjgillot · commit 3f183065a37e · 2022-04-30T15:26:02.000+02:00
diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
@@ -1,18 +1,19 @@
 use super::ResolverAstLoweringExt;
-use super::{AstOwner, ImplTraitContext, ImplTraitPosition};
 use super::{FnDeclKind, LoweringContext, ParamMode};
+use super::{ImplTraitContext, ImplTraitPosition};
 
 use rustc_ast::ptr::P;
 use rustc_ast::visit::AssocCtxt;
 use rustc_ast::*;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sorted_map::SortedMap;
+use rustc_data_structures::steal::Steal;
 use rustc_errors::struct_span_err;
 use rustc_hir as hir;
 use rustc_hir::def::{DefKind, Res};
 use rustc_hir::def_id::{LocalDefId, CRATE_DEF_ID};
 use rustc_index::vec::{Idx, IndexVec};
-use rustc_middle::ty::{ResolverOutputs, TyCtxt};
+use rustc_middle::ty::{AstOwner, ResolverOutputs, TyCtxt};
 use rustc_session::utils::NtToTokenstream;
 use rustc_span::source_map::DesugaringKind;
 use rustc_span::symbol::{kw, sym, Ident};
@@ -27,7 +28,7 @@ pub(super) struct ItemLowerer<'a, 'hir> {
     pub(super) tcx: TyCtxt<'hir>,
     pub(super) resolver: &'hir ResolverOutputs,
     pub(super) nt_to_tokenstream: NtToTokenstream,
-    pub(super) ast_index: &'a IndexVec<LocalDefId, AstOwner<'a>>,
+    pub(super) ast_index: &'a IndexVec<LocalDefId, Steal<AstOwner>>,
     pub(super) owners: &'a mut IndexVec<LocalDefId, hir::MaybeOwner<&'hir hir::OwnerInfo<'hir>>>,
 }
 
@@ -107,13 +108,13 @@ impl<'a, 'hir> ItemLowerer<'a, 'hir> {
     ) -> hir::MaybeOwner<&'hir hir::OwnerInfo<'hir>> {
         self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
         if let hir::MaybeOwner::Phantom = self.owners[def_id] {
-            let node = self.ast_index[def_id];
+            let node = self.ast_index[def_id].steal();
             match node {
-                AstOwner::NonOwner => {}
-                AstOwner::Crate(c) => self.lower_crate(c),
-                AstOwner::Item(item) => self.lower_item(item),
-                AstOwner::AssocItem(item, ctxt) => self.lower_assoc_item(item, ctxt),
-                AstOwner::ForeignItem(item) => self.lower_foreign_item(item),
+                AstOwner::NonOwner | AstOwner::Synthetic(..) => {}
+                AstOwner::Crate(c) => self.lower_crate(&c),
+                AstOwner::Item(item) => self.lower_item(&item),
+                AstOwner::AssocItem(item, ctxt) => self.lower_assoc_item(&item, ctxt),
+                AstOwner::ForeignItem(item) => self.lower_foreign_item(&item),
             }
         }
 
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
@@ -46,8 +46,10 @@ use rustc_ast_pretty::pprust;
 use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use rustc_data_structures::owning_ref::OwningRef;
 use rustc_data_structures::sorted_map::SortedMap;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
+use rustc_data_structures::steal::Steal;
 use rustc_data_structures::sync::Lrc;
 use rustc_errors::struct_span_err;
 use rustc_hir as hir;
@@ -56,7 +58,7 @@ use rustc_hir::def_id::{LocalDefId, CRATE_DEF_ID};
 use rustc_hir::definitions::DefPathData;
 use rustc_hir::{ConstArg, GenericArg, ItemLocalId, ParamName, TraitCandidate};
 use rustc_index::vec::{Idx, IndexVec};
-use rustc_middle::ty::{ResolverOutputs, TyCtxt};
+use rustc_middle::ty::{AstOwner, ResolverOutputs, TyCtxt};
 use rustc_session::parse::feature_err;
 use rustc_session::utils::{FlattenNonterminals, NtToTokenstream};
 use rustc_session::Session;
@@ -341,54 +343,80 @@ impl FnDeclKind {
     }
 }
 
-#[derive(Copy, Clone)]
-enum AstOwner<'a> {
-    NonOwner,
-    Crate(&'a ast::Crate),
-    Item(&'a ast::Item),
-    AssocItem(&'a ast::AssocItem, visit::AssocCtxt),
-    ForeignItem(&'a ast::ForeignItem),
-}
-
-fn index_crate<'a>(
+pub fn index_crate(
     node_id_to_def_id: &FxHashMap<NodeId, LocalDefId>,
-    krate: &'a Crate,
-) -> IndexVec<LocalDefId, AstOwner<'a>> {
-    let mut indexer = Indexer { node_id_to_def_id, index: IndexVec::new() };
-    indexer.index.ensure_contains_elem(CRATE_DEF_ID, || AstOwner::NonOwner);
-    indexer.index[CRATE_DEF_ID] = AstOwner::Crate(krate);
-    visit::walk_crate(&mut indexer, krate);
+    krate: Lrc<Crate>,
+) -> IndexVec<LocalDefId, Steal<AstOwner>> {
+    let mut indexer =
+        Indexer { node_id_to_def_id, krate: OwningRef::new(krate.clone()), index: IndexVec::new() };
+    indexer.index.ensure_contains_elem(CRATE_DEF_ID, || Steal::new(AstOwner::NonOwner));
+    indexer.index[CRATE_DEF_ID] = Steal::new(AstOwner::Crate(krate.clone()));
+    visit::walk_crate(&mut indexer, &krate);
     return indexer.index;
 
-    struct Indexer<'s, 'a> {
+    struct Indexer<'s> {
         node_id_to_def_id: &'s FxHashMap<NodeId, LocalDefId>,
-        index: IndexVec<LocalDefId, AstOwner<'a>>,
+        krate: OwningRef<Lrc<Crate>, Crate>,
+        index: IndexVec<LocalDefId, Steal<AstOwner>>,
+    }
+
+    impl Indexer<'_> {
+        fn visit_item_id_use_tree(&mut self, tree: &UseTree, parent: LocalDefId) {
+            match tree.kind {
+                UseTreeKind::Glob => {}
+                UseTreeKind::Simple(_, id1, id2) => {
+                    for id in &[id1, id2] {
+                        let def_id = self.node_id_to_def_id[id];
+                        self.index.ensure_contains_elem(def_id, || Steal::new(AstOwner::NonOwner));
+                        self.index[def_id] = Steal::new(AstOwner::Synthetic(parent));
+                    }
+                }
+                UseTreeKind::Nested(ref nested_vec) => {
+                    for &(ref nested, id) in nested_vec {
+                        let def_id = self.node_id_to_def_id[&id];
+                        self.index.ensure_contains_elem(def_id, || Steal::new(AstOwner::NonOwner));
+                        self.index[def_id] = Steal::new(AstOwner::Synthetic(parent));
+
+                        self.visit_item_id_use_tree(nested, def_id);
+                    }
+                }
+            }
+        }
     }
 
-    impl<'a> visit::Visitor<'a> for Indexer<'_, 'a> {
+    impl<'a> visit::Visitor<'a> for Indexer<'_> {
         fn visit_attribute(&mut self, _: &'a Attribute) {
             // We do not want to lower expressions that appear in attributes,
             // as they are not accessible to the rest of the HIR.
         }
 
         fn visit_item(&mut self, item: &'a ast::Item) {
             let def_id = self.node_id_to_def_id[&item.id];
-            self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner);
-            self.index[def_id] = AstOwner::Item(item);
+            // SAFETY: the visitor guarantees the item ref comes from krate.
+            let item_ref = self.krate.clone().map(|_| unsafe { &*(item as *const _) });
+            self.index.ensure_contains_elem(def_id, || Steal::new(AstOwner::NonOwner));
+            self.index[def_id] = Steal::new(AstOwner::Item(item_ref));
+            if let ItemKind::Use(ref use_tree) = item.kind {
+                self.visit_item_id_use_tree(use_tree, def_id);
+            }
             visit::walk_item(self, item)
         }
 
         fn visit_assoc_item(&mut self, item: &'a ast::AssocItem, ctxt: visit::AssocCtxt) {
             let def_id = self.node_id_to_def_id[&item.id];
-            self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner);
-            self.index[def_id] = AstOwner::AssocItem(item, ctxt);
+            // SAFETY: the visitor guarantees the item ref comes from krate.
+            let item_ref = self.krate.clone().map(|_| unsafe { &*(item as *const _) });
+            self.index.ensure_contains_elem(def_id, || Steal::new(AstOwner::NonOwner));
+            self.index[def_id] = Steal::new(AstOwner::AssocItem(item_ref, ctxt));
             visit::walk_assoc_item(self, item, ctxt);
         }
 
         fn visit_foreign_item(&mut self, item: &'a ast::ForeignItem) {
             let def_id = self.node_id_to_def_id[&item.id];
-            self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner);
-            self.index[def_id] = AstOwner::ForeignItem(item);
+            // SAFETY: the visitor guarantees the item ref comes from krate.
+            let item_ref = self.krate.clone().map(|_| unsafe { &*(item as *const _) });
+            self.index.ensure_contains_elem(def_id, || Steal::new(AstOwner::NonOwner));
+            self.index[def_id] = Steal::new(AstOwner::ForeignItem(item_ref));
             visit::walk_foreign_item(self, item);
         }
     }
@@ -419,13 +447,12 @@ fn compute_hir_hash(
 
 pub fn lower_crate<'hir>(
     tcx: TyCtxt<'hir>,
-    krate: &Crate,
     nt_to_tokenstream: NtToTokenstream,
 ) -> hir::Crate<'hir> {
     let _prof_timer = tcx.sess.prof.verbose_generic_activity("hir_lowering");
 
     let resolver = tcx.resolutions(());
-    let ast_index = index_crate(&resolver.node_id_to_def_id, krate);
+    let ast_index = &tcx.untracked_crate;
 
     let mut owners = IndexVec::from_fn_n(
         |_| hir::MaybeOwner::Phantom,
@@ -437,7 +464,7 @@ pub fn lower_crate<'hir>(
             tcx,
             resolver,
             nt_to_tokenstream,
-            ast_index: &ast_index,
+            ast_index: ast_index,
             owners: &mut owners,
         }
         .lower_node(def_id);
diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
@@ -492,13 +492,9 @@ pub fn configure_and_expand(
 
 fn lower_to_hir<'tcx>(tcx: TyCtxt<'tcx>, (): ()) -> Crate<'tcx> {
     let sess = tcx.sess;
-    let krate = tcx.untracked_crate.steal();
 
     // Lower AST to HIR.
-    let hir_crate = rustc_ast_lowering::lower_crate(tcx, &krate, rustc_parse::nt_to_tokenstream);
-
-    // Drop AST to free memory
-    sess.time("drop_ast", || std::mem::drop(krate));
+    let hir_crate = rustc_ast_lowering::lower_crate(tcx, rustc_parse::nt_to_tokenstream);
 
     // Discard hygiene data, which isn't required after lowering to HIR.
     if !sess.opts.debugging_opts.keep_hygiene_data {
@@ -835,6 +831,7 @@ pub fn create_global_ctxt<'tcx>(
     dep_graph.assert_ignored();
 
     let (definitions, cstore, resolver_outputs) = BoxedResolver::to_resolver_outputs(resolver);
+    let krate = rustc_ast_lowering::index_crate(&resolver_outputs.node_id_to_def_id, krate);
 
     let sess = &compiler.session();
     let query_result_on_disk_cache = rustc_incremental::load_query_result_cache(sess);
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
@@ -1029,7 +1029,7 @@ pub struct GlobalCtxt<'tcx> {
 
     /// Output of the resolver.
     pub(crate) untracked_resolutions: ty::ResolverOutputs,
-    pub untracked_crate: Steal<Lrc<ast::Crate>>,
+    pub untracked_crate: IndexVec<LocalDefId, Steal<ty::AstOwner>>,
 
     /// This provides access to the incremental compilation on-disk cache for query results.
     /// Do not access this directly. It is only meant to be used by
@@ -1170,7 +1170,7 @@ impl<'tcx> TyCtxt<'tcx> {
         definitions: Definitions,
         cstore: Box<CrateStoreDyn>,
         untracked_resolutions: ty::ResolverOutputs,
-        krate: Lrc<ast::Crate>,
+        untracked_crate: IndexVec<LocalDefId, Steal<ty::AstOwner>>,
         dep_graph: DepGraph,
         on_disk_cache: Option<&'tcx dyn OnDiskCache<'tcx>>,
         queries: &'tcx dyn query::QueryEngine<'tcx>,
@@ -1200,7 +1200,7 @@ impl<'tcx> TyCtxt<'tcx> {
             lifetimes: common_lifetimes,
             consts: common_consts,
             untracked_resolutions,
-            untracked_crate: Steal::new(krate),
+            untracked_crate,
             on_disk_cache,
             queries,
             query_caches: query::QueryCaches::default(),
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
@@ -33,7 +33,9 @@ use rustc_ast::node_id::NodeMap;
 use rustc_attr as attr;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexMap};
 use rustc_data_structures::intern::{Interned, WithStableHash};
+use rustc_data_structures::owning_ref::OwningRef;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
+use rustc_data_structures::sync::Lrc;
 use rustc_data_structures::tagged_ptr::CopyTaggedPtr;
 use rustc_hir as hir;
 use rustc_hir::def::{CtorKind, CtorOf, DefKind, LifetimeRes, Res};
@@ -128,6 +130,23 @@ mod sty;
 
 pub type RegisteredTools = FxHashSet<Ident>;
 
+/// Pointers to parts of the AST to index it.
+/// We use OwningRef to release memory when all the owners have been taken.
+pub enum AstOwner {
+    NonOwner,
+    Synthetic(LocalDefId),
+    Crate(Lrc<ast::Crate>),
+    Item(OwningRef<Lrc<ast::Crate>, ast::Item>),
+    AssocItem(OwningRef<Lrc<ast::Crate>, ast::AssocItem>, ast::visit::AssocCtxt),
+    ForeignItem(OwningRef<Lrc<ast::Crate>, ast::ForeignItem>),
+}
+
+impl AstOwner {
+    pub fn expect_crate(&self) -> &Lrc<ast::Crate> {
+        if let AstOwner::Crate(c) = self { c } else { panic!() }
+    }
+}
+
 #[derive(Debug)]
 pub struct ResolverOutputs {
     pub visibilities: FxHashMap<LocalDefId, Visibility>,
