Add queries for LocalDefId <-> HirId conversion

Author: Zoxc

src/librustc/arena.rs

@@ -125,6 +125,10 @@ macro_rules! arena_types {
             [few] hir_definitions: rustc_hir::definitions::Definitions,
             [] hir_owner: rustc::hir::Owner<$tcx>,
             [] hir_owner_nodes: rustc::hir::OwnerNodes<$tcx>,
+            [] hir_owner_defs: rustc_data_structures::fx::FxHashMap<
+                rustc_hir::ItemLocalId,
+                rustc_span::def_id::LocalDefId
+            >,
         ], $tcx);
     )
 }

src/librustc/hir/map/collector.rs

@@ -132,7 +132,7 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
             hcx,
             hir_body_nodes,
             map: (0..definitions.def_index_count())
-                .map(|_| HirOwnerData { signature: None, with_bodies: None })
+                .map(|_| HirOwnerData { signature: None, with_bodies: None, defs: None })
                 .collect(),
         };
         collector.insert_entry(
@@ -202,6 +202,17 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
 
         let data = &mut self.map[id.owner];
 
+        if i != 0 {
+            // Check if this HirId has a DefId, and insert it in the `defs` map if so.
+            let node_id = self.definitions.hir_id_to_node_id(id);
+            if let Some(def_id) = self.definitions.opt_local_def_id(node_id) {
+                if data.defs.is_none() {
+                    data.defs = Some(arena.alloc(FxHashMap::default()));
+                };
+                data.defs.as_mut().unwrap().insert(id.local_id, def_id);
+            }
+        }
+
         if data.with_bodies.is_none() {
             data.with_bodies = Some(arena.alloc(OwnerNodes {
                 hash,

src/librustc/hir/map/mod.rs

@@ -4,6 +4,7 @@ use crate::hir::{Owner, OwnerNodes};
 use crate::ty::query::Providers;
 use crate::ty::TyCtxt;
 use rustc_ast::ast::{self, Name, NodeId};
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::svh::Svh;
 use rustc_hir::def::{DefKind, Res};
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId, LOCAL_CRATE};
@@ -89,6 +90,7 @@ fn is_body_owner<'hir>(node: Node<'hir>, hir_id: HirId) -> bool {
 pub(super) struct HirOwnerData<'hir> {
     pub(super) signature: Option<&'hir Owner<'hir>>,
     pub(super) with_bodies: Option<&'hir mut OwnerNodes<'hir>>,
+    pub(super) defs: Option<&'hir mut FxHashMap<ItemLocalId, LocalDefId>>,
 }
 
 pub struct IndexedHir<'hir> {
@@ -184,8 +186,14 @@ impl<'hir> Map<'hir> {
 
     #[inline]
     pub fn opt_local_def_id(&self, hir_id: HirId) -> Option<DefId> {
-        let node_id = self.hir_id_to_node_id(hir_id);
-        self.opt_local_def_id_from_node_id(node_id)
+        if hir_id.local_id == ItemLocalId::from_u32(0) {
+            // Every HirId owner has a DefId, so we can just return it directly here
+            Some(hir_id.owner.to_def_id())
+        } else {
+            self.tcx
+                .hir_owner_defs(hir_id.owner)
+                .and_then(|map| map.get(&hir_id.local_id).map(|id| id.to_def_id()))
+        }
     }
 
     #[inline]
@@ -200,7 +208,7 @@ impl<'hir> Map<'hir> {
 
     #[inline]
     pub fn as_local_hir_id(&self, def_id: DefId) -> Option<HirId> {
-        self.tcx.definitions.as_local_hir_id(def_id)
+        def_id.as_local().and_then(|def_id| self.tcx.local_def_id_to_hir_id(def_id))
     }
 
     #[inline]

src/librustc/hir/mod.rs

@@ -81,5 +81,8 @@ pub fn provide(providers: &mut Providers<'_>) {
     providers.hir_owner = |tcx, id| tcx.index_hir(LOCAL_CRATE).map[id].signature;
     providers.hir_owner_nodes =
         |tcx, id| tcx.index_hir(LOCAL_CRATE).map[id].with_bodies.as_ref().map(|nodes| &**nodes);
+    providers.local_def_id_to_hir_id = |tcx, id| tcx.definitions.as_local_hir_id(id.to_def_id());
+    providers.hir_owner_defs =
+        |tcx, index| tcx.index_hir(LOCAL_CRATE).map[index].defs.as_ref().map(|defs| &**defs);
     map::provide(providers);
 }

src/librustc/query/mod.rs

@@ -90,6 +90,22 @@ rustc_queries! {
             desc { |tcx| "HIR owner items in `{}`", tcx.def_path_str(key.to_def_id()) }
         }
 
+        // Converts a `LocalDefId` to a `HirId`.
+        // This can be conveniently accessed by `tcx.hir().as_local_hir_id`.
+        // Avoid calling this query directly.
+        query local_def_id_to_hir_id(key: LocalDefId) -> Option<hir::HirId> {
+            eval_always
+            desc { "converting a LocalDefId to HirId" }
+        }
+
+        // A map of the HIR items which also have a `DefId` in addition to their `HirId`.
+        // This can be conveniently accessed by `tcx.hir().opt_local_def_id`.
+        // Avoid calling this query directly.
+        query hir_owner_defs(key: LocalDefId) -> Option<&'tcx FxHashMap<ItemLocalId, LocalDefId>> {
+            eval_always
+            desc { "getting a LocalDefId map" }
+        }
+
         /// Records the type of every item.
         query type_of(key: DefId) -> Ty<'tcx> {
             cache_on_disk_if { key.is_local() }