Introduce ChildFromSource

3 files changed, 317 insertions, 101 deletions

diff --git a/crates/ra_hir/src/from_source.rs b/crates/ra_hir/src/from_source.rs
index 18d87f6d7..58203c721 100644
--- a/crates/ra_hir/src/from_source.rs
+++ b/crates/ra_hir/src/from_source.rs
@@ -1,17 +1,20 @@
 //! FIXME: write short doc here
+use either::Either;
 
-use hir_def::{nameres::ModuleSource, AstItemDef, LocationCtx, ModuleId};
+use hir_def::{
+    child_from_source::ChildFromSource, nameres::ModuleSource, AstItemDef, EnumVariantId,
+    LocationCtx, ModuleId, VariantId,
+};
 use hir_expand::{name::AsName, AstId, MacroDefId, MacroDefKind};
 use ra_syntax::{
     ast::{self, AstNode, NameOwner},
-    match_ast, AstPtr, SyntaxNode,
+    match_ast, SyntaxNode,
 };
 
 use crate::{
     db::{AstDatabase, DefDatabase, HirDatabase},
-    AssocItem, Const, DefWithBody, Enum, EnumVariant, FieldSource, Function, HasSource, ImplBlock,
-    InFile, Local, MacroDef, Module, ModuleDef, Static, Struct, StructField, Trait, TypeAlias,
-    Union, VariantDef,
+    Const, DefWithBody, Enum, EnumVariant, FieldSource, Function, ImplBlock, InFile, Local,
+    MacroDef, Module, Static, Struct, StructField, Trait, TypeAlias, Union,
 };
 
 pub trait FromSource: Sized {
@@ -50,98 +53,45 @@ impl FromSource for Trait {
 impl FromSource for Function {
     type Ast = ast::FnDef;
     fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
-        let items = match Container::find(db, src.as_ref().map(|it| it.syntax()))? {
-            Container::Trait(it) => it.items(db),
-            Container::ImplBlock(it) => it.items(db),
-            Container::Module(m) => {
-                return m
-                    .declarations(db)
-                    .into_iter()
-                    .filter_map(|it| match it {
-                        ModuleDef::Function(it) => Some(it),
-                        _ => None,
-                    })
-                    .find(|it| same_source(&it.source(db), &src))
-            }
-        };
-        items
-            .into_iter()
-            .filter_map(|it| match it {
-                AssocItem::Function(it) => Some(it),
-                _ => None,
-            })
-            .find(|it| same_source(&it.source(db), &src))
+        match Container::find(db, src.as_ref().map(|it| it.syntax()))? {
+            Container::Trait(it) => it.id.child_from_source(db, src),
+            Container::ImplBlock(it) => it.id.child_from_source(db, src),
+            Container::Module(it) => it.id.child_from_source(db, src),
+        }
+        .map(Function::from)
     }
 }
 
 impl FromSource for Const {
     type Ast = ast::ConstDef;
     fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
-        let items = match Container::find(db, src.as_ref().map(|it| it.syntax()))? {
-            Container::Trait(it) => it.items(db),
-            Container::ImplBlock(it) => it.items(db),
-            Container::Module(m) => {
-                return m
-                    .declarations(db)
-                    .into_iter()
-                    .filter_map(|it| match it {
-                        ModuleDef::Const(it) => Some(it),
-                        _ => None,
-                    })
-                    .find(|it| same_source(&it.source(db), &src))
-            }
-        };
-        items
-            .into_iter()
-            .filter_map(|it| match it {
-                AssocItem::Const(it) => Some(it),
-                _ => None,
-            })
-            .find(|it| same_source(&it.source(db), &src))
+        match Container::find(db, src.as_ref().map(|it| it.syntax()))? {
+            Container::Trait(it) => it.id.child_from_source(db, src),
+            Container::ImplBlock(it) => it.id.child_from_source(db, src),
+            Container::Module(it) => it.id.child_from_source(db, src),
+        }
+        .map(Const::from)
     }
 }
 impl FromSource for Static {
     type Ast = ast::StaticDef;
     fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
-        let module = match Container::find(db, src.as_ref().map(|it| it.syntax()))? {
-            Container::Module(it) => it,
-            Container::Trait(_) | Container::ImplBlock(_) => return None,
-        };
-        module
-            .declarations(db)
-            .into_iter()
-            .filter_map(|it| match it {
-                ModuleDef::Static(it) => Some(it),
-                _ => None,
-            })
-            .find(|it| same_source(&it.source(db), &src))
+        match Container::find(db, src.as_ref().map(|it| it.syntax()))? {
+            Container::Module(it) => it.id.child_from_source(db, src).map(Static::from),
+            Container::Trait(_) | Container::ImplBlock(_) => None,
+        }
     }
 }
 
 impl FromSource for TypeAlias {
     type Ast = ast::TypeAliasDef;
     fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
-        let items = match Container::find(db, src.as_ref().map(|it| it.syntax()))? {
-            Container::Trait(it) => it.items(db),
-            Container::ImplBlock(it) => it.items(db),
-            Container::Module(m) => {
-                return m
-                    .declarations(db)
-                    .into_iter()
-                    .filter_map(|it| match it {
-                        ModuleDef::TypeAlias(it) => Some(it),
-                        _ => None,
-                    })
-                    .find(|it| same_source(&it.source(db), &src))
-            }
-        };
-        items
-            .into_iter()
-            .filter_map(|it| match it {
-                AssocItem::TypeAlias(it) => Some(it),
-                _ => None,
-            })
-            .find(|it| same_source(&it.source(db), &src))
+        match Container::find(db, src.as_ref().map(|it| it.syntax()))? {
+            Container::Trait(it) => it.id.child_from_source(db, src),
+            Container::ImplBlock(it) => it.id.child_from_source(db, src),
+            Container::Module(it) => it.id.child_from_source(db, src),
+        }
+        .map(TypeAlias::from)
     }
 }
 
