Auto merge of #15209 - Veykril:ast-id, r=Veykril

internal: Stronger typing for AstId and AstIdMap

10 files changed, 102 insertions, 50 deletions

diff --git a/crates/hir-def/src/data.rs b/crates/hir-def/src/data.rs
index 4a9f08dca6d..54fe9a2e844 100644
--- a/crates/hir-def/src/data.rs
+++ b/crates/hir-def/src/data.rs
@@ -599,7 +599,7 @@ impl<'a> AssocItemCollector<'a> {
             if !attrs.is_cfg_enabled(self.expander.cfg_options()) {
                 self.diagnostics.push(DefDiagnostic::unconfigured_code(
                     self.module_id.local_id,
-                    InFile::new(self.expander.current_file_id(), item.ast_id(item_tree).upcast()),
+                    InFile::new(self.expander.current_file_id(), item.ast_id(item_tree).erase()),
                     attrs.cfg().unwrap(),
                     self.expander.cfg_options().clone(),
                 ));
diff --git a/crates/hir-def/src/data/adt.rs b/crates/hir-def/src/data/adt.rs
index 612a5ca1a28..595b9b770b7 100644
--- a/crates/hir-def/src/data/adt.rs
+++ b/crates/hir-def/src/data/adt.rs
@@ -330,7 +330,7 @@ impl EnumData {
             } else {
                 diagnostics.push(DefDiagnostic::unconfigured_code(
                     loc.container.local_id,
-                    InFile::new(loc.id.file_id(), var.ast_id.upcast()),
+                    InFile::new(loc.id.file_id(), var.ast_id.erase()),
                     attrs.cfg().unwrap(),
                     cfg_options.clone(),
                 ))
@@ -540,8 +540,8 @@ fn lower_fields(
                         InFile::new(
                             current_file_id,
                             match field.ast_id {
-                                FieldAstId::Record(it) => it.upcast(),
-                                FieldAstId::Tuple(it) => it.upcast(),
+                                FieldAstId::Record(it) => it.erase(),
+                                FieldAstId::Tuple(it) => it.erase(),
                             },
                         ),
                         attrs.cfg().unwrap(),
@@ -564,8 +564,8 @@ fn lower_fields(
                         InFile::new(
                             current_file_id,
                             match field.ast_id {
-                                FieldAstId::Record(it) => it.upcast(),
-                                FieldAstId::Tuple(it) => it.upcast(),
+                                FieldAstId::Record(it) => it.erase(),
+                                FieldAstId::Tuple(it) => it.erase(),
                             },
                         ),
                         attrs.cfg().unwrap(),
diff --git a/crates/hir-def/src/item_tree.rs b/crates/hir-def/src/item_tree.rs
index e74b71888c2..6f80bb6e07c 100644
--- a/crates/hir-def/src/item_tree.rs
+++ b/crates/hir-def/src/item_tree.rs
@@ -46,7 +46,7 @@ use ast::{AstNode, HasName, StructKind};
 use base_db::CrateId;
 use either::Either;
 use hir_expand::{
-    ast_id_map::FileAstId,
+    ast_id_map::{AstIdNode, FileAstId},
     attrs::RawAttrs,
     hygiene::Hygiene,
     name::{name, AsName, Name},
@@ -314,7 +314,7 @@ from_attrs!(ModItem(ModItem), Variant(Idx<Variant>), Field(Idx<Field>), Param(Id
 
 /// Trait implemented by all item nodes in the item tree.
 pub trait ItemTreeNode: Clone {
-    type Source: AstNode + Into<ast::Item>;
+    type Source: AstIdNode + Into<ast::Item>;
 
     fn ast_id(&self) -> FileAstId<Self::Source>;
 
diff --git a/crates/hir-def/src/lib.rs b/crates/hir-def/src/lib.rs
index e5149ab577c..c0701e2576a 100644
--- a/crates/hir-def/src/lib.rs
+++ b/crates/hir-def/src/lib.rs
@@ -64,7 +64,7 @@ use std::{
 
 use base_db::{impl_intern_key, salsa, CrateId, ProcMacroKind};
 use hir_expand::{
-    ast_id_map::FileAstId,
+    ast_id_map::{AstIdNode, FileAstId},
     attrs::{Attr, AttrId, AttrInput},
     builtin_attr_macro::BuiltinAttrExpander,
     builtin_derive_macro::BuiltinDeriveExpander,
@@ -1124,12 +1124,12 @@ impl AsMacroCall for InFile<&ast::MacroCall> {
 
 /// Helper wrapper for `AstId` with `ModPath`
 #[derive(Clone, Debug, Eq, PartialEq)]
-struct AstIdWithPath<T: ast::AstNode> {
+struct AstIdWithPath<T: AstIdNode> {
     ast_id: AstId<T>,
     path: path::ModPath,
 }
 
-impl<T: ast::AstNode> AstIdWithPath<T> {
+impl<T: AstIdNode> AstIdWithPath<T> {
     fn new(file_id: HirFileId, ast_id: FileAstId<T>, path: path::ModPath) -> AstIdWithPath<T> {
         AstIdWithPath { ast_id: AstId::new(file_id, ast_id), path }
     }
diff --git a/crates/hir-def/src/lower.rs b/crates/hir-def/src/lower.rs
index af623fd0e5d..e523c229179 100644
--- a/crates/hir-def/src/lower.rs
+++ b/crates/hir-def/src/lower.rs
@@ -1,5 +1,9 @@
 //! Context for lowering paths.
-use hir_expand::{ast_id_map::AstIdMap, hygiene::Hygiene, AstId, HirFileId, InFile};
+use hir_expand::{
+    ast_id_map::{AstIdMap, AstIdNode},
+    hygiene::Hygiene,
+    AstId, HirFileId, InFile,
+};
 use once_cell::unsync::OnceCell;
 use syntax::ast;
 use triomphe::Arc;
@@ -37,7 +41,7 @@ impl<'a> LowerCtx<'a> {
         Path::from_src(ast, self)
     }
 
-    pub(crate) fn ast_id<N: syntax::AstNode>(&self, item: &N) -> Option<AstId<N>> {
+    pub(crate) fn ast_id<N: AstIdNode>(&self, item: &N) -> Option<AstId<N>> {
         let &(file_id, ref ast_id_map) = self.ast_id_map.as_ref()?;
         let ast_id_map = ast_id_map.get_or_init(|| self.db.ast_id_map(file_id));
         Some(InFile::new(file_id, ast_id_map.ast_id(item)))
diff --git a/crates/hir-def/src/nameres/collector.rs b/crates/hir-def/src/nameres/collector.rs
index a5cdee9628d..c048716d740 100644
--- a/crates/hir-def/src/nameres/collector.rs
+++ b/crates/hir-def/src/nameres/collector.rs
@@ -2280,7 +2280,7 @@ impl ModCollector<'_, '_> {
     fn emit_unconfigured_diagnostic(&mut self, item: ModItem, cfg: &CfgExpr) {
         let ast_id = item.ast_id(self.item_tree);
 
-        let ast_id = InFile::new(self.file_id(), ast_id.upcast());
+        let ast_id = InFile::new(self.file_id(), ast_id.erase());
         self.def_collector.def_map.diagnostics.push(DefDiagnostic::unconfigured_code(
             self.module_id,
             ast_id,
diff --git a/crates/hir-def/src/nameres/diagnostics.rs b/crates/hir-def/src/nameres/diagnostics.rs
index 18b424255cd..e82e97b628e 100644
--- a/crates/hir-def/src/nameres/diagnostics.rs
+++ b/crates/hir-def/src/nameres/diagnostics.rs
@@ -2,12 +2,9 @@
 
 use base_db::CrateId;
 use cfg::{CfgExpr, CfgOptions};
-use hir_expand::{attrs::AttrId, MacroCallKind};
+use hir_expand::{attrs::AttrId, ErasedAstId, MacroCallKind};
 use la_arena::Idx;
-use syntax::{
-    ast::{self, AnyHasAttrs},
-    SyntaxError,
-};
+use syntax::{ast, SyntaxError};
 
 use crate::{
     item_tree::{self, ItemTreeId},
@@ -24,7 +21,7 @@ pub enum DefDiagnosticKind {
 
     UnresolvedImport { id: ItemTreeId<item_tree::Import>, index: Idx<ast::UseTree> },
 
-    UnconfiguredCode { ast: AstId<AnyHasAttrs>, cfg: CfgExpr, opts: CfgOptions },
+    UnconfiguredCode { ast: ErasedAstId, cfg: CfgExpr, opts: CfgOptions },
 
     UnresolvedProcMacro { ast: MacroCallKind, krate: CrateId },
 
@@ -81,7 +78,7 @@ impl DefDiagnostic {
 
     pub fn unconfigured_code(
         container: LocalModuleId,
-        ast: AstId<ast::AnyHasAttrs>,
+        ast: ErasedAstId,
         cfg: CfgExpr,
         opts: CfgOptions,
     ) -> Self {
diff --git a/crates/hir-expand/src/ast_id_map.rs b/crates/hir-expand/src/ast_id_map.rs
index c2b0d5985e3..1906ed15bae 100644
--- a/crates/hir-expand/src/ast_id_map.rs
+++ b/crates/hir-expand/src/ast_id_map.rs
@@ -18,47 +18,89 @@ use rustc_hash::FxHasher;
 use syntax::{ast, AstNode, AstPtr, SyntaxNode, SyntaxNodePtr};
 
 /// `AstId` points to an AST node in a specific file.
-pub struct FileAstId<N: AstNode> {
+pub struct FileAstId<N: AstIdNode> {
     raw: ErasedFileAstId,
     covariant: PhantomData<fn() -> N>,
 }
 
-impl<N: AstNode> Clone for FileAstId<N> {
+impl<N: AstIdNode> Clone for FileAstId<N> {
     fn clone(&self) -> FileAstId<N> {
         *self
     }
 }
-impl<N: AstNode> Copy for FileAstId<N> {}
+impl<N: AstIdNode> Copy for FileAstId<N> {}
 
-impl<N: AstNode> PartialEq for FileAstId<N> {
+impl<N: AstIdNode> PartialEq for FileAstId<N> {
     fn eq(&self, other: &Self) -> bool {
         self.raw == other.raw
     }
 }
-impl<N: AstNode> Eq for FileAstId<N> {}
-impl<N: AstNode> Hash for FileAstId<N> {
+impl<N: AstIdNode> Eq for FileAstId<N> {}
+impl<N: AstIdNode> Hash for FileAstId<N> {
     fn hash<H: Hasher>(&self, hasher: &mut H) {
         self.raw.hash(hasher);
     }
 }
 
-impl<N: AstNode> fmt::Debug for FileAstId<N> {
+impl<N: AstIdNode> fmt::Debug for FileAstId<N> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "FileAstId::<{}>({})", type_name::<N>(), self.raw.into_raw())
     }
 }
 
-impl<N: AstNode> FileAstId<N> {
+impl<N: AstIdNode> FileAstId<N> {
     // Can't make this a From implementation because of coherence
-    pub fn upcast<M: AstNode>(self) -> FileAstId<M>
+    pub fn upcast<M: AstIdNode>(self) -> FileAstId<M>
     where
         N: Into<M>,
     {
         FileAstId { raw: self.raw, covariant: PhantomData }
     }
+
+    pub fn erase(self) -> ErasedFileAstId {
+        self.raw
+    }
 }
 
-type ErasedFileAstId = Idx<SyntaxNodePtr>;
+pub type ErasedFileAstId = Idx<SyntaxNodePtr>;
+
+pub trait AstIdNode: AstNode {}
+macro_rules! register_ast_id_node {
+    (impl AstIdNode for $($ident:ident),+ ) => {
+        $(
+            impl AstIdNode for ast::$ident {}
+        )+
+        fn should_alloc_id(kind: syntax::SyntaxKind) -> bool {
+            $(
+                ast::$ident::can_cast(kind)
+            )||+
+        }
+    };
+}
+register_ast_id_node! {
+    impl AstIdNode for
+    Item,
+        Adt,
+            Enum,
+            Struct,
+            Union,
+        Const,
+        ExternBlock,
+        ExternCrate,
+        Fn,
+        Impl,
+        Macro,
+            MacroDef,
+            MacroRules,
+        MacroCall,
+        Module,
+        Static,
+        Trait,
+        TraitAlias,
+        TypeAlias,
+        Use,
+    AssocItem, BlockExpr, Variant, RecordField, TupleField, ConstArg
+}
 
 /// Maps items' `SyntaxNode`s to `ErasedFileAstId`s and back.
 #[derive(Default)]
@@ -92,14 +134,7 @@ impl AstIdMap {
         // change parent's id. This means that, say, adding a new function to a
         // trait does not change ids of top-level items, which helps caching.
         bdfs(node, |it| {
-            let kind = it.kind();
-            if ast::Item::can_cast(kind)
-                || ast::BlockExpr::can_cast(kind)
-                || ast::Variant::can_cast(kind)
-                || ast::RecordField::can_cast(kind)
-                || ast::TupleField::can_cast(kind)
-                || ast::ConstArg::can_cast(kind)
-            {
+            if should_alloc_id(it.kind()) {
                 res.alloc(&it);
                 true
             } else {
@@ -120,15 +155,19 @@ impl AstIdMap {
         res
     }
 
-    pub fn ast_id<N: AstNode>(&self, item: &N) -> FileAstId<N> {
+    pub fn ast_id<N: AstIdNode>(&self, item: &N) -> FileAstId<N> {
         let raw = self.erased_ast_id(item.syntax());
         FileAstId { raw, covariant: PhantomData }
     }
 
-    pub fn get<N: AstNode>(&self, id: FileAstId<N>) -> AstPtr<N> {
+    pub fn get<N: AstIdNode>(&self, id: FileAstId<N>) -> AstPtr<N> {
         AstPtr::try_from_raw(self.arena[id.raw].clone()).unwrap()
     }
 
+    pub(crate) fn get_raw(&self, id: ErasedFileAstId) -> SyntaxNodePtr {
+        self.arena[id].clone()
+    }
+
     fn erased_ast_id(&self, item: &SyntaxNode) -> ErasedFileAstId {
         let ptr = SyntaxNodePtr::new(item);
         let hash = hash_ptr(&ptr);
diff --git a/crates/hir-expand/src/lib.rs b/crates/hir-expand/src/lib.rs
index b0dc8e1b5c9..3a534f56e4f 100644
--- a/crates/hir-expand/src/lib.rs
+++ b/crates/hir-expand/src/lib.rs
@@ -37,11 +37,11 @@ use either::Either;
 use syntax::{
     algo::{self, skip_trivia_token},
     ast::{self, AstNode, HasDocComments},
-    Direction, SyntaxNode, SyntaxToken,
+    AstPtr, Direction, SyntaxNode, SyntaxNodePtr, SyntaxToken,
 };
 
 use crate::{
-    ast_id_map::FileAstId,
+    ast_id_map::{AstIdNode, ErasedFileAstId, FileAstId},
     attrs::AttrId,
     builtin_attr_macro::BuiltinAttrExpander,
     builtin_derive_macro::BuiltinDeriveExpander,
@@ -551,9 +551,9 @@ impl MacroCallKind {
         };
 
         let range = match kind {
-            MacroCallKind::FnLike { ast_id, .. } => ast_id.to_node(db).syntax().text_range(),
-            MacroCallKind::Derive { ast_id, .. } => ast_id.to_node(db).syntax().text_range(),
-            MacroCallKind::Attr { ast_id, .. } => ast_id.to_node(db).syntax().text_range(),
+            MacroCallKind::FnLike { ast_id, .. } => ast_id.to_ptr(db).text_range(),
+            MacroCallKind::Derive { ast_id, .. } => ast_id.to_ptr(db).text_range(),
+            MacroCallKind::Attr { ast_id, .. } => ast_id.to_ptr(db).text_range(),
         };
 
         FileRange { range, file_id }
@@ -801,13 +801,25 @@ impl ExpansionInfo {
 /// It is stable across reparses, and can be used as salsa key/value.
 pub type AstId<N> = InFile<FileAstId<N>>;
 
-impl<N: AstNode> AstId<N> {
+impl<N: AstIdNode> AstId<N> {
     pub fn to_node(&self, db: &dyn db::ExpandDatabase) -> N {
-        let root = db.parse_or_expand(self.file_id);
-        db.ast_id_map(self.file_id).get(self.value).to_node(&root)
+        self.to_ptr(db).to_node(&db.parse_or_expand(self.file_id))
+    }
+    pub fn to_ptr(&self, db: &dyn db::ExpandDatabase) -> AstPtr<N> {
+        db.ast_id_map(self.file_id).get(self.value)
     }
 }
 
+pub type ErasedAstId = InFile<ErasedFileAstId>;
+
+impl ErasedAstId {
+    pub fn to_node(&self, db: &dyn db::ExpandDatabase) -> SyntaxNode {
+        self.to_ptr(db).to_node(&db.parse_or_expand(self.file_id))
+    }
+    pub fn to_ptr(&self, db: &dyn db::ExpandDatabase) -> SyntaxNodePtr {
+        db.ast_id_map(self.file_id).get_raw(self.value)
+    }
+}
 /// `InFile<T>` stores a value of `T` inside a particular file/syntax tree.
 ///
 /// Typical usages are:
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 645ea42f538..3688fd0e837 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -754,7 +754,7 @@ fn emit_def_diagnostic_(
             let item = ast.to_node(db.upcast());
             acc.push(
                 InactiveCode {
-                    node: ast.with_value(AstPtr::new(&item).into()),
+                    node: ast.with_value(SyntaxNodePtr::new(&item).into()),
                     cfg: cfg.clone(),
                     opts: opts.clone(),
                 }