Perform HIR indexing per-owner

compiler/rustc_hir/src/definitions.rs

@@ -92,6 +92,12 @@ impl DefPathTable {
             .iter_enumerated()
             .map(move |(index, key)| (index, key, &self.def_path_hashes[index]))
     }
+
+    pub fn all_def_path_hashes_and_def_ids(
+        &self,
+    ) -> impl Iterator<Item = (&DefPathHash, DefIndex)> + '_ {
+        self.def_path_hashes.iter_enumerated().map(move |(index, hash)| (hash, index))
+    }
 }
 
 /// The definition table containing node definitions.

compiler/rustc_middle/src/arena.rs

@@ -92,10 +92,7 @@ macro_rules! arena_types {
             [] predicates: rustc_middle::ty::PredicateInner<$tcx>,
 
             // HIR query types
-            [few] indexed_hir: rustc_middle::hir::IndexedHir<$tcx>,
-            [few] hir_definitions: rustc_hir::definitions::Definitions,
-            [] hir_owner: rustc_middle::hir::Owner<$tcx>,
-            [] hir_owner_nodes: rustc_middle::hir::OwnerNodes<$tcx>,
+            [] indexed_hir: rustc_middle::hir::IndexedHir<$tcx>,
 
             // Note that this deliberately duplicates items in the `rustc_hir::arena`,
             // since we need to allocate this type on both the `rustc_hir` arena

compiler/rustc_middle/src/hir/map/collector.rs

@@ -1,4 +1,3 @@
-use crate::arena::Arena;
 use crate::hir::map::Map;
 use crate::hir::{IndexedHir, OwnerNodes, ParentedNode};
 use crate::ich::StableHashingContext;
@@ -7,7 +6,6 @@ use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_hir as hir;
 use rustc_hir::def_id::LocalDefId;
-use rustc_hir::def_id::CRATE_DEF_ID;
 use rustc_hir::definitions;
 use rustc_hir::intravisit::{self, NestedVisitorMap, Visitor};
 use rustc_hir::*;
@@ -20,25 +18,21 @@ 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,
 
-    map: IndexVec<LocalDefId, Option<&'hir mut OwnerNodes<'hir>>>,
-    parenting: FxHashMap<LocalDefId, HirId>,
+    nodes: OwnerNodes<'hir>,
+    parenting: FxHashMap<LocalDefId, ItemLocalId>,
 
     /// The parent of this node
-    parent_node: hir::HirId,
+    parent_node: hir::ItemLocalId,
 
-    current_dep_node_owner: LocalDefId,
+    owner: LocalDefId,
 
     definitions: &'a definitions::Definitions,
-
-    hcx: StableHashingContext<'a>,
 }
 
 fn insert_vec_map<K: Idx, V: Clone>(map: &mut IndexVec<K, Option<V>>, k: K, v: V) {
@@ -62,58 +56,47 @@ fn hash_body(
     stable_hasher.finish()
 }
 
-impl<'a, 'hir> NodeCollector<'a, 'hir> {
-    pub(super) fn root(
-        sess: &'a Session,
-        arena: &'hir Arena<'hir>,
-        krate: &'hir Crate<'hir>,
-        definitions: &'a definitions::Definitions,
-        hcx: StableHashingContext<'a>,
-    ) -> NodeCollector<'a, 'hir> {
-        let mut collector = NodeCollector {
-            arena,
-            krate,
-            source_map: sess.source_map(),
-            parent_node: hir::CRATE_HIR_ID,
-            current_dep_node_owner: CRATE_DEF_ID,
-            definitions,
-            hcx,
-            map: IndexVec::from_fn_n(|_| None, definitions.def_index_count()),
-            parenting: FxHashMap::default(),
-        };
-        collector.insert_owner(CRATE_DEF_ID, OwnerNode::Crate(krate.module()));
-
-        collector
-    }
-
-    pub(super) fn finalize_and_compute_crate_hash(mut self) -> IndexedHir<'hir> {
-        // Insert bodies into the map
-        for (id, body) in self.krate.bodies.iter() {
-            let bodies = &mut self.map[id.hir_id.owner].as_mut().unwrap().bodies;
-            assert!(bodies.insert(id.hir_id.local_id, body).is_none());
-        }
-        IndexedHir { map: self.map, parenting: self.parenting }
-    }
-
-    fn insert_owner(&mut self, owner: LocalDefId, node: OwnerNode<'hir>) {
-        let hash = hash_body(&mut self.hcx, node);
-
-        let mut nodes = IndexVec::new();
-        nodes.push(Some(ParentedNode { parent: ItemLocalId::new(0), node: node.into() }));
-
-        debug_assert!(self.map[owner].is_none());
-        self.map[owner] =
-            Some(self.arena.alloc(OwnerNodes { hash, nodes, bodies: FxHashMap::default() }));
-    }
+pub(super) fn collect(
+    sess: &'a Session,
+    krate: &'hir Crate<'hir>,
+    definitions: &'a definitions::Definitions,
+    mut hcx: StableHashingContext<'a>,
+    owner: LocalDefId,
+) -> Option<IndexedHir<'hir>> {
+    let item = *krate.owners.get(owner)?.as_ref()?;
+    let hash = hash_body(&mut hcx, item);
+    let mut nodes = IndexVec::new();
+    nodes.push(Some(ParentedNode { parent: ItemLocalId::new(0), node: item.into() }));
+    let mut collector = NodeCollector {
+        krate,
+        source_map: sess.source_map(),
+        owner,
+        parent_node: ItemLocalId::new(0),
+        definitions,
+        nodes: OwnerNodes { hash, nodes, bodies: FxHashMap::default() },
+        parenting: FxHashMap::default(),
+    };
+
+    match item {
+        OwnerNode::Crate(citem) => collector.visit_mod(&citem, citem.inner, hir::CRATE_HIR_ID),
+        OwnerNode::Item(item) => collector.visit_item(item),
+        OwnerNode::TraitItem(item) => collector.visit_trait_item(item),
+        OwnerNode::ImplItem(item) => collector.visit_impl_item(item),
+        OwnerNode::ForeignItem(item) => collector.visit_foreign_item(item),
+    };
+
+    Some(IndexedHir { nodes: collector.nodes, parenting: collector.parenting })
+}
 
