7 files changed, 88 insertions, 45 deletions

diff --git a/crates/ra_hir/src/code_model_impl/module.rs b/crates/ra_hir/src/code_model_impl/module.rs
index 790e2b80f..0edb8ade5 100644
--- a/crates/ra_hir/src/code_model_impl/module.rs
+++ b/crates/ra_hir/src/code_model_impl/module.rs
@@ -1,18 +1,18 @@
 use ra_db::FileId;
-use ra_syntax::{ast, TreeArc, AstNode};
+use ra_syntax::{ast, TreeArc};
 
 use crate::{
-    Module, ModuleSource, Name,
+    Module, ModuleSource, Name, AstId,
     nameres::{CrateModuleId, ImportId},
     HirDatabase, DefDatabase,
-    HirFileId, SourceItemId,
+    HirFileId,
 };
 
 impl ModuleSource {
     pub(crate) fn new(
         db: &impl DefDatabase,
         file_id: Option<FileId>,
-        decl_id: Option<SourceItemId>,
+        decl_id: Option<AstId<ast::Module>>,
     ) -> ModuleSource {
         match (file_id, decl_id) {
             (Some(file_id), _) => {
@@ -20,8 +20,7 @@ impl ModuleSource {
                 ModuleSource::SourceFile(source_file)
             }
             (None, Some(item_id)) => {
-                let module = db.file_item(item_id);
-                let module = ast::Module::cast(&*module).unwrap();
+                let module = item_id.to_node(db);
                 assert!(module.item_list().is_some(), "expected inline module");
                 ModuleSource::Module(module.to_owned())
             }
@@ -55,7 +54,7 @@ impl Module {
         let decl_id = def_map[self.module_id].declaration;
         let file_id = def_map[self.module_id].definition;
         let module_source = ModuleSource::new(db, file_id, decl_id);
-        let file_id = file_id.map(HirFileId::from).unwrap_or_else(|| decl_id.unwrap().file_id);
+        let file_id = file_id.map(HirFileId::from).unwrap_or_else(|| decl_id.unwrap().file_id());
         (file_id, module_source)
     }
 
@@ -65,9 +64,8 @@ impl Module {
     ) -> Option<(HirFileId, TreeArc<ast::Module>)> {
         let def_map = db.crate_def_map(self.krate);
         let decl = def_map[self.module_id].declaration?;
-        let syntax_node = db.file_item(decl);
-        let ast = ast::Module::cast(&syntax_node).unwrap().to_owned();
-        Some((decl.file_id, ast))
+        let ast = decl.to_node(db);
+        Some((decl.file_id(), ast))
     }
 
     pub(crate) fn import_source_impl(
diff --git a/crates/ra_hir/src/lib.rs b/crates/ra_hir/src/lib.rs
index ac2585de0..3fc08c55c 100644
--- a/crates/ra_hir/src/lib.rs
+++ b/crates/ra_hir/src/lib.rs
@@ -48,7 +48,7 @@ mod marks;
 use crate::{
     db::{HirDatabase, DefDatabase},
     name::{AsName, KnownName},
-    source_id::SourceFileItemId,
+    source_id::{SourceFileItemId, FileAstId, AstId},
 };
 
 pub use self::{
diff --git a/crates/ra_hir/src/nameres.rs b/crates/ra_hir/src/nameres.rs
index e962bbd31..67b9d6986 100644
--- a/crates/ra_hir/src/nameres.rs
+++ b/crates/ra_hir/src/nameres.rs
@@ -59,13 +59,15 @@ use rustc_hash::FxHashMap;
 use ra_arena::{Arena, RawId, impl_arena_id};
 use ra_db::{FileId, Edition};
 use test_utils::tested_by;
+use ra_syntax::ast;
 
 use crate::{
-    ModuleDef, Name, Crate, Module, SourceItemId,
+    ModuleDef, Name, Crate, Module,
     DefDatabase, Path, PathKind, HirFileId, Trait,
     ids::MacroDefId,
     diagnostics::DiagnosticSink,
     nameres::diagnostics::DefDiagnostic,
+    AstId,
 };
 
 pub(crate) use self::raw::{RawItems, ImportId, ImportSourceMap};
@@ -106,7 +108,7 @@ pub(crate) struct ModuleData {
     pub(crate) children: FxHashMap<Name, CrateModuleId>,
     pub(crate) scope: ModuleScope,
     /// None for root
-    pub(crate) declaration: Option<SourceItemId>,
+    pub(crate) declaration: Option<AstId<ast::Module>>,
     /// None for inline modules.
     ///
     /// Note that non-inline modules, by definition, live inside non-macro file.
@@ -225,7 +227,7 @@ impl CrateDefMap {
     pub(crate) fn find_module_by_source(
         &self,
         file_id: HirFileId,
-        decl_id: Option<SourceItemId>,
+        decl_id: Option<AstId<ast::Module>>,
     ) -> Option<CrateModuleId> {
         let (module_id, _module_data) = self.modules.iter().find(|(_module_id, module_data)| {
             if decl_id.is_some() {
@@ -429,10 +431,10 @@ impl CrateDefMap {
 
 mod diagnostics {
     use relative_path::RelativePathBuf;
-    use ra_syntax::{AstPtr, AstNode, ast};
+    use ra_syntax::{AstPtr, ast};
 
     use crate::{
-        SourceItemId, DefDatabase,
+        AstId, DefDatabase,
         nameres::CrateModuleId,
         diagnostics::{DiagnosticSink, UnresolvedModule},
 };
@@ -441,7 +443,7 @@ mod diagnostics {
     pub(super) enum DefDiagnostic {
         UnresolvedModule {
             module: CrateModuleId,
-            declaration: SourceItemId,
+            declaration: AstId<ast::Module>,
             candidate: RelativePathBuf,
         },
     }
@@ -458,10 +460,9 @@ mod diagnostics {
                     if *module != target_module {
                         return;
                     }
-                    let syntax = db.file_item(*declaration);
-                    let decl = ast::Module::cast(&syntax).unwrap();
+                    let decl = declaration.to_node(db);
                     sink.push(UnresolvedModule {
-                        file: declaration.file_id,
+                        file: declaration.file_id(),
                         decl: AstPtr::new(&decl),
                         candidate: candidate.clone(),
                     })
@@ -469,5 +470,4 @@ mod diagnostics {
             }
         }
     }
-
 }
diff --git a/crates/ra_hir/src/nameres/collector.rs b/crates/ra_hir/src/nameres/collector.rs
index 4fb298155..e6fd8632a 100644
--- a/crates/ra_hir/src/nameres/collector.rs
+++ b/crates/ra_hir/src/nameres/collector.rs
@@ -3,6 +3,7 @@ use rustc_hash::FxHashMap;
 use relative_path::RelativePathBuf;
 use test_utils::tested_by;
 use ra_db::FileId;
+use ra_syntax::ast;
 
 use crate::{
     Function, Module, Struct, Enum, Const, Static, Trait, TypeAlias,
@@ -15,6 +16,7 @@ use crate::{
         raw,
     },
     ids::{AstItemDef, LocationCtx, MacroCallLoc, MacroCallId, MacroDefId},
+    AstId,
 };
 
 pub(super) fn collect_defs(db: &impl DefDatabase, mut def_map: CrateDefMap) -> CrateDefMap {
@@ -364,12 +366,9 @@ where
     fn collect_module(&mut self, module: &raw::ModuleData) {
         match module {
             // inline module, just recurse
-            raw::ModuleData::Definition { name, items, source_item_id } => {
-                let module_id = self.push_child_module(
-                    name.clone(),
-                    source_item_id.with_file_id(self.file_id),
-                    None,
-                );
+            raw::ModuleData::Definition { name, items, ast_id } => {
+                let module_id =
+                    self.push_child_module(name.clone(), ast_id.with_file_id(self.file_id), None);
                 ModCollector {
                     def_collector: &mut *self.def_collector,
                     module_id,
@@ -379,13 +378,12 @@ where
                 .collect(&*items);
             }
             // out of line module, resovle, parse and recurse
-            raw::ModuleData::Declaration { name, source_item_id } => {
-                let source_item_id = source_item_id.with_file_id(self.file_id);
+            raw::ModuleData::Declaration { name, ast_id } => {
+                let ast_id = ast_id.with_file_id(self.file_id);
                 let is_root = self.def_collector.def_map.modules[self.module_id].parent.is_none();
                 match resolve_submodule(self.def_collector.db, self.file_id, name, is_root) {
                     Ok(file_id) => {
-                        let module_id =
-                            self.push_child_module(name.clone(), source_item_id, Some(file_id));
+                        let module_id = self.push_child_module(name.clone(), ast_id, Some(file_id));
                         let raw_items = self.def_collector.db.raw_items(file_id.into());
                         ModCollector {
                             def_collector: &mut *self.def_collector,
@@ -398,7 +396,7 @@ where
                     Err(candidate) => self.def_collector.def_map.diagnostics.push(
                         DefDiagnostic::UnresolvedModule {
                             module: self.module_id,
-                            declaration: source_item_id,
+                            declaration: ast_id,
                             candidate,
                         },
                     ),
@@ -410,7 +408,7 @@ where
     fn push_child_module(
         &mut self,
         name: Name,
-        declaration: SourceItemId,
+        declaration: AstId<ast::Module>,
         definition: Option<FileId>,
     ) -> CrateModuleId {
         let modules = &mut self.def_collector.def_map.modules;
diff --git a/crates/ra_hir/src/nameres/raw.rs b/crates/ra_hir/src/nameres/raw.rs
index f32004601..09acd5a98 100644
--- a/crates/ra_hir/src/nameres/raw.rs
+++ b/crates/ra_hir/src/nameres/raw.rs
@@ -12,7 +12,7 @@ use ra_syntax::{
 
 use crate::{
     DefDatabase, Name, AsName, Path, HirFileId, ModuleSource,
-    SourceFileItemId, SourceFileItems,
+    SourceFileItemId, SourceFileItems, FileAstId,
 };
 
 /// `RawItems` is a set of top-level items in a file (except for impls).
@@ -115,8 +115,8 @@ impl_arena_id!(Module);
 
 #[derive(Debug, PartialEq, Eq)]
 pub(super) enum ModuleData {
-    Declaration { name: Name, source_item_id: SourceFileItemId },
-    Definition { name: Name, source_item_id: SourceFileItemId, items: Vec<RawItem> },
+    Declaration { name: Name, ast_id: FileAstId<ast::Module> },
+    Definition { name: Name, ast_id: FileAstId<ast::Module>, items: Vec<RawItem> },
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -221,10 +221,9 @@ impl RawItemsCollector {
             Some(it) => it.as_name(),
             None => return,
         };
-        let source_item_id = self.source_file_items.id_of_unchecked(module.syntax());
+        let ast_id = self.source_file_items.ast_id(module);
         if module.has_semi() {
-            let item =
-                self.raw_items.modules.alloc(ModuleData::Declaration { name, source_item_id });
+            let item = self.raw_items.modules.alloc(ModuleData::Declaration { name, ast_id });
             self.push_item(current_module, RawItem::Module(item));
             return;
         }
@@ -232,7 +231,7 @@ impl RawItemsCollector {
         if let Some(item_list) = module.item_list() {
             let item = self.raw_items.modules.alloc(ModuleData::Definition {
                 name,
-                source_item_id,
+                ast_id,
                 items: Vec::new(),
             });
             self.process_module(Some(item), item_list);
diff --git a/crates/ra_hir/src/source_binder.rs b/crates/ra_hir/src/source_binder.rs
index db9e3a22e..f9d2d0247 100644
--- a/crates/ra_hir/src/source_binder.rs
+++ b/crates/ra_hir/src/source_binder.rs
@@ -13,10 +13,10 @@ use ra_syntax::{
 };
 
 use crate::{
-    HirDatabase, Function, Struct, Enum, SourceItemId,
+    HirDatabase, Function, Struct, Enum,
     AsName, Module, HirFileId, Crate, Trait, Resolver,
     ids::LocationCtx,
-    expr
+    expr, AstId
 };
 
 /// Locates the module by `FileId`. Picks topmost module in the file.
@@ -55,7 +55,7 @@ fn module_from_inline(
     assert!(!module.has_semi());
     let file_id = file_id.into();
     let file_items = db.file_items(file_id);
-    let item_id = file_items.id_of(file_id, module.syntax()).with_file_id(file_id);
+    let item_id = file_items.ast_id(module).with_file_id(file_id);
     module_from_source(db, file_id, Some(item_id))
 }
 
@@ -75,7 +75,7 @@ pub fn module_from_child_node(
 fn module_from_source(
     db: &impl HirDatabase,
     file_id: HirFileId,
-    decl_id: Option<SourceItemId>,
+    decl_id: Option<AstId<ast::Module>>,
 ) -> Option<Module> {
     let source_root_id = db.file_source_root(file_id.as_original_file());
     db.source_root_crates(source_root_id).iter().map(|&crate_id| Crate { crate_id }).find_map(
diff --git a/crates/ra_hir/src/source_id.rs b/crates/ra_hir/src/source_id.rs
index 62707ba6a..881cc590e 100644
--- a/crates/ra_hir/src/source_id.rs
+++ b/crates/ra_hir/src/source_id.rs
@@ -1,10 +1,55 @@
-use std::sync::Arc;
+use std::{marker::PhantomData, sync::Arc};
 
 use ra_arena::{Arena, RawId, impl_arena_id};
 use ra_syntax::{SyntaxNodePtr, TreeArc, SyntaxNode, SourceFile, AstNode, ast};
 
 use crate::{HirFileId, DefDatabase};
 
+#[derive(Debug, PartialEq, Eq)]
+pub(crate) struct AstId<N: AstNode> {
+    file_id: HirFileId,
+    file_ast_id: FileAstId<N>,
+}
+
+impl<N: AstNode> Clone for AstId<N> {
+    fn clone(&self) -> AstId<N> {
+        *self
+    }
+}
+
+impl<N: AstNode> Copy for AstId<N> {}
+
+impl<N: AstNode> AstId<N> {
+    pub(crate) fn file_id(&self) -> HirFileId {
+        self.file_id
+    }
+
+    pub(crate) fn to_node(&self, db: &impl DefDatabase) -> TreeArc<N> {
+        let syntax_node = db.file_item(self.file_ast_id.raw.with_file_id(self.file_id));
+        N::cast(&syntax_node).unwrap().to_owned()
+    }
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub(crate) struct FileAstId<N: AstNode> {
+    raw: SourceFileItemId,
+    _ty: PhantomData<N>,
+}
+
+impl<N: AstNode> Clone for FileAstId<N> {
+    fn clone(&self) -> FileAstId<N> {
+        *self
+    }
+}
+
+impl<N: AstNode> Copy for FileAstId<N> {}
+
+impl<N: AstNode> FileAstId<N> {
+    pub(crate) fn with_file_id(self, file_id: HirFileId) -> AstId<N> {
+        AstId { file_id, file_ast_id: self }
+    }
+}
+
 /// Identifier of item within a specific file. This is stable over reparses, so
 /// it's OK to use it as a salsa key/value.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -90,6 +135,9 @@ impl SourceFileItems {
             self.arena.iter().map(|(_id, i)| i).collect::<Vec<_>>(),
         );
     }
+    pub(crate) fn ast_id<N: AstNode>(&self, item: &N) -> FileAstId<N> {
+        FileAstId { raw: self.id_of_unchecked(item.syntax()), _ty: PhantomData }
+    }
 }
 
 impl std::ops::Index<SourceFileItemId> for SourceFileItems {