@@ -174,34 +124,33 @@ impl FromSource for EnumVariant {
     fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
         let parent_enum = src.value.parent_enum();
         let src_enum = InFile { file_id: src.file_id, value: parent_enum };
-        let variants = Enum::from_source(db, src_enum)?.variants(db);
-        variants.into_iter().find(|v| same_source(&v.source(db), &src))
+        let parent_enum = Enum::from_source(db, src_enum)?;
+        parent_enum.id.child_from_source(db, src).map(EnumVariant::from)
     }
 }
 
 impl FromSource for StructField {
     type Ast = FieldSource;
     fn from_source(db: &(impl DefDatabase + AstDatabase), src: InFile<Self::Ast>) -> Option<Self> {
-        let variant_def: VariantDef = match src.value {
+        let variant_id: VariantId = match src.value {
             FieldSource::Named(ref field) => {
                 let value = field.syntax().ancestors().find_map(ast::StructDef::cast)?;
                 let src = InFile { file_id: src.file_id, value };
                 let def = Struct::from_source(db, src)?;
-                VariantDef::from(def)
+                def.id.into()
             }
             FieldSource::Pos(ref field) => {
                 let value = field.syntax().ancestors().find_map(ast::EnumVariant::cast)?;
                 let src = InFile { file_id: src.file_id, value };
                 let def = EnumVariant::from_source(db, src)?;
-                VariantDef::from(def)
+                EnumVariantId::from(def).into()
             }
         };
-        variant_def
-            .variant_data(db)
-            .fields()
-            .iter()
-            .map(|(id, _)| StructField { parent: variant_def, id })
-            .find(|f| f.source(db) == src)
+        let src = src.map(|field_source| match field_source {
+            FieldSource::Pos(it) => Either::Left(it),
+            FieldSource::Named(it) => Either::Right(it),
+        });
+        variant_id.child_from_source(db, src).map(StructField::from)
     }
 }
 
@@ -314,13 +263,3 @@ impl Container {
         Some(Container::Module(c))
     }
 }
