5 files changed, 71 insertions, 55 deletions

diff --git a/crates/ra_hir_def/src/body.rs b/crates/ra_hir_def/src/body.rs
index d4cab0561..332c509e1 100644
--- a/crates/ra_hir_def/src/body.rs
+++ b/crates/ra_hir_def/src/body.rs
@@ -17,7 +17,7 @@ use crate::{
     nameres::{BuiltinShadowMode, CrateDefMap},
     path::{ModPath, Path},
     src::HasSource,
-    DefWithBodyId, HasModule, Lookup, ModuleId,
+    DefWithBodyId, HasModule, Lookup, ModuleDefId, ModuleId,
 };
 
 struct Expander {
@@ -119,6 +119,7 @@ pub struct Body {
     pub params: Vec<PatId>,
     /// The `ExprId` of the actual body expression.
     pub body_expr: ExprId,
+    pub defs: Vec<ModuleDefId>,
 }
 
 pub type ExprPtr = Either<AstPtr<ast::Expr>, AstPtr<ast::RecordField>>;
diff --git a/crates/ra_hir_def/src/body/lower.rs b/crates/ra_hir_def/src/body/lower.rs
index 61193b4d8..86960186f 100644
--- a/crates/ra_hir_def/src/body/lower.rs
+++ b/crates/ra_hir_def/src/body/lower.rs
@@ -41,6 +41,7 @@ pub(super) fn lower(
             pats: Arena::default(),
             params: Vec::new(),
             body_expr: ExprId::dummy(),
+            defs: Vec::new(),
         },
     }
     .collect(params, body)
