3 files changed, 130 insertions, 53 deletions

diff --git a/crates/ra_hir_def/src/item_tree.rs b/crates/ra_hir_def/src/item_tree.rs
index a13a989dd..9656f845e 100644
--- a/crates/ra_hir_def/src/item_tree.rs
+++ b/crates/ra_hir_def/src/item_tree.rs
@@ -3,6 +3,9 @@
 mod lower;
 
 use std::{
+    fmt::{self, Debug},
+    hash::{Hash, Hasher},
+    marker::PhantomData,
     ops::{Index, Range},
     sync::Arc,
 };
@@ -109,6 +112,68 @@ impl ItemTree {
     }
 }
 
+pub trait ItemTreeNode: Sized {
+    fn lookup(tree: &ItemTree, index: Idx<Self>) -> &Self;
+}
+
+pub struct FileItemTreeId<N: ItemTreeNode> {
+    index: Idx<N>,
+    _p: PhantomData<N>,
+}
+
+impl<N: ItemTreeNode> Clone for FileItemTreeId<N> {
+    fn clone(&self) -> Self {
+        Self { index: self.index, _p: PhantomData }
+    }
+}
+impl<N: ItemTreeNode> Copy for FileItemTreeId<N> {}
+
+impl<N: ItemTreeNode> PartialEq for FileItemTreeId<N> {
+    fn eq(&self, other: &FileItemTreeId<N>) -> bool {
+        self.index == other.index
+    }
+}
+impl<N: ItemTreeNode> Eq for FileItemTreeId<N> {}
+
+impl<N: ItemTreeNode> Hash for FileItemTreeId<N> {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.index.hash(state)
+    }
+}
+
+impl<N: ItemTreeNode> fmt::Debug for FileItemTreeId<N> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.index.fmt(f)
+    }
+}
+
+pub type ItemTreeId<N> = InFile<FileItemTreeId<N>>;
+
+macro_rules! nodes {
+    ( $($node:ident in $fld:ident),+ $(,)? ) => { $(
+        impl ItemTreeNode for $node {
+            fn lookup(tree: &ItemTree, index: Idx<Self>) -> &Self {
+                &tree.$fld[index]
+            }
+        }
+    )+ };
+}
+
+nodes!(
+    Import in imports,
+    Function in functions,
+    Struct in structs,
+    Union in unions,
+    Enum in enums,
+    Const in consts,
+    Static in statics,
+    Trait in traits,
+    Impl in impls,
+    TypeAlias in type_aliases,
+    Mod in mods,
+    MacroCall in macro_calls,
+);
+
 macro_rules! impl_index {
     ( $($fld:ident: $t:ty),+ $(,)? ) => {
         $(
@@ -141,6 +206,13 @@ impl_index!(
     exprs: Expr,
 );
 
+impl<N: ItemTreeNode> Index<FileItemTreeId<N>> for ItemTree {
+    type Output = N;
+    fn index(&self, id: FileItemTreeId<N>) -> &N {
+        N::lookup(self, id.index)
+    }
+}
+
 /// A desugared `extern crate` or `use` import.
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Import {
@@ -304,48 +376,48 @@ macro_rules! impl_froms {
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
 pub enum ModItem {
-    Import(Idx<Import>),
-    Function(Idx<Function>),
-    Struct(Idx<Struct>),
-    Union(Idx<Union>),
-    Enum(Idx<Enum>),
-    Const(Idx<Const>),
-    Static(Idx<Static>),
-    Trait(Idx<Trait>),
-    Impl(Idx<Impl>),
-    TypeAlias(Idx<TypeAlias>),
-    Mod(Idx<Mod>),
-    MacroCall(Idx<MacroCall>),
+    Import(FileItemTreeId<Import>),
+    Function(FileItemTreeId<Function>),
+    Struct(FileItemTreeId<Struct>),
+    Union(FileItemTreeId<Union>),
+    Enum(FileItemTreeId<Enum>),
+    Const(FileItemTreeId<Const>),
+    Static(FileItemTreeId<Static>),
+    Trait(FileItemTreeId<Trait>),
+    Impl(FileItemTreeId<Impl>),
+    TypeAlias(FileItemTreeId<TypeAlias>),
+    Mod(FileItemTreeId<Mod>),
+    MacroCall(FileItemTreeId<MacroCall>),
 }
 
 impl_froms!(ModItem {
-    Import(Idx<Import>),
-    Function(Idx<Function>),
-    Struct(Idx<Struct>),
-    Union(Idx<Union>),
-    Enum(Idx<Enum>),
-    Const(Idx<Const>),
-    Static(Idx<Static>),
-    Trait(Idx<Trait>),
-    Impl(Idx<Impl>),
-    TypeAlias(Idx<TypeAlias>),
-    Mod(Idx<Mod>),
-    MacroCall(Idx<MacroCall>),
+    Import(FileItemTreeId<Import>),
+    Function(FileItemTreeId<Function>),
+    Struct(FileItemTreeId<Struct>),
+    Union(FileItemTreeId<Union>),
+    Enum(FileItemTreeId<Enum>),
+    Const(FileItemTreeId<Const>),
+    Static(FileItemTreeId<Static>),
+    Trait(FileItemTreeId<Trait>),
+    Impl(FileItemTreeId<Impl>),
+    TypeAlias(FileItemTreeId<TypeAlias>),
+    Mod(FileItemTreeId<Mod>),
+    MacroCall(FileItemTreeId<MacroCall>),
 });
 
 #[derive(Debug, Eq, PartialEq)]
 pub enum AssocItem {
-    Function(Idx<Function>),
-    TypeAlias(Idx<TypeAlias>),
-    Const(Idx<Const>),
-    MacroCall(Idx<MacroCall>),
+    Function(FileItemTreeId<Function>),
+    TypeAlias(FileItemTreeId<TypeAlias>),
+    Const(FileItemTreeId<Const>),
+    MacroCall(FileItemTreeId<MacroCall>),
 }
 
 impl_froms!(AssocItem {
-    Function(Idx<Function>),
-    TypeAlias(Idx<TypeAlias>),
-    Const(Idx<Const>),
-    MacroCall(Idx<MacroCall>),
+    Function(FileItemTreeId<Function>),
+    TypeAlias(FileItemTreeId<TypeAlias>),
+    Const(FileItemTreeId<Const>),
+    MacroCall(FileItemTreeId<MacroCall>),
 });
 
 #[derive(Debug, Eq, PartialEq)]
diff --git a/crates/ra_hir_def/src/item_tree/lower.rs b/crates/ra_hir_def/src/item_tree/lower.rs
index d123a7310..0c9454848 100644
--- a/crates/ra_hir_def/src/item_tree/lower.rs
+++ b/crates/ra_hir_def/src/item_tree/lower.rs
@@ -5,6 +5,10 @@ use ra_syntax::ast::{self, ModuleItemOwner};
 use smallvec::SmallVec;
 use std::sync::Arc;
 
+fn id<N: ItemTreeNode>(index: Idx<N>) -> FileItemTreeId<N> {
+    FileItemTreeId { index, _p: PhantomData }
+}
+
 struct ModItems(SmallVec<[ModItem; 1]>);
 
 impl<T> From<T> for ModItems
@@ -38,54 +42,54 @@ impl Ctx {
         let attrs = Attrs::new(item, &self.hygiene);
         let items = match item {
             ast::ModuleItem::StructDef(ast) => {
-                self.lower_struct(ast).map(|data| self.tree.structs.alloc(data).into())
+                self.lower_struct(ast).map(|data| id(self.tree.structs.alloc(data)).into())
             }
             ast::ModuleItem::UnionDef(ast) => {
-                self.lower_union(ast).map(|data| self.tree.unions.alloc(data).into())
+                self.lower_union(ast).map(|data| id(self.tree.unions.alloc(data)).into())
             }
             ast::ModuleItem::EnumDef(ast) => {
-                self.lower_enum(ast).map(|data| self.tree.enums.alloc(data).into())
+                self.lower_enum(ast).map(|data| id(self.tree.enums.alloc(data)).into())
             }
             ast::ModuleItem::FnDef(ast) => {
-                self.lower_function(ast).map(|data| self.tree.functions.alloc(data).into())
+                self.lower_function(ast).map(|data| id(self.tree.functions.alloc(data)).into())
             }
             ast::ModuleItem::TypeAliasDef(ast) => {
-                self.lower_type_alias(ast).map(|data| self.tree.type_aliases.alloc(data).into())
+                self.lower_type_alias(ast).map(|data| id(self.tree.type_aliases.alloc(data)).into())
             }
             ast::ModuleItem::StaticDef(ast) => {
-                self.lower_static(ast).map(|data| self.tree.statics.alloc(data).into())
+                self.lower_static(ast).map(|data| id(self.tree.statics.alloc(data)).into())
             }
             ast::ModuleItem::ConstDef(ast) => {
                 let data = self.lower_const(ast);
-                Some(self.tree.consts.alloc(data).into())
+                Some(id(self.tree.consts.alloc(data)).into())
             }
             ast::ModuleItem::Module(ast) => {
-                self.lower_module(ast).map(|data| self.tree.mods.alloc(data).into())
+                self.lower_module(ast).map(|data| id(self.tree.mods.alloc(data)).into())
             }
             ast::ModuleItem::TraitDef(ast) => {
-                self.lower_trait(ast).map(|data| self.tree.traits.alloc(data).into())
+                self.lower_trait(ast).map(|data| id(self.tree.traits.alloc(data)).into())
             }
             ast::ModuleItem::ImplDef(ast) => {
-                self.lower_impl(ast).map(|data| self.tree.impls.alloc(data).into())
+                self.lower_impl(ast).map(|data| id(self.tree.impls.alloc(data)).into())
             }
             ast::ModuleItem::UseItem(ast) => Some(ModItems(
                 self.lower_use(ast)
                     .into_iter()
-                    .map(|data| self.tree.imports.alloc(data).into())
+                    .map(|data| id(self.tree.imports.alloc(data)).into())
                     .collect::<SmallVec<_>>(),
             )),
             ast::ModuleItem::ExternCrateItem(ast) => {
-                self.lower_extern_crate(ast).map(|data| self.tree.imports.alloc(data).into())
+                self.lower_extern_crate(ast).map(|data| id(self.tree.imports.alloc(data)).into())
             }
             ast::ModuleItem::MacroCall(ast) => {
-                self.lower_macro_call(ast).map(|data| self.tree.macro_calls.alloc(data).into())
+                self.lower_macro_call(ast).map(|data| id(self.tree.macro_calls.alloc(data)).into())
             }
             ast::ModuleItem::ExternBlock(ast) => Some(ModItems(
                 self.lower_extern_block(ast)
                     .into_iter()
                     .map(|item| match item {
-                        Either::Left(func) => self.tree.functions.alloc(func).into(),
-                        Either::Right(statik) => self.tree.statics.alloc(statik).into(),
+                        Either::Left(func) => id(self.tree.functions.alloc(func)).into(),
+                        Either::Right(statik) => id(self.tree.statics.alloc(statik)).into(),
                     })
                     .collect::<SmallVec<_>>(),
             )),
@@ -103,14 +107,14 @@ impl Ctx {
     fn lower_assoc_item(&mut self, item: &ast::AssocItem) -> Option<AssocItem> {
         match item {
             ast::AssocItem::FnDef(ast) => {
-                self.lower_function(ast).map(|data| self.tree.functions.alloc(data).into())
+                self.lower_function(ast).map(|data| id(self.tree.functions.alloc(data)).into())
             }
             ast::AssocItem::TypeAliasDef(ast) => {
-                self.lower_type_alias(ast).map(|data| self.tree.type_aliases.alloc(data).into())
+                self.lower_type_alias(ast).map(|data| id(self.tree.type_aliases.alloc(data)).into())
             }
             ast::AssocItem::ConstDef(ast) => {
                 let data = self.lower_const(ast);
-                Some(self.tree.consts.alloc(data).into())
+                Some(id(self.tree.consts.alloc(data)).into())
             }
         }
     }
diff --git a/crates/ra_hir_def/src/nameres/collector.rs b/crates/ra_hir_def/src/nameres/collector.rs
index b899a5fb3..c227b6da1 100644
--- a/crates/ra_hir_def/src/nameres/collector.rs
+++ b/crates/ra_hir_def/src/nameres/collector.rs
@@ -20,7 +20,9 @@ use test_utils::mark;
 use crate::{
     attr::Attrs,
     db::DefDatabase,
-    item_tree::{Import, ItemTree, MacroCall, Mod, ModItem, ModKind, StructDefKind},
+    item_tree::{
+        FileItemTreeId, Import, ItemTree, MacroCall, Mod, ModItem, ModKind, StructDefKind,
+    },
     nameres::{
         diagnostics::DefDiagnostic, mod_resolution::ModDir, path_resolution::ReachedFixedPoint,
         BuiltinShadowMode, CrateDefMap, ModuleData, ModuleOrigin, ResolveMode,
@@ -32,7 +34,6 @@ use crate::{
     FunctionLoc, ImplLoc, Intern, LocalModuleId, ModuleDefId, ModuleId, StaticLoc, StructLoc,
     TraitLoc, TypeAliasLoc, UnionLoc,
 };
-use ra_arena::Idx;
 
 pub(super) fn collect_defs(db: &dyn DefDatabase, mut def_map: CrateDefMap) -> CrateDefMap {
     let crate_graph = db.crate_graph();
@@ -107,7 +108,7 @@ impl PartialResolvedImport {
 #[derive(Clone, Debug, Eq, PartialEq)]
 struct ImportDirective {
     module_id: LocalModuleId,
-    import_id: Idx<Import>,
+    import_id: FileItemTreeId<Import>,
     import: Import,
     status: PartialResolvedImport,
 }