-
-/// 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 guarantee, but, for
-/// the time being, we will use identity-less AstPtr comparison.
-fn same_source<N: AstNode>(s1: &InFile<N>, s2: &InFile<N>) -> bool {
-    s1.as_ref().map(AstPtr::new) == s2.as_ref().map(AstPtr::new)
-}
diff --git a/crates/ra_hir_def/src/child_from_source.rs b/crates/ra_hir_def/src/child_from_source.rs
new file mode 100644
index 000000000..37d4b7870
--- /dev/null
+++ b/crates/ra_hir_def/src/child_from_source.rs
@@ -0,0 +1,276 @@
+//! When *constructing* `hir`, we start at some parent syntax node and recursively
+//! lower the children.
+//!
+//! This modules allows one to go in the opposite direction: start with a syntax
+//! node for a *child*, and get its hir.
+
+use either::Either;
+use hir_expand::InFile;
+use ra_syntax::{ast, AstNode, AstPtr};
+
+use crate::{
+    db::DefDatabase,
+    src::{HasChildSource, HasSource},
+    AssocItemId, ConstId, EnumId, EnumVariantId, FunctionId, ImplId, Lookup, ModuleDefId, ModuleId,
+    StaticId, StructFieldId, TraitId, TypeAliasId, VariantId,
+};
+
+pub trait ChildFromSource<CHILD, SOURCE> {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<SOURCE>,
+    ) -> Option<CHILD>;
+}
+
+impl ChildFromSource<FunctionId, ast::FnDef> for TraitId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::FnDef>,
+    ) -> Option<FunctionId> {
+        let data = db.trait_data(*self);
+        data.items
+            .iter()
+            .filter_map(|(_, item)| match item {
+                AssocItemId::FunctionId(it) => Some(*it),
+                _ => None,
+            })
+            .find(|func| {
+                let source = func.lookup(db).source(db);
+                same_source(&source, &child_source)
+            })
+    }
+}
+
+impl ChildFromSource<FunctionId, ast::FnDef> for ImplId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::FnDef>,
+    ) -> Option<FunctionId> {
+        let data = db.impl_data(*self);
+        data.items
+            .iter()
+            .filter_map(|item| match item {
+                AssocItemId::FunctionId(it) => Some(*it),
+                _ => None,
+            })
+            .find(|func| {
+                let source = func.lookup(db).source(db);
+                same_source(&source, &child_source)
+            })
+    }
+}
+
+impl ChildFromSource<FunctionId, ast::FnDef> for ModuleId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::FnDef>,
+    ) -> Option<FunctionId> {
+        let crate_def_map = db.crate_def_map(self.krate);
+        let res = crate_def_map[self.local_id]
+            .scope
+            .declarations()
+            .filter_map(|item| match item {
+                ModuleDefId::FunctionId(it) => Some(it),
+                _ => None,
+            })
+            .find(|func| {
+                let source = func.lookup(db).source(db);
+                same_source(&source, &child_source)
+            });
+        res
+    }
+}
+
+impl ChildFromSource<ConstId, ast::ConstDef> for TraitId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::ConstDef>,
+    ) -> Option<ConstId> {
+        let data = db.trait_data(*self);
+        data.items
+            .iter()
+            .filter_map(|(_, item)| match item {
+                AssocItemId::ConstId(it) => Some(*it),
+                _ => None,
+            })
+            .find(|func| {
+                let source = func.lookup(db).source(db);
+                same_source(&source, &child_source)
+            })
+    }
+}
+
+impl ChildFromSource<ConstId, ast::ConstDef> for ImplId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::ConstDef>,
+    ) -> Option<ConstId> {
+        let data = db.impl_data(*self);
+        data.items
+            .iter()
+            .filter_map(|item| match item {
+                AssocItemId::ConstId(it) => Some(*it),
+                _ => None,
+            })
+            .find(|func| {
+                let source = func.lookup(db).source(db);
+                same_source(&source, &child_source)
+            })
+    }
+}
+
+impl ChildFromSource<ConstId, ast::ConstDef> for ModuleId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::ConstDef>,
+    ) -> Option<ConstId> {
+        let crate_def_map = db.crate_def_map(self.krate);
+        let res = crate_def_map[self.local_id]
+            .scope
+            .declarations()
+            .filter_map(|item| match item {
+                ModuleDefId::ConstId(it) => Some(it),
+                _ => None,
+            })
+            .find(|func| {
+                let source = func.lookup(db).source(db);
+                same_source(&source, &child_source)
+            });
+        res
+    }
+}
+
+impl ChildFromSource<TypeAliasId, ast::TypeAliasDef> for TraitId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::TypeAliasDef>,
+    ) -> Option<TypeAliasId> {
+        let data = db.trait_data(*self);
+        data.items
+            .iter()
+            .filter_map(|(_, item)| match item {
+                AssocItemId::TypeAliasId(it) => Some(*it),
+                _ => None,
+            })
+            .find(|func| {
+                let source = func.lookup(db).source(db);
+                same_source(&source, &child_source)
+            })
+    }
+}
+
+impl ChildFromSource<TypeAliasId, ast::TypeAliasDef> for ImplId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::TypeAliasDef>,
+    ) -> Option<TypeAliasId> {
+        let data = db.impl_data(*self);
+        data.items
+            .iter()
+            .filter_map(|item| match item {
+                AssocItemId::TypeAliasId(it) => Some(*it),
+                _ => None,
+            })
+            .find(|func| {
+                let source = func.lookup(db).source(db);
+                same_source(&source, &child_source)
+            })
+    }
+}
+
+impl ChildFromSource<TypeAliasId, ast::TypeAliasDef> for ModuleId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::TypeAliasDef>,
+    ) -> Option<TypeAliasId> {
+        let crate_def_map = db.crate_def_map(self.krate);
+        let res = crate_def_map[self.local_id]
+            .scope
+            .declarations()
+            .filter_map(|item| match item {
+                ModuleDefId::TypeAliasId(it) => Some(it),
+                _ => None,
+            })
+            .find(|func| {
+                let source = func.lookup(db).source(db);
+                same_source(&source, &child_source)
+            });
+        res
+    }
+}
+
+impl ChildFromSource<StaticId, ast::StaticDef> for ModuleId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::StaticDef>,
+    ) -> Option<StaticId> {
+        let crate_def_map = db.crate_def_map(self.krate);
+        let res = crate_def_map[self.local_id]
+            .scope
+            .declarations()
+            .filter_map(|item| match item {
+                ModuleDefId::StaticId(it) => Some(it),
+                _ => None,
+            })
+            .find(|func| {
+                let source = func.lookup(db).source(db);
+                same_source(&source, &child_source)
+            });
+        res
+    }
+}
+
+impl ChildFromSource<StructFieldId, Either<ast::TupleFieldDef, ast::RecordFieldDef>> for VariantId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<Either<ast::TupleFieldDef, ast::RecordFieldDef>>,
+    ) -> Option<StructFieldId> {
+        let arena_map = self.child_source(db);
+        let (local_id, _) = arena_map.as_ref().value.iter().find(|(_local_id, source)| {
+            child_source.file_id == arena_map.file_id
+                && match (source, &child_source.value) {
+                    (Either::Left(a), Either::Left(b)) => AstPtr::new(a) == AstPtr::new(b),
+                    (Either::Right(a), Either::Right(b)) => AstPtr::new(a) == AstPtr::new(b),
+                    _ => false,
+                }
+        })?;
+        Some(StructFieldId { parent: *self, local_id })
+    }
+}
+
+impl ChildFromSource<EnumVariantId, ast::EnumVariant> for EnumId {
+    fn child_from_source(
+        &self,
+        db: &impl DefDatabase,
+        child_source: InFile<ast::EnumVariant>,
+    ) -> Option<EnumVariantId> {
+        let arena_map = self.child_source(db);
+        let (local_id, _) = arena_map.as_ref().value.iter().find(|(_local_id, source)| {
+            child_source.file_id == arena_map.file_id
+                && AstPtr::new(*source) == AstPtr::new(&child_source.value)
+        })?;
+        Some(EnumVariantId { parent: *self, local_id })
+    }
+}
+
+/// 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 guarantee, but, for
+/// the time being, we will use identity-less AstPtr comparison.
+fn same_source<N: AstNode>(s1: &InFile<N>, s2: &InFile<N>) -> bool {
+    s1.as_ref().map(AstPtr::new) == s2.as_ref().map(AstPtr::new)
+}
diff --git a/crates/ra_hir_def/src/lib.rs b/crates/ra_hir_def/src/lib.rs
index cfeacfded..e02622f62 100644
--- a/crates/ra_hir_def/src/lib.rs
+++ b/crates/ra_hir_def/src/lib.rs
@@ -30,6 +30,7 @@ mod trace;
 pub mod nameres;
 
 pub mod src;
+pub mod child_from_source;
 
 #[cfg(test)]
 mod test_db;