diff --git a/crates/ra_hir_def/src/child_by_source.rs b/crates/ra_hir_def/src/child_by_source.rs
index 3c9379b15..f5a65ad40 100644
--- a/crates/ra_hir_def/src/child_by_source.rs
+++ b/crates/ra_hir_def/src/child_by_source.rs
@@ -11,8 +11,8 @@ use crate::{
     dyn_map::DynMap,
     keys,
     src::{HasChildSource, HasSource},
-    AdtId, AssocItemId, EnumId, EnumVariantId, ImplId, Lookup, ModuleDefId, ModuleId,
-    StructFieldId, TraitId, VariantId,
+    AdtId, AssocItemId, DefWithBodyId, EnumId, EnumVariantId, ImplId, Lookup, ModuleDefId,
+    ModuleId, StructFieldId, TraitId, VariantId,
 };
 
 pub trait ChildBySource {
@@ -76,47 +76,11 @@ impl ChildBySource for ModuleId {
         let mut res = DynMap::default();
 
         let crate_def_map = db.crate_def_map(self.krate);
-        for item in crate_def_map[self.local_id].scope.declarations() {
-            match item {
-                ModuleDefId::FunctionId(func) => {
-                    let src = func.lookup(db).source(db);
-                    res[keys::FUNCTION].insert(src, func)
-                }
-                ModuleDefId::ConstId(konst) => {
-                    let src = konst.lookup(db).source(db);
-                    res[keys::CONST].insert(src, konst)
-                }
-                ModuleDefId::StaticId(statik) => {
-                    let src = statik.lookup(db).source(db);
-                    res[keys::STATIC].insert(src, statik)
-                }
-                ModuleDefId::TypeAliasId(ty) => {
-                    let src = ty.lookup(db).source(db);
-                    res[keys::TYPE_ALIAS].insert(src, ty)
-                }
-                ModuleDefId::TraitId(trait_) => {
-                    let src = trait_.lookup(db).source(db);
-                    res[keys::TRAIT].insert(src, trait_)
-                }
-                ModuleDefId::AdtId(adt) => match adt {
-                    AdtId::StructId(strukt) => {
-                        let src = strukt.lookup(db).source(db);
-                        res[keys::STRUCT].insert(src, strukt)
-                    }
-                    AdtId::UnionId(union_) => {
-                        let src = union_.lookup(db).source(db);
-                        res[keys::UNION].insert(src, union_)
-                    }
-                    AdtId::EnumId(enum_) => {
-                        let src = enum_.lookup(db).source(db);
-                        res[keys::ENUM].insert(src, enum_)
-                    }
-                },
-                _ => (),
-            }
-        }
+        let module_data = &crate_def_map[self.local_id];
 
-        for &impl_ in crate_def_map[self.local_id].impls.iter() {
+        module_data.scope.declarations().for_each(|item| add_module_def(db, &mut res, item));
+
+        for &impl_ in module_data.impls.iter() {
             let src = impl_.lookup(db).source(db);
             res[keys::IMPL].insert(src, impl_)
         }
@@ -125,6 +89,46 @@ impl ChildBySource for ModuleId {
     }
 }
 
+fn add_module_def(db: &impl DefDatabase, map: &mut DynMap, item: ModuleDefId) {
+    match item {
+        ModuleDefId::FunctionId(func) => {
+            let src = func.lookup(db).source(db);
+            map[keys::FUNCTION].insert(src, func)
+        }
+        ModuleDefId::ConstId(konst) => {
+            let src = konst.lookup(db).source(db);
+            map[keys::CONST].insert(src, konst)
+        }
+        ModuleDefId::StaticId(statik) => {
+            let src = statik.lookup(db).source(db);
+            map[keys::STATIC].insert(src, statik)
+        }
+        ModuleDefId::TypeAliasId(ty) => {
+            let src = ty.lookup(db).source(db);
+            map[keys::TYPE_ALIAS].insert(src, ty)
+        }
+        ModuleDefId::TraitId(trait_) => {
+            let src = trait_.lookup(db).source(db);
+            map[keys::TRAIT].insert(src, trait_)
+        }
+        ModuleDefId::AdtId(adt) => match adt {
+            AdtId::StructId(strukt) => {
+                let src = strukt.lookup(db).source(db);
+                map[keys::STRUCT].insert(src, strukt)
+            }
+            AdtId::UnionId(union_) => {
+                let src = union_.lookup(db).source(db);
+                map[keys::UNION].insert(src, union_)
+            }
+            AdtId::EnumId(enum_) => {
+                let src = enum_.lookup(db).source(db);
+                map[keys::ENUM].insert(src, enum_)
+            }
+        },
+        _ => (),
+    }
+}
+
 impl ChildBySource for VariantId {
     fn child_by_source(&self, db: &impl DefDatabase) -> DynMap {
         let mut res = DynMap::default();
@@ -160,3 +164,12 @@ impl ChildBySource for EnumId {
         res
     }
 }
+
+impl ChildBySource for DefWithBodyId {
+    fn child_by_source(&self, db: &impl DefDatabase) -> DynMap {
+        let mut res = DynMap::default();
+        let body = db.body(*self);
+        body.defs.iter().copied().for_each(|item| add_module_def(db, &mut res, item));
+        res
+    }
+}
diff --git a/crates/ra_hir_def/src/lib.rs b/crates/ra_hir_def/src/lib.rs
index f085bbe87..4fc3127c4 100644
--- a/crates/ra_hir_def/src/lib.rs
+++ b/crates/ra_hir_def/src/lib.rs
@@ -335,6 +335,7 @@ pub enum ContainerId {
     ModuleId(ModuleId),
     ImplId(ImplId),
     TraitId(TraitId),
+    DefWithBodyId(DefWithBodyId),
 }
 
 /// A Data Type
@@ -478,33 +479,32 @@ pub trait HasModule {
     fn module(&self, db: &impl db::DefDatabase) -> ModuleId;
 }
 
-impl HasModule for FunctionLoc {
+impl HasModule for ContainerId {
     fn module(&self, db: &impl db::DefDatabase) -> ModuleId {
-        match self.container {
+        match *self {
             ContainerId::ModuleId(it) => it,
             ContainerId::ImplId(it) => it.lookup(db).container,
             ContainerId::TraitId(it) => it.lookup(db).container,
+            ContainerId::DefWithBodyId(it) => it.module(db),
         }
     }
 }
 
+impl HasModule for FunctionLoc {
+    fn module(&self, db: &impl db::DefDatabase) -> ModuleId {
+        self.container.module(db)
+    }
+}
+
 impl HasModule for TypeAliasLoc {
     fn module(&self, db: &impl db::DefDatabase) -> ModuleId {
-        match self.container {
-            ContainerId::ModuleId(it) => it,
-            ContainerId::ImplId(it) => it.lookup(db).container,
-            ContainerId::TraitId(it) => it.lookup(db).container,
-        }
+        self.container.module(db)
     }
 }
 
 impl HasModule for ConstLoc {
     fn module(&self, db: &impl db::DefDatabase) -> ModuleId {
-        match self.container {
-            ContainerId::ModuleId(it) => it,
-            ContainerId::ImplId(it) => it.lookup(db).container,
-            ContainerId::TraitId(it) => it.lookup(db).container,
-        }
+        self.container.module(db)
     }
 }
 
diff --git a/crates/ra_hir_def/src/resolver.rs b/crates/ra_hir_def/src/resolver.rs
index 2694c0438..250329271 100644
--- a/crates/ra_hir_def/src/resolver.rs
+++ b/crates/ra_hir_def/src/resolver.rs
@@ -586,6 +586,7 @@ impl HasResolver for ContainerId {
             ContainerId::TraitId(it) => it.resolver(db),
             ContainerId::ImplId(it) => it.resolver(db),
             ContainerId::ModuleId(it) => it.resolver(db),
+            ContainerId::DefWithBodyId(it) => it.resolver(db),
         }
     }
 }