Implement ChildBySource for DefWithBody

1 files changed, 55 insertions, 42 deletions

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
+    }
+}