8 files changed, 101 insertions, 134 deletions

diff --git a/crates/ra_hir/src/code_model.rs b/crates/ra_hir/src/code_model.rs
index cb990f4e2..c49190a0f 100644
--- a/crates/ra_hir/src/code_model.rs
+++ b/crates/ra_hir/src/code_model.rs
@@ -10,8 +10,10 @@ use hir_def::{
     adt::VariantData,
     body::scope::ExprScopes,
     builtin_type::BuiltinType,
+    traits::TraitData,
     type_ref::{Mutability, TypeRef},
-    CrateModuleId, ImplId, LocalEnumVariantId, LocalStructFieldId, ModuleId, UnionId,
+    AssocItemId, CrateModuleId, FunctionContainerId, HasModule, ImplId, LocalEnumVariantId,
+    LocalStructFieldId, Lookup, ModuleId, UnionId,
 };
 use hir_expand::{
     diagnostics::DiagnosticSink,
@@ -30,7 +32,6 @@ use crate::{
         TypeAliasId,
     },
     resolve::{Resolver, Scope, TypeNs},
-    traits::TraitData,
     ty::{InferenceResult, Namespace, TraitRef},
     Either, HasSource, ImportId, Name, ScopeDef, Source, Ty,
 };
@@ -230,15 +231,7 @@ impl Module {
 
     pub fn declarations(self, db: &impl DefDatabase) -> Vec<ModuleDef> {
         let def_map = db.crate_def_map(self.id.krate);
-        def_map[self.id.module_id]
-            .scope
-            .entries()
-            .filter_map(|(_name, res)| if res.import.is_none() { Some(res.def) } else { None })
-            .flat_map(|per_ns| {
-                per_ns.take_types().into_iter().chain(per_ns.take_values().into_iter())
-            })
-            .map(ModuleDef::from)
-            .collect()
+        def_map[self.id.module_id].scope.declarations().map(ModuleDef::from).collect()
     }
 
     pub fn impl_blocks(self, db: &impl DefDatabase) -> Vec<ImplBlock> {
@@ -655,7 +648,7 @@ impl FnData {
 
 impl Function {
     pub fn module(self, db: &impl DefDatabase) -> Module {
-        Module { id: self.id.module(db) }
+        self.id.lookup(db).module(db).into()
     }
 
     pub fn krate(self, db: &impl DefDatabase) -> Option<Crate> {
@@ -688,21 +681,25 @@ impl Function {
 
     /// The containing impl block, if this is a method.
     pub fn impl_block(self, db: &impl DefDatabase) -> Option<ImplBlock> {
-        ImplBlock::containing(db, self.into())
+        match self.container(db) {
+            Some(Container::ImplBlock(it)) => Some(it),
+            _ => None,
+        }
     }
 
     /// The containing trait, if this is a trait method definition.
     pub fn parent_trait(self, db: &impl DefDatabase) -> Option<Trait> {
-        db.trait_items_index(self.module(db)).get_parent_trait(self.into())
+        match self.container(db) {
+            Some(Container::Trait(it)) => Some(it),
+            _ => None,
+        }
     }
 
     pub fn container(self, db: &impl DefDatabase) -> Option<Container> {
-        if let Some(impl_block) = self.impl_block(db) {
-            Some(impl_block.into())
-        } else if let Some(trait_) = self.parent_trait(db) {
-            Some(trait_.into())
-        } else {
-            None
+        match self.id.lookup(db).container {
+            FunctionContainerId::TraitId(it) => Some(Container::Trait(it.into())),
+            FunctionContainerId::ImplId(it) => Some(Container::ImplBlock(it.into())),
+            FunctionContainerId::ModuleId(_) => None,
         }
     }
 
@@ -757,7 +754,7 @@ impl Const {
     }
 
     pub fn parent_trait(self, db: &impl DefDatabase) -> Option<Trait> {
-        db.trait_items_index(self.module(db)).get_parent_trait(self.into())
+        db.trait_items_index(self.module(db).id).get_parent_trait(self.id.into()).map(Trait::from)
     }
 
     pub fn container(self, db: &impl DefDatabase) -> Option<Container> {
@@ -861,11 +858,11 @@ impl Trait {
     }
 
     pub fn name(self, db: &impl DefDatabase) -> Option<Name> {
-        self.trait_data(db).name().clone()
+        self.trait_data(db).name.clone()
     }
 
     pub fn items(self, db: &impl DefDatabase) -> Vec<AssocItem> {
-        self.trait_data(db).items().to_vec()
+        self.trait_data(db).items.iter().map(|it| (*it).into()).collect()
     }
 
     fn direct_super_traits(self, db: &impl HirDatabase) -> Vec<Trait> {
@@ -912,10 +909,10 @@ impl Trait {
     pub fn associated_type_by_name(self, db: &impl DefDatabase, name: &Name) -> Option<TypeAlias> {
         let trait_data = self.trait_data(db);
         trait_data
-            .items()
+            .items
             .iter()
             .filter_map(|item| match item {
-                AssocItem::TypeAlias(t) => Some(*t),
+                AssocItemId::TypeAliasId(t) => Some(TypeAlias::from(*t)),
                 _ => None,
             })
             .find(|t| &t.name(db) == name)
@@ -930,7 +927,7 @@ impl Trait {
     }
 
     pub(crate) fn trait_data(self, db: &impl DefDatabase) -> Arc<TraitData> {
-        db.trait_data(self)
+        db.trait_data(self.id)
     }
 
     pub fn trait_ref(self, db: &impl HirDatabase) -> TraitRef {
@@ -938,7 +935,7 @@ impl Trait {
     }
 
     pub fn is_auto(self, db: &impl DefDatabase) -> bool {
-        self.trait_data(db).is_auto()
+        self.trait_data(db).auto
     }
 
     pub(crate) fn resolver(self, db: &impl DefDatabase) -> Resolver {
@@ -971,7 +968,7 @@ impl TypeAlias {
 
     /// The containing trait, if this is a trait method definition.
     pub fn parent_trait(self, db: &impl DefDatabase) -> Option<Trait> {
-        db.trait_items_index(self.module(db)).get_parent_trait(self.into())
+        db.trait_items_index(self.module(db).id).get_parent_trait(self.id.into()).map(Trait::from)
     }
 
     pub fn container(self, db: &impl DefDatabase) -> Option<Container> {
diff --git a/crates/ra_hir/src/code_model/src.rs b/crates/ra_hir/src/code_model/src.rs
index 556417b0f..91cab7414 100644
--- a/crates/ra_hir/src/code_model/src.rs
+++ b/crates/ra_hir/src/code_model/src.rs
@@ -1,5 +1,6 @@
 //! FIXME: write short doc here
 
+use hir_def::{HasSource as _, Lookup};
 use ra_syntax::ast::{self, AstNode};
 
 use crate::{
@@ -113,7 +114,7 @@ impl HasSource for EnumVariant {
 impl HasSource for Function {
     type Ast = ast::FnDef;
     fn source(self, db: &(impl DefDatabase + AstDatabase)) -> Source<ast::FnDef> {
-        self.id.source(db)
+        self.id.lookup(db).source(db)
     }
 }
 impl HasSource for Const {
diff --git a/crates/ra_hir/src/db.rs b/crates/ra_hir/src/db.rs
index 0d35014a0..73859e1e9 100644
--- a/crates/ra_hir/src/db.rs
+++ b/crates/ra_hir/src/db.rs
@@ -11,7 +11,6 @@ use crate::{
     generics::{GenericDef, GenericParams},
     ids,
     lang_item::{LangItemTarget, LangItems},
-    traits::TraitData,
     ty::{
         method_resolution::CrateImplBlocks,
         traits::{AssocTyValue, Impl},
@@ -26,7 +25,8 @@ use crate::{
 pub use hir_def::db::{
     BodyQuery, BodyWithSourceMapQuery, CrateDefMapQuery, DefDatabase2, DefDatabase2Storage,
     EnumDataQuery, ExprScopesQuery, ImplDataQuery, InternDatabase, InternDatabaseStorage,
-    RawItemsQuery, RawItemsWithSourceMapQuery, StructDataQuery,
+    RawItemsQuery, RawItemsWithSourceMapQuery, StructDataQuery, TraitDataQuery,
+    TraitItemsIndexQuery,
 };
 pub use hir_expand::db::{
     AstDatabase, AstDatabaseStorage, AstIdMapQuery, MacroArgQuery, MacroDefQuery, MacroExpandQuery,
@@ -37,12 +37,6 @@ pub use hir_expand::db::{
 #[salsa::query_group(DefDatabaseStorage)]
 #[salsa::requires(AstDatabase)]
 pub trait DefDatabase: HirDebugDatabase + DefDatabase2 {
-    #[salsa::invoke(crate::traits::TraitData::trait_data_query)]
-    fn trait_data(&self, t: Trait) -> Arc<TraitData>;
-
-    #[salsa::invoke(crate::traits::TraitItemsIndex::trait_items_index)]
-    fn trait_items_index(&self, module: Module) -> crate::traits::TraitItemsIndex;
-
     #[salsa::invoke(crate::generics::generic_params_query)]
     fn generic_params(&self, def: GenericDef) -> Arc<GenericParams>;
 
diff --git a/crates/ra_hir/src/from_source.rs b/crates/ra_hir/src/from_source.rs
index f4dca25cb..303d5f138 100644
--- a/crates/ra_hir/src/from_source.rs
+++ b/crates/ra_hir/src/from_source.rs
@@ -4,15 +4,15 @@ use hir_def::{ModuleId, StructId, StructOrUnionId, UnionId};
 use hir_expand::{name::AsName, AstId, MacroDefId, MacroDefKind};
 use ra_syntax::{
     ast::{self, AstNode, NameOwner},
-    match_ast,
+    match_ast, AstPtr,
 };
 
 use crate::{
     db::{AstDatabase, DefDatabase, HirDatabase},
     ids::{AstItemDef, LocationCtx},
-    Const, DefWithBody, Enum, EnumVariant, FieldSource, Function, HasBody, HasSource, ImplBlock,
-    Local, MacroDef, Module, ModuleSource, Source, Static, Struct, StructField, Trait, TypeAlias,
-    Union, VariantDef,
+    AssocItem, Const, DefWithBody, Enum, EnumVariant, FieldSource, Function, HasBody, HasSource,
+    ImplBlock, Local, MacroDef, Module, ModuleDef, ModuleSource, Source, Static, Struct,
+    StructField, Trait, TypeAlias, Union, VariantDef,
 };
 
 pub trait FromSource: Sized {
@@ -52,10 +52,51 @@ impl FromSource for Trait {
 impl FromSource for Function {
     type Ast = ast::FnDef;
     fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self> {
-        let id = from_source(db, src)?;
-        Some(Function { id })
+        // FIXME: this doesn't try to handle nested declarations
+        for container in src.value.syntax().ancestors() {
+            let res = match_ast! {
+                match container {
+                    ast::TraitDef(it) => {
+                        let c = Trait::from_source(db, src.with_value(it))?;
+                        c.items(db)
+                            .into_iter()
+                            .filter_map(|it| match it {
+                                AssocItem::Function(it) => Some(it),
+                                _ => None
+                            })
+                            .find(|it| same_source(&it.source(db), &src))?
+                    },
+                    ast::ImplBlock(it) => {
+                        let c = ImplBlock::from_source(db, src.with_value(it))?;
+                        c.items(db)
+                            .into_iter()
+                            .filter_map(|it| match it {
+                                AssocItem::Function(it) => Some(it),
+                                _ => None
+                            })
+                            .find(|it| same_source(&it.source(db), &src))?
+
+                     },
+                    _ => { continue },
+                }
+            };
+            return Some(res);
+        }
+
+        let module_source = ModuleSource::from_child_node(db, src.as_ref().map(|it| it.syntax()));
+        let c = Module::from_definition(db, src.with_value(module_source))?;
+        let res = c
+            .declarations(db)
+            .into_iter()
+            .filter_map(|it| match it {
+                ModuleDef::Function(it) => Some(it),
+                _ => None,
+            })
+            .find(|it| same_source(&it.source(db), &src));
+        res
     }
 }
+
 impl FromSource for Const {
     type Ast = ast::ConstDef;
     fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self> {
@@ -108,7 +149,7 @@ impl FromSource for EnumVariant {
         let parent_enum = src.value.parent_enum();
         let src_enum = Source { file_id: src.file_id, value: parent_enum };
         let variants = Enum::from_source(db, src_enum)?.variants(db);
-        variants.into_iter().find(|v| v.source(db) == src)
+        variants.into_iter().find(|v| same_source(&v.source(db), &src))
     }
 }
 
@@ -216,3 +257,13 @@ where
     let ctx = LocationCtx::new(db, module.id, src.file_id);
     Some(DEF::from_ast(ctx, &src.value))
 }
+
+/// XXX: AST Nodes and SyntaxNodes have identity equality semantics: nodes are
+/// equal if they point to exactly the same object.
+///
+/// In general, we do not guarantee that we have exactly one instance of a
+/// syntax tree for each file. We probably should add such guanratree, but, for
+/// the time being, we will use identity-less AstPtr comparison.
+fn same_source<N: AstNode>(s1: &Source<N>, s2: &Source<N>) -> bool {
+    s1.as_ref().map(AstPtr::new) == s2.as_ref().map(AstPtr::new)
+}
diff --git a/crates/ra_hir/src/lib.rs b/crates/ra_hir/src/lib.rs
index da33c9591..7ac9a9041 100644
--- a/crates/ra_hir/src/lib.rs
+++ b/crates/ra_hir/src/lib.rs
@@ -33,7 +33,6 @@ pub mod source_binder;
 
 mod ids;
 mod adt;
-mod traits;
 mod type_alias;
 mod ty;
 mod impl_block;
diff --git a/crates/ra_hir/src/source_binder.rs b/crates/ra_hir/src/source_binder.rs
index a1c1daacd..caa8d0082 100644
--- a/crates/ra_hir/src/source_binder.rs
+++ b/crates/ra_hir/src/source_binder.rs
@@ -70,7 +70,7 @@ fn def_with_body_from_child_node(
     child.value.ancestors().find_map(|node| {
         match_ast! {
             match node {
-                ast::FnDef(def)  => { Some(Function {id: ctx.to_def(&def) }.into()) },
+                ast::FnDef(def)  => { return Function::from_source(db, child.with_value(def)).map(DefWithBody::from); },
                 ast::ConstDef(def) => { Some(Const { id: ctx.to_def(&def) }.into()) },
                 ast::StaticDef(def) => { Some(Static { id: ctx.to_def(&def) }.into()) },
                 _ => { None },
@@ -405,9 +405,16 @@ impl SourceAnalyzer {
         implements_trait(&canonical_ty, db, &self.resolver, krate, std_future_trait)
     }
 
-    pub fn expand(&self, db: &impl HirDatabase, macro_call: &ast::MacroCall) -> Option<Expansion> {
-        let def = self.resolve_macro_call(db, macro_call)?.id;
-        let ast_id = AstId::new(self.file_id, db.ast_id_map(self.file_id).ast_id(macro_call));
+    pub fn expand(
+        &self,
+        db: &impl HirDatabase,
+        macro_call: Source<&ast::MacroCall>,
+    ) -> Option<Expansion> {
+        let def = self.resolve_macro_call(db, macro_call.value)?.id;
+        let ast_id = AstId::new(
+            macro_call.file_id,
+            db.ast_id_map(macro_call.file_id).ast_id(macro_call.value),
+        );
         let macro_call_loc = MacroCallLoc { def, ast_id };
         Some(Expansion { macro_call_id: db.intern_macro(macro_call_loc) })
     }
diff --git a/crates/ra_hir/src/traits.rs b/crates/ra_hir/src/traits.rs
deleted file mode 100644
index 858972c6f..000000000
--- a/crates/ra_hir/src/traits.rs
+++ /dev/null
@@ -1,82 +0,0 @@
-//! HIR for trait definitions.
-
-use std::sync::Arc;
-
-use hir_expand::name::AsName;
-
-use ra_syntax::ast::{self, NameOwner};
-use rustc_hash::FxHashMap;
-
-use crate::{
-    db::{AstDatabase, DefDatabase},
-    ids::LocationCtx,
-    AssocItem, Const, Function, HasSource, Module, Name, Trait, TypeAlias,
-};
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct TraitData {
-    name: Option<Name>,
-    items: Vec<AssocItem>,
-    auto: bool,
-}
-
-impl TraitData {
-    pub(crate) fn trait_data_query(
-        db: &(impl DefDatabase + AstDatabase),
-        tr: Trait,
-    ) -> Arc<TraitData> {
-        let src = tr.source(db);
-        let name = src.value.name().map(|n| n.as_name());
-        let module = tr.module(db);
-        let ctx = LocationCtx::new(db, module.id, src.file_id);
-        let auto = src.value.is_auto();
-        let items = if let Some(item_list) = src.value.item_list() {
-            item_list
-                .impl_items()
-                .map(|item_node| match item_node {
-                    ast::ImplItem::FnDef(it) => Function { id: ctx.to_def(&it) }.into(),
-                    ast::ImplItem::ConstDef(it) => Const { id: ctx.to_def(&it) }.into(),
-                    ast::ImplItem::TypeAliasDef(it) => TypeAlias { id: ctx.to_def(&it) }.into(),
-                })
-                .collect()
-        } else {
-            Vec::new()
-        };
-        Arc::new(TraitData { name, items, auto })
-    }
-
-    pub(crate) fn name(&self) -> &Option<Name> {
-        &self.name
-    }
-
-    pub(crate) fn items(&self) -> &[AssocItem] {
-        &self.items
-    }
-
-    pub(crate) fn is_auto(&self) -> bool {
-        self.auto
-    }
-}
-
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct TraitItemsIndex {
-    traits_by_def: FxHashMap<AssocItem, Trait>,
-}
-
-impl TraitItemsIndex {
-    pub(crate) fn trait_items_index(db: &impl DefDatabase, module: Module) -> TraitItemsIndex {
-        let mut index = TraitItemsIndex { traits_by_def: FxHashMap::default() };
-        for decl in module.declarations(db) {
-            if let crate::ModuleDef::Trait(tr) = decl {
-                for item in tr.trait_data(db).items() {
-                    index.traits_by_def.insert(*item, tr);
-                }
-            }
-        }
-        index
-    }
-
-    pub(crate) fn get_parent_trait(&self, item: AssocItem) -> Option<Trait> {
-        self.traits_by_def.get(&item).cloned()
-    }
-}
diff --git a/crates/ra_hir/src/ty/method_resolution.rs b/crates/ra_hir/src/ty/method_resolution.rs
index d20aeaacf..f377fca48 100644
--- a/crates/ra_hir/src/ty/method_resolution.rs
+++ b/crates/ra_hir/src/ty/method_resolution.rs
@@ -232,8 +232,8 @@ fn iterate_trait_method_candidates<T>(
         // trait, but if we find out it doesn't, we'll skip the rest of the
         // iteration
         let mut known_implemented = false;
-        for &item in data.items() {
-            if !is_valid_candidate(db, name, mode, item) {
+        for &item in data.items.iter() {
+            if !is_valid_candidate(db, name, mode, item.into()) {
                 continue;
             }
             if !known_implemented {
@@ -243,7 +243,7 @@ fn iterate_trait_method_candidates<T>(
                 }
             }
             known_implemented = true;
-            if let Some(result) = callback(&ty.value, item) {
+            if let Some(result) = callback(&ty.value, item.into()) {
                 return Some(result);
             }
         }