Collect the new maps

Author: Zoxc

src/librustc/arena.rs

@@ -161,6 +161,10 @@ macro_rules! arena_types {
             [] type_binding: rustc_hir::TypeBinding<$tcx>,
             [] variant: rustc_hir::Variant<$tcx>,
             [] where_predicate: rustc_hir::WherePredicate<$tcx>,
+
+            // HIR query types
+            [] hir_owner: rustc::hir::HirOwner<$tcx>,
+            [] hir_owner_items: rustc::hir::HirOwnerItems<$tcx>,
         ], $tcx);
     )
 }

src/librustc/hir/map/collector.rs

@@ -1,6 +1,8 @@
+use crate::arena::Arena;
 use crate::dep_graph::{DepGraph, DepKind, DepNode, DepNodeIndex};
 use crate::hir::map::definitions::{self, DefPathHash};
 use crate::hir::map::{Entry, HirEntryMap, Map};
+use crate::hir::{HirItem, HirOwner, HirOwnerItems};
 use crate::ich::StableHashingContext;
 use crate::middle::cstore::CrateStore;
 use rustc_ast::ast::NodeId;
@@ -22,12 +24,17 @@ use std::iter::repeat;
 
 /// A visitor that walks over the HIR and collects `Node`s into a HIR map.
 pub(super) struct NodeCollector<'a, 'hir> {
+    arena: &'hir Arena<'hir>,
+
     /// The crate
     krate: &'hir Crate<'hir>,
 
     /// Source map
     source_map: &'a SourceMap,
 
+    owner_map: FxHashMap<DefIndex, &'hir HirOwner<'hir>>,
+    owner_items_map: FxHashMap<DefIndex, &'hir mut HirOwnerItems<'hir>>,
+
     /// The node map
     map: HirEntryMap<'hir>,
     /// The parent of this node
@@ -112,6 +119,7 @@ fn upstream_crates(cstore: &dyn CrateStore) -> Vec<(Symbol, Fingerprint, Svh)> {
 impl<'a, 'hir> NodeCollector<'a, 'hir> {
     pub(super) fn root(
         sess: &'a Session,
+        arena: &'hir Arena<'hir>,
         krate: &'hir Crate<'hir>,
         dep_graph: &'a DepGraph,
         definitions: &'a definitions::Definitions,
@@ -161,6 +169,7 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
         }
 
         let mut collector = NodeCollector {
+            arena,
             krate,
             source_map: sess.source_map(),
             map: IndexVec::from_elem_n(IndexVec::new(), definitions.def_index_count()),
@@ -174,6 +183,8 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
             hir_to_node_id,
             hcx,
             hir_body_nodes,
+            owner_map: FxHashMap::default(),
+            owner_items_map: FxHashMap::default(),
         };
         collector.insert_entry(
             hir::CRATE_HIR_ID,
@@ -192,7 +203,12 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
         crate_disambiguator: CrateDisambiguator,
         cstore: &dyn CrateStore,
         commandline_args_hash: u64,
-    ) -> (HirEntryMap<'hir>, Svh) {
+    ) -> (
+        HirEntryMap<'hir>,
+        FxHashMap<DefIndex, &'hir HirOwner<'hir>>,
+        FxHashMap<DefIndex, &'hir mut HirOwnerItems<'hir>>,
+        Svh,
+    ) {
         self.hir_body_nodes.sort_unstable_by_key(|bn| bn.0);
 
         let node_hashes = self.hir_body_nodes.iter().fold(
@@ -229,13 +245,36 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
         let crate_hash: Fingerprint = stable_hasher.finish();
 
         let svh = Svh::new(crate_hash.to_smaller_hash());
-        (self.map, svh)
+        (self.map, self.owner_map, self.owner_items_map, svh)
     }
 
     fn insert_entry(&mut self, id: HirId, entry: Entry<'hir>) {
+        let i = id.local_id.as_u32() as usize;
+
+        let owner = HirOwner { parent: entry.parent, node: entry.node };
+
+        let arena = self.arena;
+
+        let items = self.owner_items_map.entry(id.owner).or_insert_with(|| {
+            arena.alloc(HirOwnerItems { items: IndexVec::new(), bodies: FxHashMap::default() })
+        });
+
+        if i == 0 {
+            self.owner_map.insert(id.owner, self.arena.alloc(owner));
+        // FIXME: feature(impl_trait_in_bindings) broken and trigger this assert
+        //assert!(self.owner_map.insert(id.owner, self.arena.alloc(owner)).is_none());
+        } else {
+            let len = items.items.len();
+            if i >= len {
+                items.items.extend(repeat(None).take(i - len + 1));
+            }
+            assert_eq!(entry.parent.owner, id.owner);
+            items.items[id.local_id] =
+                Some(HirItem { parent: entry.parent.local_id, node: entry.node });
+        }
+
         debug!("hir_map: {:?} => {:?}", id, entry);
         let local_map = &mut self.map[id.owner];
-        let i = id.local_id.as_u32() as usize;
         let len = local_map.len();
         if i >= len {
             local_map.extend(repeat(None).take(i - len + 1));

src/librustc/hir/map/mod.rs

@@ -3,7 +3,9 @@ pub use self::definitions::{
     DefKey, DefPath, DefPathData, DefPathHash, Definitions, DisambiguatedDefPathData,
 };
 
+use crate::arena::Arena;
 use crate::dep_graph::{DepGraph, DepKind, DepNode, DepNodeIndex};
+use crate::hir::{HirOwner, HirOwnerItems};
 use crate::middle::cstore::CrateStoreDyn;
 use crate::ty::query::Providers;
 use rustc_ast::ast::{self, Name, NodeId};
@@ -145,6 +147,9 @@ pub struct Map<'hir> {
     /// The SVH of the local crate.
     pub crate_hash: Svh,
 
+    pub(super) owner_map: FxHashMap<DefIndex, &'hir HirOwner<'hir>>,
+    pub(super) owner_items_map: FxHashMap<DefIndex, &'hir HirOwnerItems<'hir>>,
+
     map: HirEntryMap<'hir>,
 
     definitions: Definitions,
@@ -1201,6 +1206,7 @@ impl Named for ImplItem<'_> {
 
 pub fn map_crate<'hir>(
     sess: &rustc_session::Session,
+    arena: &'hir Arena<'hir>,
     cstore: &CrateStoreDyn,
     krate: &'hir Crate<'hir>,
     dep_graph: DepGraph,
@@ -1215,19 +1221,28 @@ pub fn map_crate<'hir>(
         .map(|(node_id, &hir_id)| (hir_id, node_id))
         .collect();
 
-    let (map, crate_hash) = {
+    let (map, owner_map, owner_items_map, crate_hash) = {
         let hcx = crate::ich::StableHashingContext::new(sess, krate, &definitions, cstore);
 
         let mut collector =
-            NodeCollector::root(sess, krate, &dep_graph, &definitions, &hir_to_node_id, hcx);
+            NodeCollector::root(sess, arena, krate, &dep_graph, &definitions, &hir_to_node_id, hcx);
         intravisit::walk_crate(&mut collector, krate);
 
         let crate_disambiguator = sess.local_crate_disambiguator();
         let cmdline_args = sess.opts.dep_tracking_hash();
         collector.finalize_and_compute_crate_hash(crate_disambiguator, cstore, cmdline_args)
     };
 
-    let map = Map { krate, dep_graph, crate_hash, map, hir_to_node_id, definitions };
+    let map = Map {
+        krate,
+        dep_graph,
+        crate_hash,
+        map,
+        owner_map,
+        owner_items_map: owner_items_map.into_iter().map(|(k, v)| (k, &*v)).collect(),
+        hir_to_node_id,
+        definitions,
+    };
 
     sess.time("validate_HIR_map", || {
         hir_id_validator::check_crate(&map, sess);

src/librustc/hir/mod.rs

@@ -24,15 +24,15 @@ pub struct HirOwner<'tcx> {
     node: Node<'tcx>,
 }
 
-#[derive(HashStable)]
+#[derive(HashStable, Clone)]
 pub struct HirItem<'tcx> {
     parent: ItemLocalId,
     node: Node<'tcx>,
 }
 
 #[derive(HashStable)]
 pub struct HirOwnerItems<'tcx> {
-    owner: &'tcx HirOwner<'tcx>,
+    //owner: &'tcx HirOwner<'tcx>,
     items: IndexVec<ItemLocalId, Option<HirItem<'tcx>>>,
     bodies: FxHashMap<ItemLocalId, &'tcx Body<'tcx>>,
 }

src/librustc_interface/passes.rs

@@ -716,7 +716,7 @@ pub fn create_global_ctxt<'tcx>(
     let defs = mem::take(&mut resolver_outputs.definitions);
 
     // Construct the HIR map.
-    let hir_map = map::map_crate(sess, &*resolver_outputs.cstore, krate, dep_graph, defs);
+    let hir_map = map::map_crate(sess, &**arena, &*resolver_outputs.cstore, krate, dep_graph, defs);
 
     let query_result_on_disk_cache = rustc_incremental::load_query_result_cache(sess);