1 files changed, 150 insertions, 212 deletions

diff --git a/crates/ra_hir/src/from_source.rs b/crates/ra_hir/src/from_source.rs
index 9f7c22b21..6314be8d4 100644
--- a/crates/ra_hir/src/from_source.rs
+++ b/crates/ra_hir/src/from_source.rs
@@ -1,219 +1,140 @@
-//! FIXME: write short doc here
+//! Finds a corresponding hir data structure for a syntax node in a specific
+//! file.
 
-use hir_def::{AstItemDef, LocationCtx, ModuleId};
+use hir_def::{
+    child_by_source::ChildBySource, dyn_map::DynMap, keys, keys::Key, nameres::ModuleSource,
+    ConstId, DefWithBodyId, EnumId, EnumVariantId, FunctionId, GenericDefId, ImplId, ModuleId,
+    StaticId, StructId, TraitId, TypeAliasId, UnionId, VariantId,
+};
 use hir_expand::{name::AsName, AstId, MacroDefId, MacroDefKind};
+use ra_db::FileId;
+use ra_prof::profile;
 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,
-    Local, MacroDef, Module, ModuleDef, ModuleSource, Source, Static, Struct, StructField, Trait,
-    TypeAlias, Union, VariantDef,
+    db::{DefDatabase, HirDatabase},
+    Const, DefWithBody, Enum, EnumVariant, FieldSource, Function, ImplBlock, InFile, Local,
+    MacroDef, Module, Static, Struct, StructField, Trait, TypeAlias, TypeParam, Union,
 };
 
 pub trait FromSource: Sized {
     type Ast;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self>;
+    fn from_source(db: &impl DefDatabase, src: InFile<Self::Ast>) -> Option<Self>;
 }
 
-impl FromSource for Struct {
-    type Ast = ast::StructDef;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self> {
-        let id = from_source(db, src)?;
-        Some(Struct { id })
-    }
-}
-impl FromSource for Union {
-    type Ast = ast::UnionDef;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self> {
-        let id = from_source(db, src)?;
-        Some(Union { id })
-    }
-}
-impl FromSource for Enum {
-    type Ast = ast::EnumDef;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self> {
-        let id = from_source(db, src)?;
-        Some(Enum { id })
-    }
-}
-impl FromSource for Trait {
-    type Ast = ast::TraitDef;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self> {
-        let id = from_source(db, src)?;
-        Some(Trait { id })
-    }
-}
-impl FromSource for Function {
-    type Ast = ast::FnDef;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<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))
-    }
+pub trait FromSourceByContainer: Sized {
+    type Ast: AstNode + 'static;
+    type Id: Copy + 'static;
+    const KEY: Key<Self::Ast, Self::Id>;
 }
 
-impl FromSource for Const {
-    type Ast = ast::ConstDef;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<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))
-    }
-}
-impl FromSource for Static {
-    type Ast = ast::StaticDef;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<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))
+impl<T: FromSourceByContainer> FromSource for T
+where
+    T: From<<T as FromSourceByContainer>::Id>,
+{
+    type Ast = <T as FromSourceByContainer>::Ast;
+    fn from_source(db: &impl DefDatabase, src: InFile<Self::Ast>) -> Option<Self> {
+        analyze_container(db, src.as_ref().map(|it| it.syntax()))[T::KEY]
+            .get(&src)
+            .copied()
+            .map(Self::from)
     }
 }
 
-impl FromSource for TypeAlias {
-    type Ast = ast::TypeAliasDef;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<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))
-    }
+macro_rules! from_source_by_container_impls {
+    ($(($hir:ident, $id:ident, $ast:path, $key:path)),* ,) => {$(
+        impl FromSourceByContainer for $hir {
+            type Ast = $ast;
+            type Id = $id;
+            const KEY: Key<Self::Ast, Self::Id> = $key;
+        }
+    )*}
 }
 
