Query-ify LocalDefId <-> HirId conversion

Based on PR rust-lang#70039

Author: Aaron1011

src/librustc_middle/arena.rs

@@ -94,6 +94,7 @@ macro_rules! arena_types {
             [few] hir_definitions: rustc_hir::definitions::Definitions,
             [] hir_owner: rustc_middle::hir::Owner<$tcx>,
             [] hir_owner_nodes: rustc_middle::hir::OwnerNodes<$tcx>,
+            [] hir_owner_defs: rustc_data_structures::fx::FxHashMap<rustc_hir::ItemLocalId, rustc_span::def_id::LocalDefId>,
 
             // Note that this deliberately duplicates items in the `rustc_hir::arena`,
             // since we need to allocate this type on both the `rustc_hir` arena

src/librustc_middle/hir/map/collector.rs

@@ -133,7 +133,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(
@@ -229,6 +229,14 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
                 id.local_id,
                 ParentedNode { parent: entry.parent.local_id, node: entry.node },
             );
+
+            // Check if this HirId has a DefId, and insert it in the `defs` map if so.
+            if let Some(def_id) = self.definitions.opt_hir_id_to_local_def_id(id) {
+                if data.defs.is_none() {
+                    data.defs = Some(arena.alloc(FxHashMap::default()));
+                };
+                data.defs.as_mut().unwrap().insert(id.local_id, def_id);
+            }
         }
     }
 

src/librustc_middle/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 as ast;
+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, CRATE_DEF_INDEX, 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> {
@@ -170,17 +172,22 @@ impl<'hir> Map<'hir> {
 
     #[inline]
     pub fn opt_local_def_id(&self, hir_id: HirId) -> Option<LocalDefId> {
-        self.tcx.definitions.opt_hir_id_to_local_def_id(hir_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)
+        } else {
+            self.tcx.hir_owner_defs(hir_id.owner).and_then(|map| map.get(&hir_id.local_id).copied())
+        }
     }
 
     #[inline]
     pub fn local_def_id_to_hir_id(&self, def_id: LocalDefId) -> HirId {
-        self.tcx.definitions.local_def_id_to_hir_id(def_id)
+        self.tcx.local_def_id_to_hir_id(def_id).unwrap()
     }
 
     #[inline]
     pub fn opt_local_def_id_to_hir_id(&self, def_id: LocalDefId) -> Option<HirId> {
-        self.tcx.definitions.opt_local_def_id_to_hir_id(def_id)
+        self.tcx.local_def_id_to_hir_id(def_id)
     }
 
     pub fn def_kind(&self, local_def_id: LocalDefId) -> DefKind {

src/librustc_middle/hir/mod.rs

@@ -92,5 +92,8 @@ pub fn provide(providers: &mut Providers) {
             span_bug!(hir.span(hir_id), "fn_arg_names: unexpected item {:?}", id);
         }
     };
+    providers.local_def_id_to_hir_id = |tcx, id| tcx.definitions.opt_local_def_id_to_hir_id(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_middle/query/mod.rs

@@ -112,6 +112,22 @@ rustc_queries! {
             desc { |tcx| "computing the optional const parameter of `{}`", 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> {
             desc { |tcx| "computing type of `{}`", tcx.def_path_str(key) }