Auto merge of #80919 - cjgillot:defkey-span, r=oli-obk

Generate metadata by iterating on DefId instead of traversing the HIR tree 1/N

Sample from #80347.

4 files changed, 74 insertions, 48 deletions

diff --git a/compiler/rustc_middle/src/hir/map/mod.rs b/compiler/rustc_middle/src/hir/map/mod.rs
index 06bb1347dc1..9b01a64de84 100644
--- a/compiler/rustc_middle/src/hir/map/mod.rs
+++ b/compiler/rustc_middle/src/hir/map/mod.rs
@@ -1,7 +1,6 @@
 use self::collector::NodeCollector;
 
 use crate::hir::{Owner, OwnerNodes};
-use crate::ty::query::Providers;
 use crate::ty::TyCtxt;
 use rustc_ast as ast;
 use rustc_data_structures::svh::Svh;
@@ -183,14 +182,18 @@ impl<'hir> Map<'hir> {
         self.tcx.definitions.opt_local_def_id_to_hir_id(def_id)
     }
 
-    pub fn def_kind(&self, local_def_id: LocalDefId) -> DefKind {
+    pub fn iter_local_def_id(&self) -> impl Iterator<Item = LocalDefId> + '_ {
+        self.tcx.definitions.iter_local_def_id()
+    }
+
+    pub fn opt_def_kind(&self, local_def_id: LocalDefId) -> Option<DefKind> {
         // FIXME(eddyb) support `find` on the crate root.
         if local_def_id.to_def_id().index == CRATE_DEF_INDEX {
-            return DefKind::Mod;
+            return Some(DefKind::Mod);
         }
 
         let hir_id = self.local_def_id_to_hir_id(local_def_id);
-        match self.get(hir_id) {
+        let def_kind = match self.find(hir_id)? {
             Node::Item(item) => match item.kind {
                 ItemKind::Static(..) => DefKind::Static,
                 ItemKind::Const(..) => DefKind::Const,
@@ -249,6 +252,7 @@ impl<'hir> Map<'hir> {
                 GenericParamKind::Type { .. } => DefKind::TyParam,
                 GenericParamKind::Const { .. } => DefKind::ConstParam,
             },
+            Node::Crate(_) => DefKind::Mod,
             Node::Stmt(_)
             | Node::PathSegment(_)
             | Node::Ty(_)
@@ -260,9 +264,14 @@ impl<'hir> Map<'hir> {
             | Node::Arm(_)
             | Node::Lifetime(_)
             | Node::Visibility(_)
-            | Node::Block(_)
-            | Node::Crate(_) => bug!("def_kind: unsupported node: {}", self.node_to_string(hir_id)),
-        }
+            | Node::Block(_) => return None,
+        };
+        Some(def_kind)
+    }
+
+    pub fn def_kind(&self, local_def_id: LocalDefId) -> DefKind {
+        self.opt_def_kind(local_def_id)
+            .unwrap_or_else(|| bug!("def_kind: unsupported node: {:?}", local_def_id))
     }
 
     fn find_entry(&self, id: HirId) -> Option<Entry<'hir>> {
@@ -514,9 +523,7 @@ impl<'hir> Map<'hir> {
 
     /// Retrieves the `Node` corresponding to `id`, returning `None` if cannot be found.
     pub fn find(&self, hir_id: HirId) -> Option<Node<'hir>> {
-        self.find_entry(hir_id).and_then(|entry| {
-            if let Node::Crate(..) = entry.node { None } else { Some(entry.node) }
-        })
+        self.find_entry(hir_id).map(|entry| entry.node)
     }
 
     /// Similar to `get_parent`; returns the parent HIR Id, or just `hir_id` if there
@@ -848,50 +855,55 @@ impl<'hir> Map<'hir> {
     /// Gets the span of the definition of the specified HIR node.
     /// This is used by `tcx.get_span`
     pub fn span(&self, hir_id: HirId) -> Span {
-        match self.find_entry(hir_id).map(|entry| entry.node) {
-            Some(Node::Param(param)) => param.span,
-            Some(Node::Item(item)) => match &item.kind {
+        self.opt_span(hir_id)
+            .unwrap_or_else(|| bug!("hir::map::Map::span: id not in map: {:?}", hir_id))
+    }
+
+    pub fn opt_span(&self, hir_id: HirId) -> Option<Span> {
+        let span = match self.find_entry(hir_id)?.node {
+            Node::Param(param) => param.span,
+            Node::Item(item) => match &item.kind {
                 ItemKind::Fn(sig, _, _) => sig.span,
                 _ => item.span,
             },
-            Some(Node::ForeignItem(foreign_item)) => foreign_item.span,
-            Some(Node::TraitItem(trait_item)) => match &trait_item.kind {
+            Node::ForeignItem(foreign_item) => foreign_item.span,
+            Node::TraitItem(trait_item) => match &trait_item.kind {
                 TraitItemKind::Fn(sig, _) => sig.span,
                 _ => trait_item.span,
             },
-            Some(Node::ImplItem(impl_item)) => match &impl_item.kind {
+            Node::ImplItem(impl_item) => match &impl_item.kind {
                 ImplItemKind::Fn(sig, _) => sig.span,
                 _ => impl_item.span,
             },
-            Some(Node::Variant(variant)) => variant.span,
-            Some(Node::Field(field)) => field.span,
-            Some(Node::AnonConst(constant)) => self.body(constant.body).value.span,
-            Some(Node::Expr(expr)) => expr.span,
-            Some(Node::Stmt(stmt)) => stmt.span,
-            Some(Node::PathSegment(seg)) => seg.ident.span,
-            Some(Node::Ty(ty)) => ty.span,
-            Some(Node::TraitRef(tr)) => tr.path.span,
-            Some(Node::Binding(pat)) => pat.span,
-            Some(Node::Pat(pat)) => pat.span,
-            Some(Node::Arm(arm)) => arm.span,
-            Some(Node::Block(block)) => block.span,
-            Some(Node::Ctor(..)) => match self.find(self.get_parent_node(hir_id)) {
-                Some(Node::Item(item)) => item.span,
-                Some(Node::Variant(variant)) => variant.span,
+            Node::Variant(variant) => variant.span,
+            Node::Field(field) => field.span,
+            Node::AnonConst(constant) => self.body(constant.body).value.span,
+            Node::Expr(expr) => expr.span,
+            Node::Stmt(stmt) => stmt.span,
+            Node::PathSegment(seg) => seg.ident.span,
+            Node::Ty(ty) => ty.span,
+            Node::TraitRef(tr) => tr.path.span,
+            Node::Binding(pat) => pat.span,
+            Node::Pat(pat) => pat.span,
+            Node::Arm(arm) => arm.span,
+            Node::Block(block) => block.span,
+            Node::Ctor(..) => match self.find(self.get_parent_node(hir_id))? {
+                Node::Item(item) => item.span,
+                Node::Variant(variant) => variant.span,
                 _ => unreachable!(),
             },
-            Some(Node::Lifetime(lifetime)) => lifetime.span,
-            Some(Node::GenericParam(param)) => param.span,
-            Some(Node::Visibility(&Spanned {
+            Node::Lifetime(lifetime) => lifetime.span,
+            Node::GenericParam(param) => param.span,
+            Node::Visibility(&Spanned {
                 node: VisibilityKind::Restricted { ref path, .. },
                 ..
-            })) => path.span,
-            Some(Node::Visibility(v)) => bug!("unexpected Visibility {:?}", v),
-            Some(Node::Local(local)) => local.span,
-            Some(Node::MacroDef(macro_def)) => macro_def.span,
-            Some(Node::Crate(item)) => item.span,
-            None => bug!("hir::map::Map::span: id not in map: {:?}", hir_id),
-        }
+            }) => path.span,
+            Node::Visibility(v) => bug!("unexpected Visibility {:?}", v),
+            Node::Local(local) => local.span,
+            Node::MacroDef(macro_def) => macro_def.span,
+            Node::Crate(item) => item.span,
+        };
+        Some(span)
     }
 
     /// Like `hir.span()`, but includes the body of function items
@@ -907,7 +919,7 @@ impl<'hir> Map<'hir> {
     }
 
     pub fn span_if_local(&self, id: DefId) -> Option<Span> {
-        id.as_local().map(|id| self.span(self.local_def_id_to_hir_id(id)))
+        id.as_local().and_then(|id| self.opt_span(self.local_def_id_to_hir_id(id)))
     }
 
     pub fn res_span(&self, res: Res) -> Option<Span> {
@@ -1101,7 +1113,3 @@ fn hir_id_to_string(map: &Map<'_>, id: HirId) -> String {
         None => format!("unknown node{}", id_str),
     }
 }
-
-pub fn provide(providers: &mut Providers) {
-    providers.def_kind = |tcx, def_id| tcx.hir().def_kind(def_id.expect_local());
-}
diff --git a/compiler/rustc_middle/src/hir/mod.rs b/compiler/rustc_middle/src/hir/mod.rs
index ae3b30217cc..e96f3f2a294 100644
--- a/compiler/rustc_middle/src/hir/mod.rs
+++ b/compiler/rustc_middle/src/hir/mod.rs
@@ -15,6 +15,7 @@ use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_hir::def_id::{LocalDefId, LOCAL_CRATE};
 use rustc_hir::*;
 use rustc_index::vec::IndexVec;
+use rustc_span::DUMMY_SP;
 
 pub struct Owner<'tcx> {
     parent: HirId,
@@ -77,6 +78,7 @@ 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_deref();
+    providers.def_span = |tcx, def_id| tcx.hir().span_if_local(def_id).unwrap_or(DUMMY_SP);
     providers.fn_arg_names = |tcx, id| {
         let hir = tcx.hir();
         let hir_id = hir.local_def_id_to_hir_id(id.expect_local());
@@ -92,5 +94,5 @@ pub fn provide(providers: &mut Providers) {
             span_bug!(hir.span(hir_id), "fn_arg_names: unexpected item {:?}", id);
         }
     };
-    map::provide(providers);
+    providers.opt_def_kind = |tcx, def_id| tcx.hir().opt_def_kind(def_id.expect_local());
 }
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 00ee7b8ec77..e23d04f8211 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -878,7 +878,7 @@ rustc_queries! {
             cache_on_disk_if { true }
         }
 
-        query def_kind(def_id: DefId) -> DefKind {
+        query opt_def_kind(def_id: DefId) -> Option<DefKind> {
             desc { |tcx| "looking up definition kind of `{}`", tcx.def_path_str(def_id) }
         }
 
diff --git a/compiler/rustc_middle/src/ty/query/mod.rs b/compiler/rustc_middle/src/ty/query/mod.rs
index acfa58e511e..f580cb14dc9 100644
--- a/compiler/rustc_middle/src/ty/query/mod.rs
+++ b/compiler/rustc_middle/src/ty/query/mod.rs
@@ -130,3 +130,19 @@ mod sealed {
 }
 
 use sealed::IntoQueryParam;
+
+impl TyCtxt<'tcx> {
+    pub fn def_kind(self, def_id: impl IntoQueryParam<DefId>) -> DefKind {
+        let def_id = def_id.into_query_param();
+        self.opt_def_kind(def_id)
+            .unwrap_or_else(|| bug!("def_kind: unsupported node: {:?}", def_id))
+    }
+}
+
+impl TyCtxtAt<'tcx> {
+    pub fn def_kind(self, def_id: impl IntoQueryParam<DefId>) -> DefKind {
+        let def_id = def_id.into_query_param();
+        self.opt_def_kind(def_id)
+            .unwrap_or_else(|| bug!("def_kind: unsupported node: {:?}", def_id))
+    }
+}