+from_source_by_container_impls![
+    (Struct, StructId, ast::StructDef, keys::STRUCT),
+    (Union, UnionId, ast::UnionDef, keys::UNION),
+    (Enum, EnumId, ast::EnumDef, keys::ENUM),
+    (Trait, TraitId, ast::TraitDef, keys::TRAIT),
+    (Function, FunctionId, ast::FnDef, keys::FUNCTION),
+    (Static, StaticId, ast::StaticDef, keys::STATIC),
+    (Const, ConstId, ast::ConstDef, keys::CONST),
+    (TypeAlias, TypeAliasId, ast::TypeAliasDef, keys::TYPE_ALIAS),
+    (ImplBlock, ImplId, ast::ImplBlock, keys::IMPL),
+];
+
 impl FromSource for MacroDef {
     type Ast = ast::MacroCall;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self> {
+    fn from_source(db: &impl DefDatabase, src: InFile<Self::Ast>) -> Option<Self> {
         let kind = MacroDefKind::Declarative;
 
         let module_src = ModuleSource::from_child_node(db, src.as_ref().map(|it| it.syntax()));
-        let module = Module::from_definition(db, Source::new(src.file_id, module_src))?;
-        let krate = module.krate().crate_id();
+        let module = Module::from_definition(db, InFile::new(src.file_id, module_src))?;
+        let krate = Some(module.krate().id);
 
-        let ast_id = AstId::new(src.file_id, db.ast_id_map(src.file_id).ast_id(&src.value));
+        let ast_id = Some(AstId::new(src.file_id, db.ast_id_map(src.file_id).ast_id(&src.value)));
 
         let id: MacroDefId = MacroDefId { krate, ast_id, kind };
         Some(MacroDef { id })
     }
 }
 
-impl FromSource for ImplBlock {
-    type Ast = ast::ImplBlock;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self> {
-        let id = from_source(db, src)?;
-        Some(ImplBlock { id })
-    }
-}
-
 impl FromSource for EnumVariant {
     type Ast = ast::EnumVariant;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self> {
+    fn from_source(db: &impl DefDatabase, src: InFile<Self::Ast>) -> Option<Self> {
         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| same_source(&v.source(db), &src))
+        let src_enum = InFile { file_id: src.file_id, value: parent_enum };
+        let parent_enum = Enum::from_source(db, src_enum)?;
+        parent_enum.id.child_by_source(db)[keys::ENUM_VARIANT]
+            .get(&src)
+            .copied()
+            .map(EnumVariant::from)
     }
 }
 
 impl FromSource for StructField {
     type Ast = FieldSource;
-    fn from_source(db: &(impl DefDatabase + AstDatabase), src: Source<Self::Ast>) -> Option<Self> {
-        let variant_def: VariantDef = match src.value {
-            FieldSource::Named(ref field) => {
+    fn from_source(db: &impl DefDatabase, src: InFile<Self::Ast>) -> Option<Self> {
+        let src = src.as_ref();
+
+        // FIXME this is buggy
+        let variant_id: VariantId = match src.value {
+            FieldSource::Named(field) => {
                 let value = field.syntax().ancestors().find_map(ast::StructDef::cast)?;
-                let src = Source { file_id: src.file_id, value };
+                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) => {
+            FieldSource::Pos(field) => {
                 let value = field.syntax().ancestors().find_map(ast::EnumVariant::cast)?;
-                let src = Source { file_id: src.file_id, value };
+                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 dyn_map = variant_id.child_by_source(db);
+        match src.value {
+            FieldSource::Pos(it) => dyn_map[keys::TUPLE_FIELD].get(&src.with_value(it.clone())),
+            FieldSource::Named(it) => dyn_map[keys::RECORD_FIELD].get(&src.with_value(it.clone())),
+        }
+        .copied()
+        .map(StructField::from)
     }
 }
 
 impl Local {
-    pub fn from_source(db: &impl HirDatabase, src: Source<ast::BindPat>) -> Option<Self> {
+    pub fn from_source(db: &impl HirDatabase, src: InFile<ast::BindPat>) -> Option<Self> {
         let file_id = src.file_id;
         let parent: DefWithBody = src.value.syntax().ancestors().find_map(|it| {
             let res = match_ast! {
                 match it {
-                    ast::ConstDef(value) => { Const::from_source(db, Source { value, file_id})?.into() },
-                    ast::StaticDef(value) => { Static::from_source(db, Source { value, file_id})?.into() },
-                    ast::FnDef(value) => { Function::from_source(db, Source { value, file_id})?.into() },
+                    ast::ConstDef(value) => { Const::from_source(db, InFile { value, file_id})?.into() },
+                    ast::StaticDef(value) => { Static::from_source(db, InFile { value, file_id})?.into() },
+                    ast::FnDef(value) => { Function::from_source(db, InFile { value, file_id})?.into() },
                     _ => return None,
                 }
             };
@@ -226,84 +147,111 @@ impl Local {
     }
 }
 
+impl TypeParam {
+    pub fn from_source(db: &impl HirDatabase, src: InFile<ast::TypeParam>) -> Option<Self> {
+        let file_id = src.file_id;
+        let parent: GenericDefId = src.value.syntax().ancestors().find_map(|it| {
+            let res = match_ast! {
+                match it {
+                    ast::FnDef(value) => { Function::from_source(db, InFile { value, file_id})?.id.into() },
+                    ast::StructDef(value) => { Struct::from_source(db, InFile { value, file_id})?.id.into() },
+                    ast::EnumDef(value) => { Enum::from_source(db, InFile { value, file_id})?.id.into() },
+                    ast::TraitDef(value) => { Trait::from_source(db, InFile { value, file_id})?.id.into() },
+                    ast::TypeAliasDef(value) => { TypeAlias::from_source(db, InFile { value, file_id})?.id.into() },
+                    ast::ImplBlock(value) => { ImplBlock::from_source(db, InFile { value, file_id})?.id.into() },
+                    _ => return None,
+                }
+            };
+            Some(res)
+        })?;
+        let &id = parent.child_by_source(db)[keys::TYPE_PARAM].get(&src)?;
+        Some(TypeParam { id })
+    }
+}
+
 impl Module {
-    pub fn from_declaration(db: &impl DefDatabase, src: Source<ast::Module>) -> Option<Self> {
+    pub fn from_declaration(db: &impl DefDatabase, src: InFile<ast::Module>) -> Option<Self> {
+        let _p = profile("Module::from_declaration");
         let parent_declaration = src.value.syntax().ancestors().skip(1).find_map(ast::Module::cast);
 
         let parent_module = match parent_declaration {
             Some(parent_declaration) => {
-                let src_parent = Source { file_id: src.file_id, value: parent_declaration };
+                let src_parent = InFile { file_id: src.file_id, value: parent_declaration };
                 Module::from_declaration(db, src_parent)
             }
-            _ => {
-                let src_parent = Source {
-                    file_id: src.file_id,
-                    value: ModuleSource::new(db, Some(src.file_id.original_file(db)), None),
-                };
+            None => {
+                let source_file = db.parse(src.file_id.original_file(db)).tree();
+                let src_parent =
+                    InFile { file_id: src.file_id, value: ModuleSource::SourceFile(source_file) };
                 Module::from_definition(db, src_parent)
             }
         }?;
 
-        let child_name = src.value.name()?;
-        parent_module.child(db, &child_name.as_name())
+        let child_name = src.value.name()?.as_name();
+        let def_map = db.crate_def_map(parent_module.id.krate);
+        let child_id = def_map[parent_module.id.local_id].children.get(&child_name)?;
+        Some(parent_module.with_module_id(*child_id))
     }
 
-    pub fn from_definition(db: &impl DefDatabase, src: Source<ModuleSource>) -> Option<Self> {
+    pub fn from_definition(db: &impl DefDatabase, src: InFile<ModuleSource>) -> Option<Self> {
+        let _p = profile("Module::from_definition");
         match src.value {
             ModuleSource::Module(ref module) => {
                 assert!(!module.has_semi());
                 return Module::from_declaration(
                     db,
-                    Source { file_id: src.file_id, value: module.clone() },
+                    InFile { file_id: src.file_id, value: module.clone() },
                 );
             }
             ModuleSource::SourceFile(_) => (),
         };
 
         let original_file = src.file_id.original_file(db);
+        Module::from_file(db, original_file)
+    }
 
-        let (krate, local_id) = db.relevant_crates(original_file).iter().find_map(|&crate_id| {
+    fn from_file(db: &impl DefDatabase, file: FileId) -> Option<Self> {
+        let _p = profile("Module::from_file");
+        let (krate, local_id) = db.relevant_crates(file).iter().find_map(|&crate_id| {
             let crate_def_map = db.crate_def_map(crate_id);
-            let local_id = crate_def_map.modules_for_file(original_file).next()?;
+            let local_id = crate_def_map.modules_for_file(file).next()?;
             Some((crate_id, local_id))
         })?;
         Some(Module { id: ModuleId { krate, local_id } })
     }
 }
 
-fn from_source<N, DEF>(db: &(impl DefDatabase + AstDatabase), src: Source<N>) -> Option<DEF>
-where
-    N: AstNode,
-    DEF: AstItemDef<N>,
-{
-    let module_src = ModuleSource::from_child_node(db, src.as_ref().map(|it| it.syntax()));
-    let module = Module::from_definition(db, Source::new(src.file_id, module_src))?;
-    let ctx = LocationCtx::new(db, module.id, src.file_id);
-    let items = db.ast_id_map(src.file_id);
-    let item_id = items.ast_id(&src.value);
-    Some(DEF::from_ast_id(ctx, item_id))
-}
-
-enum Container {
-    Trait(Trait),
-    ImplBlock(ImplBlock),
-    Module(Module),
-}
+fn analyze_container(db: &impl DefDatabase, src: InFile<&SyntaxNode>) -> DynMap {
+    let _p = profile("analyze_container");
+    return child_by_source(db, src).unwrap_or_default();
 
-impl Container {
-    fn find(db: &impl DefDatabase, src: Source<&SyntaxNode>) -> Option<Container> {
-        // FIXME: this doesn't try to handle nested declarations
-        for container in src.value.ancestors() {
+    fn child_by_source(db: &impl DefDatabase, src: InFile<&SyntaxNode>) -> Option<DynMap> {
+        for container in src.value.ancestors().skip(1) {
             let res = match_ast! {
                 match container {
                     ast::TraitDef(it) => {
-                        let c = Trait::from_source(db, src.with_value(it))?;
-                        Container::Trait(c)
+                        let def = Trait::from_source(db, src.with_value(it))?;
+                        def.id.child_by_source(db)
                     },
                     ast::ImplBlock(it) => {
-                        let c = ImplBlock::from_source(db, src.with_value(it))?;
-                        Container::ImplBlock(c)
-                     },
+                        let def = ImplBlock::from_source(db, src.with_value(it))?;
+                        def.id.child_by_source(db)
+                    },
+                    ast::FnDef(it) => {
+                        let def = Function::from_source(db, src.with_value(it))?;
+                        DefWithBodyId::from(def.id)
+                            .child_by_source(db)
+                    },
+                    ast::StaticDef(it) => {
+                        let def = Static::from_source(db, src.with_value(it))?;
+                        DefWithBodyId::from(def.id)
+                            .child_by_source(db)
+                    },
+                    ast::ConstDef(it) => {
+                        let def = Const::from_source(db, src.with_value(it))?;
+                        DefWithBodyId::from(def.id)
+                            .child_by_source(db)
+                    },
                     _ => { continue },
                 }
             };
@@ -312,16 +260,6 @@ impl Container {
 
         let module_source = ModuleSource::from_child_node(db, src);
         let c = Module::from_definition(db, src.with_value(module_source))?;
-        Some(Container::Module(c))
+        Some(c.id.child_by_source(db))
     }
 }
-
-/// 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: &Source<N>, s2: &Source<N>) -> bool {
-    s1.as_ref().map(AstPtr::new) == s2.as_ref().map(AstPtr::new)
-}