+impl<'a, 'hir> NodeCollector<'a, 'hir> {
     fn insert(&mut self, span: Span, hir_id: HirId, node: Node<'hir>) {
-        debug_assert_eq!(self.current_dep_node_owner, hir_id.owner);
+        debug_assert_eq!(self.owner, hir_id.owner);
         debug_assert_ne!(hir_id.local_id.as_u32(), 0);
 
         // Make sure that the DepNode of some node coincides with the HirId
         // owner of that node.
         if cfg!(debug_assertions) {
-            if hir_id.owner != self.current_dep_node_owner {
+            if hir_id.owner != self.owner {
                 let node_str = match self.definitions.opt_hir_id_to_local_def_id(hir_id) {
                     Some(def_id) => self.definitions.def_path(def_id).to_string_no_crate_verbose(),
                     None => format!("{:?}", node),
@@ -125,58 +108,37 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
                      current_dep_node_owner={} ({:?}), hir_id.owner={} ({:?})",
                     self.source_map.span_to_diagnostic_string(span),
                     node_str,
-                    self.definitions
-                        .def_path(self.current_dep_node_owner)
-                        .to_string_no_crate_verbose(),
-                    self.current_dep_node_owner,
+                    self.definitions.def_path(self.owner).to_string_no_crate_verbose(),
+                    self.owner,
                     self.definitions.def_path(hir_id.owner).to_string_no_crate_verbose(),
                     hir_id.owner,
                 )
             }
         }
 
-        let nodes = self.map[hir_id.owner].as_mut().unwrap();
-
-        debug_assert_eq!(self.parent_node.owner, self.current_dep_node_owner);
         insert_vec_map(
-            &mut nodes.nodes,
+            &mut self.nodes.nodes,
             hir_id.local_id,
-            ParentedNode { parent: self.parent_node.local_id, node: node },
+            ParentedNode { parent: self.parent_node, node: node },
         );
     }
 
     fn with_parent<F: FnOnce(&mut Self)>(&mut self, parent_node_id: HirId, f: F) {
+        debug_assert_eq!(parent_node_id.owner, self.owner);
         let parent_node = self.parent_node;
-        self.parent_node = parent_node_id;
+        self.parent_node = parent_node_id.local_id;
         f(self);
         self.parent_node = parent_node;
     }
 
-    fn with_dep_node_owner(&mut self, dep_node_owner: LocalDefId, f: impl FnOnce(&mut Self)) {
-        let prev_owner = self.current_dep_node_owner;
-        let prev_parent = self.parent_node;
-
-        self.current_dep_node_owner = dep_node_owner;
-        self.parent_node = HirId::make_owner(dep_node_owner);
-        f(self);
-        self.current_dep_node_owner = prev_owner;
-        self.parent_node = prev_parent;
-    }
-
     fn insert_nested(&mut self, item: LocalDefId) {
-        #[cfg(debug_assertions)]
-        {
-            let dk_parent = self.definitions.def_key(item).parent.unwrap();
-            let dk_parent = LocalDefId { local_def_index: dk_parent };
-            let dk_parent = self.definitions.local_def_id_to_hir_id(dk_parent);
-            debug_assert_eq!(
-                dk_parent.owner, self.parent_node.owner,
-                "Different parents for {:?}",
-                item
-            )
+        let dk_parent = self.definitions.def_key(item).parent.unwrap();
+        let dk_parent = LocalDefId { local_def_index: dk_parent };
+        let dk_parent = self.definitions.local_def_id_to_hir_id(dk_parent);
+        debug_assert_eq!(dk_parent.owner, self.owner, "Different parents for {:?}", item);
+        if dk_parent.local_id != self.parent_node {
+            self.parenting.insert(item, self.parent_node);
         }
-
-        assert_eq!(self.parenting.insert(item, self.parent_node), None);
     }
 }
 
@@ -194,26 +156,25 @@ impl<'a, 'hir> Visitor<'hir> for NodeCollector<'a, 'hir> {
     fn visit_nested_item(&mut self, item: ItemId) {
         debug!("visit_nested_item: {:?}", item);
         self.insert_nested(item.def_id);
-        self.visit_item(self.krate.item(item));
     }
 
     fn visit_nested_trait_item(&mut self, item_id: TraitItemId) {
         self.insert_nested(item_id.def_id);
-        self.visit_trait_item(self.krate.trait_item(item_id));
     }
 
     fn visit_nested_impl_item(&mut self, item_id: ImplItemId) {
         self.insert_nested(item_id.def_id);
-        self.visit_impl_item(self.krate.impl_item(item_id));
     }
 
     fn visit_nested_foreign_item(&mut self, foreign_id: ForeignItemId) {
         self.insert_nested(foreign_id.def_id);
-        self.visit_foreign_item(self.krate.foreign_item(foreign_id));
     }
 
     fn visit_nested_body(&mut self, id: BodyId) {
-        self.visit_body(self.krate.body(id));
+        let body = self.krate.body(id);
+        debug_assert_eq!(id.hir_id.owner, self.owner);
+        assert!(self.nodes.bodies.insert(id.hir_id.local_id, body).is_none());
+        self.visit_body(body);
     }
 
     fn visit_param(&mut self, param: &'hir Param<'hir>) {
@@ -226,8 +187,8 @@ impl<'a, 'hir> Visitor<'hir> for NodeCollector<'a, 'hir> {
 
     fn visit_item(&mut self, i: &'hir Item<'hir>) {
         debug!("visit_item: {:?}", i);
-        self.insert_owner(i.def_id, OwnerNode::Item(i));
-        self.with_dep_node_owner(i.def_id, |this| {
+        debug_assert_eq!(i.def_id, self.owner);
+        self.with_parent(i.hir_id(), |this| {
             if let ItemKind::Struct(ref struct_def, _) = i.kind {
                 // If this is a tuple or unit-like struct, register the constructor.
                 if let Some(ctor_hir_id) = struct_def.ctor_hir_id() {
@@ -239,8 +200,8 @@ impl<'a, 'hir> Visitor<'hir> for NodeCollector<'a, 'hir> {
     }
 
     fn visit_foreign_item(&mut self, fi: &'hir ForeignItem<'hir>) {
-        self.insert_owner(fi.def_id, OwnerNode::ForeignItem(fi));
-        self.with_dep_node_owner(fi.def_id, |this| {
+        debug_assert_eq!(fi.def_id, self.owner);
+        self.with_parent(fi.hir_id(), |this| {
             intravisit::walk_foreign_item(this, fi);
         });
     }
@@ -257,15 +218,15 @@ impl<'a, 'hir> Visitor<'hir> for NodeCollector<'a, 'hir> {
     }
 
     fn visit_trait_item(&mut self, ti: &'hir TraitItem<'hir>) {
-        self.insert_owner(ti.def_id, OwnerNode::TraitItem(ti));
-        self.with_dep_node_owner(ti.def_id, |this| {
+        debug_assert_eq!(ti.def_id, self.owner);
+        self.with_parent(ti.hir_id(), |this| {
             intravisit::walk_trait_item(this, ti);
         });
     }
 
     fn visit_impl_item(&mut self, ii: &'hir ImplItem<'hir>) {
-        self.insert_owner(ii.def_id, OwnerNode::ImplItem(ii));
-        self.with_dep_node_owner(ii.def_id, |this| {
+        debug_assert_eq!(ii.def_id, self.owner);
+        self.with_parent(ii.hir_id(), |this| {
             intravisit::walk_impl_item(this, ii);
         });
     }
@@ -353,7 +314,8 @@ impl<'a, 'hir> Visitor<'hir> for NodeCollector<'a, 'hir> {
         s: Span,
         id: HirId,
     ) {
-        assert_eq!(self.parent_node, id);
+        assert_eq!(self.owner, id.owner);
+        assert_eq!(self.parent_node, id.local_id);
         intravisit::walk_fn(self, fk, fd, b, s, id);
     }
 

compiler/rustc_middle/src/hir/map/mod.rs

@@ -1,5 +1,3 @@
-use self::collector::NodeCollector;
-
 use crate::hir::{AttributeMap, IndexedHir, Owner};
 use crate::ty::TyCtxt;
 use rustc_ast as ast;
@@ -306,7 +304,7 @@ impl<'hir> Map<'hir> {
     }
 
     pub fn get_parent_node(&self, hir_id: HirId) -> HirId {
-        self.find_parent_node(hir_id).unwrap_or(CRATE_HIR_ID)
+        self.find_parent_node(hir_id).unwrap()
     }
 
     /// Retrieves the `Node` corresponding to `id`, returning `None` if cannot be found.
@@ -938,38 +936,31 @@ impl<'hir> intravisit::Map<'hir> for Map<'hir> {
     }
 }
 
-pub(super) fn index_hir<'tcx>(tcx: TyCtxt<'tcx>, (): ()) -> &'tcx IndexedHir<'tcx> {
-    let _prof_timer = tcx.sess.prof.generic_activity("build_hir_map");
-
-    // We can access untracked state since we are an eval_always query.
-    let hcx = tcx.create_stable_hashing_context();
-    let mut collector = NodeCollector::root(
+pub(super) fn index_hir<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    owner: LocalDefId,
+) -> Option<&'tcx IndexedHir<'tcx>> {
+    let map = collector::collect(
         tcx.sess,
-        &**tcx.arena,
         tcx.untracked_crate,
         &tcx.untracked_resolutions.definitions,
-        hcx,
-    );
-    let top_mod = tcx.untracked_crate.module();
-    collector.visit_mod(top_mod, top_mod.inner, CRATE_HIR_ID);
+        tcx.create_stable_hashing_context(),
+        owner,
+    )?;
 
-    let map = collector.finalize_and_compute_crate_hash();
-    tcx.arena.alloc(map)
+    Some(&*tcx.arena.alloc(map))
 }
 
 pub(super) fn crate_hash(tcx: TyCtxt<'_>, crate_num: CrateNum) -> Svh {
-    assert_eq!(crate_num, LOCAL_CRATE);
-
-    // We can access untracked state since we are an eval_always query.
-    let mut hcx = tcx.create_stable_hashing_context();
-
+    debug_assert_eq!(crate_num, LOCAL_CRATE);
     let mut hir_body_nodes: Vec<_> = tcx
-        .index_hir(())
-        .map
-        .iter_enumerated()
-        .filter_map(|(def_id, hod)| {
-            let def_path_hash = tcx.untracked_resolutions.definitions.def_path_hash(def_id);
-            let hash = hod.as_ref()?.hash;
+        .untracked_resolutions
+        .definitions
+        .def_path_table()
+        .all_def_path_hashes_and_def_ids()
+        .filter_map(|(def_path_hash, local_def_index)| {
+            let def_id = LocalDefId { local_def_index };
+            let hash = tcx.index_hir(def_id).as_ref()?.nodes.hash;
             Some((def_path_hash, hash, def_id))
         })
         .collect();
@@ -993,6 +984,7 @@ pub(super) fn crate_hash(tcx: TyCtxt<'_>, crate_num: CrateNum) -> Svh {
 
     source_file_names.sort_unstable();
 
+    let mut hcx = tcx.create_stable_hashing_context();
     let mut stable_hasher = StableHasher::new();
     for (def_path_hash, fingerprint, def_id) in hir_body_nodes.iter() {
         def_path_hash.0.hash_stable(&mut hcx, &mut stable_hasher);

compiler/rustc_middle/src/hir/mod.rs

@@ -19,16 +19,13 @@ use rustc_index::vec::{Idx, IndexVec};
 use rustc_span::DUMMY_SP;
 use std::collections::BTreeMap;
 
-/// Result of HIR indexing.
-#[derive(Debug)]
+/// Result of HIR indexing for a given HIR owner.
+#[derive(Debug, HashStable)]
 pub struct IndexedHir<'hir> {
-    /// Contents of the HIR owned by each definition. None for definitions that are not HIR owners.
-    // The `mut` comes from construction time, and is harmless since we only ever hand out
-    // immutable refs to IndexedHir.
-    map: IndexVec<LocalDefId, Option<&'hir mut OwnerNodes<'hir>>>,
-    /// Map from each owner to its parent's HirId inside another owner.
-    // This map is separate from `map` to eventually allow for per-owner indexing.
-    parenting: FxHashMap<LocalDefId, HirId>,
+    /// Contents of the HIR.
+    nodes: OwnerNodes<'hir>,
+    /// Map from each nested owner to its parent's local id.
+    parenting: FxHashMap<LocalDefId, ItemLocalId>,
 }
 
 /// Top-level HIR node for current owner. This only contains the node for which
@@ -140,19 +137,28 @@ pub fn provide(providers: &mut Providers) {
     providers.hir_crate = |tcx, ()| tcx.untracked_crate;
     providers.index_hir = map::index_hir;
     providers.crate_hash = map::crate_hash;
-    providers.hir_module_items = |tcx, id| &tcx.untracked_crate.modules[&id];
+    providers.hir_module_items = |tcx, id| &tcx.hir_crate(()).modules[&id];
     providers.hir_owner = |tcx, id| {
-        let owner = tcx.index_hir(()).map[id].as_ref()?;
-        let node = owner.nodes[ItemLocalId::new(0)].as_ref().unwrap().node;
+        let owner = tcx.index_hir(id)?;
+        let node = owner.nodes.nodes[ItemLocalId::new(0)].as_ref().unwrap().node;
         let node = node.as_owner().unwrap(); // Indexing must ensure it is an OwnerNode.
         Some(Owner { node })
     };
-    providers.hir_owner_nodes = |tcx, id| tcx.index_hir(()).map[id].as_deref();
+    providers.hir_owner_nodes = |tcx, id| tcx.index_hir(id).map(|i| &i.nodes);
     providers.hir_owner_parent = |tcx, id| {
-        let index = tcx.index_hir(());
-        index.parenting.get(&id).copied().unwrap_or(CRATE_HIR_ID)
+        let parent = tcx.untracked_resolutions.definitions.def_key(id).parent;
+        let parent = parent.map_or(CRATE_HIR_ID, |local_def_index| {
+            let def_id = LocalDefId { local_def_index };
+            let mut parent_hir_id =
+                tcx.untracked_resolutions.definitions.local_def_id_to_hir_id(def_id);
+            if let Some(local_id) = tcx.index_hir(parent_hir_id.owner).unwrap().parenting.get(&id) {
+                parent_hir_id.local_id = *local_id;
+            }
+            parent_hir_id
+        });
+        parent
     };
-    providers.hir_attrs = |tcx, id| AttributeMap { map: &tcx.untracked_crate.attrs, prefix: id };
+    providers.hir_attrs = |tcx, id| AttributeMap { map: &tcx.hir_crate(()).attrs, prefix: id };
     providers.source_span = |tcx, def_id| tcx.resolutions(()).definitions.def_span(def_id);
     providers.def_span = |tcx, def_id| tcx.hir().span_if_local(def_id).unwrap_or(DUMMY_SP);
     providers.fn_arg_names = |tcx, id| {

compiler/rustc_middle/src/query/mod.rs

@@ -42,9 +42,8 @@ rustc_queries! {
 
     /// The indexed HIR. This can be conveniently accessed by `tcx.hir()`.
     /// Avoid calling this query directly.
-    query index_hir(_: ()) -> &'tcx crate::hir::IndexedHir<'tcx> {
+    query index_hir(_: LocalDefId) -> Option<&'tcx crate::hir::IndexedHir<'tcx>> {
         eval_always
-        no_hash
         desc { "index HIR" }
     }
 
@@ -53,7 +52,6 @@ rustc_queries! {
     /// This can be conveniently accessed by `tcx.hir().visit_item_likes_in_module`.
     /// Avoid calling this query directly.
     query hir_module_items(key: LocalDefId) -> &'tcx hir::ModuleItems {
-        eval_always
         desc { |tcx| "HIR module items in `{}`", tcx.def_path_str(key.to_def_id()) }
     }
 
@@ -62,7 +60,6 @@ rustc_queries! {
     /// This can be conveniently accessed by methods on `tcx.hir()`.
     /// Avoid calling this query directly.
     query hir_owner(key: LocalDefId) -> Option<crate::hir::Owner<'tcx>> {
-        eval_always
         desc { |tcx| "HIR owner of `{}`", tcx.def_path_str(key.to_def_id()) }
     }
 
@@ -71,7 +68,6 @@ rustc_queries! {
     /// This can be conveniently accessed by methods on `tcx.hir()`.
     /// Avoid calling this query directly.
     query hir_owner_parent(key: LocalDefId) -> hir::HirId {
-        eval_always
         desc { |tcx| "HIR parent of `{}`", tcx.def_path_str(key.to_def_id()) }
     }
 
@@ -80,7 +76,6 @@ rustc_queries! {
     /// This can be conveniently accessed by methods on `tcx.hir()`.
     /// Avoid calling this query directly.
     query hir_owner_nodes(key: LocalDefId) -> Option<&'tcx crate::hir::OwnerNodes<'tcx>> {
-        eval_always
         desc { |tcx| "HIR owner items in `{}`", tcx.def_path_str(key.to_def_id()) }
     }
 
@@ -89,7 +84,6 @@ rustc_queries! {
     /// This can be conveniently accessed by methods on `tcx.hir()`.
     /// Avoid calling this query directly.
     query hir_attrs(key: LocalDefId) -> rustc_middle::hir::AttributeMap<'tcx> {
-        eval_always
         desc { |tcx| "HIR owner attributes in `{}`", tcx.def_path_str(key.to_def_id()) }
     }
 
@@ -941,12 +935,6 @@ rustc_queries! {
 
     query def_span(def_id: DefId) -> Span {
         desc { |tcx| "looking up span for `{}`", tcx.def_path_str(def_id) }
-        // FIXME(mw): DefSpans are not really inputs since they are derived from
-        // HIR. But at the moment HIR hashing still contains some hacks that allow
-        // to make type debuginfo to be source location independent. Declaring
-        // DefSpan an input makes sure that changes to these are always detected
-        // regardless of HIR hashing.
-        eval_always
     }
 
     query def_ident_span(def_id: DefId) -> Option<Span> {