2 files changed, 545 insertions, 45 deletions

diff --git a/crates/ra_assists/src/handlers/auto_import.rs b/crates/ra_assists/src/handlers/auto_import.rs
index 1fb701da5..c4aea2a06 100644
--- a/crates/ra_assists/src/handlers/auto_import.rs
+++ b/crates/ra_assists/src/handlers/auto_import.rs
@@ -1,10 +1,18 @@
-use ra_ide_db::imports_locator::ImportsLocator;
-use ra_syntax::ast::{self, AstNode};
-
 use crate::{
     assist_ctx::{Assist, AssistCtx},
     insert_use_statement, AssistId,
 };
+use hir::{
+    db::HirDatabase, AsAssocItem, AssocItemContainer, ModPath, Module, ModuleDef, PathResolution,
+    SourceAnalyzer, Trait, Type,
+};
+use ra_ide_db::{imports_locator::ImportsLocator, RootDatabase};
+use ra_prof::profile;
+use ra_syntax::{
+    ast::{self, AstNode},
+    SyntaxNode,
+};
+use rustc_hash::FxHashSet;
 use std::collections::BTreeSet;
 
 // Assist: auto_import
@@ -27,52 +35,24 @@ use std::collections::BTreeSet;
 // # pub mod std { pub mod collections { pub struct HashMap { } } }
 // ```
 pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
-    let path_under_caret: ast::Path = ctx.find_node_at_offset()?;
-    if path_under_caret.syntax().ancestors().find_map(ast::UseItem::cast).is_some() {
-        return None;
-    }
-
-    let module = path_under_caret.syntax().ancestors().find_map(ast::Module::cast);
-    let position = match module.and_then(|it| it.item_list()) {
-        Some(item_list) => item_list.syntax().clone(),
-        None => {
-            let current_file =
-                path_under_caret.syntax().ancestors().find_map(ast::SourceFile::cast)?;
-            current_file.syntax().clone()
-        }
-    };
-    let source_analyzer = ctx.source_analyzer(&position, None);
-    let module_with_name_to_import = source_analyzer.module()?;
-
-    let name_ref_to_import =
-        path_under_caret.syntax().descendants().find_map(ast::NameRef::cast)?;
-    if source_analyzer
-        .resolve_path(ctx.db, &name_ref_to_import.syntax().ancestors().find_map(ast::Path::cast)?)
-        .is_some()
-    {
-        return None;
-    }
-
-    let name_to_import = name_ref_to_import.syntax().to_string();
-    let proposed_imports = ImportsLocator::new(ctx.db)
-        .find_imports(&name_to_import)
-        .into_iter()
-        .filter_map(|module_def| module_with_name_to_import.find_use_path(ctx.db, module_def))
-        .filter(|use_path| !use_path.segments.is_empty())
-        .take(20)
-        .collect::<BTreeSet<_>>();
-
+    let auto_import_assets = AutoImportAssets::new(&ctx)?;
+    let proposed_imports = auto_import_assets.search_for_imports(ctx.db);
     if proposed_imports.is_empty() {
         return None;
     }
 
-    let mut group = ctx.add_assist_group(format!("Import {}", name_to_import));
+    let assist_group_name = if proposed_imports.len() == 1 {
+        format!("Import `{}`", proposed_imports.iter().next().unwrap())
+    } else {
+        auto_import_assets.get_import_group_message()
+    };
+    let mut group = ctx.add_assist_group(assist_group_name);
     for import in proposed_imports {
         group.add_assist(AssistId("auto_import"), format!("Import `{}`", &import), |edit| {
-            edit.target(path_under_caret.syntax().text_range());
+            edit.target(auto_import_assets.syntax_under_caret.text_range());
             insert_use_statement(
-                &position,
-                path_under_caret.syntax(),
+                &auto_import_assets.syntax_under_caret,
+                &auto_import_assets.syntax_under_caret,
                 &import,
                 edit.text_edit_builder(),
             );
@@ -81,11 +61,232 @@ pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
     group.finish()
 }
 
+struct AutoImportAssets {
+    import_candidate: ImportCandidate,
+    module_with_name_to_import: Module,
+    syntax_under_caret: SyntaxNode,
+}
+
+impl AutoImportAssets {
+    fn new(ctx: &AssistCtx) -> Option<Self> {
+        if let Some(path_under_caret) = ctx.find_node_at_offset::<ast::Path>() {
+            Self::for_regular_path(path_under_caret, &ctx)
+        } else {
+            Self::for_method_call(ctx.find_node_at_offset()?, &ctx)
+        }
+    }
+
+    fn for_method_call(method_call: ast::MethodCallExpr, ctx: &AssistCtx) -> Option<Self> {
+        let syntax_under_caret = method_call.syntax().to_owned();
+        let source_analyzer = ctx.source_analyzer(&syntax_under_caret, None);
+        let module_with_name_to_import = source_analyzer.module()?;
+        Some(Self {
+            import_candidate: ImportCandidate::for_method_call(
+                &method_call,
+                &source_analyzer,
+                ctx.db,
+            )?,
+            module_with_name_to_import,
+            syntax_under_caret,
+        })
+    }
+
+    fn for_regular_path(path_under_caret: ast::Path, ctx: &AssistCtx) -> Option<Self> {
+        let syntax_under_caret = path_under_caret.syntax().to_owned();
+        if syntax_under_caret.ancestors().find_map(ast::UseItem::cast).is_some() {
+            return None;
+        }
+
+        let source_analyzer = ctx.source_analyzer(&syntax_under_caret, None);
+        let module_with_name_to_import = source_analyzer.module()?;
+        Some(Self {
+            import_candidate: ImportCandidate::for_regular_path(
+                &path_under_caret,
+                &source_analyzer,
+                ctx.db,
+            )?,
+            module_with_name_to_import,
+            syntax_under_caret,
+        })
+    }
+
+    fn get_search_query(&self) -> &str {
+        match &self.import_candidate {
+            ImportCandidate::UnqualifiedName(name) => name,
+            ImportCandidate::QualifierStart(qualifier_start) => qualifier_start,
+            ImportCandidate::TraitAssocItem(_, trait_assoc_item_name) => trait_assoc_item_name,
+            ImportCandidate::TraitMethod(_, trait_method_name) => trait_method_name,
+        }
+    }
+
+    fn get_import_group_message(&self) -> String {
+        match &self.import_candidate {
+            ImportCandidate::UnqualifiedName(name) => format!("Import {}", name),
+            ImportCandidate::QualifierStart(qualifier_start) => {
+                format!("Import {}", qualifier_start)
+            }
+            ImportCandidate::TraitAssocItem(_, trait_assoc_item_name) => {
+                format!("Import a trait for item {}", trait_assoc_item_name)
+            }
+            ImportCandidate::TraitMethod(_, trait_method_name) => {
+                format!("Import a trait for method {}", trait_method_name)
+            }
+        }
+    }
+
+    fn search_for_imports(&self, db: &RootDatabase) -> BTreeSet<ModPath> {
+        let _p = profile("auto_import::search_for_imports");
+        let current_crate = self.module_with_name_to_import.krate();
+        ImportsLocator::new(db)
+            .find_imports(&self.get_search_query())
+            .into_iter()
+            .filter_map(|module_def| match &self.import_candidate {
+                ImportCandidate::TraitAssocItem(assoc_item_type, _) => {
+                    let located_assoc_item = match module_def {
+                        ModuleDef::Function(located_function) => located_function
+                            .as_assoc_item(db)
+                            .map(|assoc| assoc.container(db))
+                            .and_then(Self::assoc_to_trait),
+                        ModuleDef::Const(located_const) => located_const
+                            .as_assoc_item(db)
+                            .map(|assoc| assoc.container(db))
+                            .and_then(Self::assoc_to_trait),
+                        _ => None,
+                    }?;
+
+                    let mut trait_candidates = FxHashSet::default();
+                    trait_candidates.insert(located_assoc_item.into());
+
+                    assoc_item_type
+                        .iterate_path_candidates(
+                            db,
+                            current_crate,
+                            &trait_candidates,
+                            None,
+                            |_, assoc| Self::assoc_to_trait(assoc.container(db)),
+                        )
+                        .map(ModuleDef::from)
+                }
+                ImportCandidate::TraitMethod(function_callee, _) => {
+                    let located_assoc_item =
+                        if let ModuleDef::Function(located_function) = module_def {
+                            located_function
+                                .as_assoc_item(db)
+                                .map(|assoc| assoc.container(db))
+                                .and_then(Self::assoc_to_trait)
+                        } else {
+                            None
+                        }?;
+
+                    let mut trait_candidates = FxHashSet::default();
+                    trait_candidates.insert(located_assoc_item.into());
+
+                    function_callee
+                        .iterate_method_candidates(
+                            db,
+                            current_crate,
+                            &trait_candidates,
+                            None,
+                            |_, function| {
+                                Self::assoc_to_trait(function.as_assoc_item(db)?.container(db))
+                            },
+                        )
+                        .map(ModuleDef::from)
+                }
+                _ => Some(module_def),
+            })
+            .filter_map(|module_def| self.module_with_name_to_import.find_use_path(db, module_def))
+            .filter(|use_path| !use_path.segments.is_empty())
+            .take(20)
+            .collect::<BTreeSet<_>>()
+    }
+
+    fn assoc_to_trait(assoc: AssocItemContainer) -> Option<Trait> {
+        if let AssocItemContainer::Trait(extracted_trait) = assoc {
+            Some(extracted_trait)
+        } else {
+            None
+        }
+    }
+}
+
+#[derive(Debug)]
+enum ImportCandidate {
+    /// Simple name like 'HashMap'
+    UnqualifiedName(String),
+    /// First part of the qualified name.
+    /// For 'std::collections::HashMap', that will be 'std'.
+    QualifierStart(String),
+    /// A trait associated function (with no self parameter) or associated constant.
+    /// For 'test_mod::TestEnum::test_function', `Type` is the `test_mod::TestEnum` expression type
+    /// and `String` is the `test_function`
+    TraitAssocItem(Type, String),
+    /// A trait method with self parameter.
+    /// For 'test_enum.test_method()', `Type` is the `test_enum` expression type
+    /// and `String` is the `test_method`
+    TraitMethod(Type, String),
+}
+
+impl ImportCandidate {
+    fn for_method_call(
+        method_call: &ast::MethodCallExpr,
+        source_analyzer: &SourceAnalyzer,
+        db: &impl HirDatabase,
+    ) -> Option<Self> {
+        if source_analyzer.resolve_method_call(method_call).is_some() {
+            return None;
+        }
+        Some(Self::TraitMethod(
+            source_analyzer.type_of(db, &method_call.expr()?)?,
+            method_call.name_ref()?.syntax().to_string(),
+        ))
+    }
+
+    fn for_regular_path(
+        path_under_caret: &ast::Path,
+        source_analyzer: &SourceAnalyzer,
+        db: &impl HirDatabase,
+    ) -> Option<Self> {
+        if source_analyzer.resolve_path(db, path_under_caret).is_some() {
+            return None;
+        }
+
+        let segment = path_under_caret.segment()?;
+        if let Some(qualifier) = path_under_caret.qualifier() {
+            let qualifier_start = qualifier.syntax().descendants().find_map(ast::NameRef::cast)?;
+            let qualifier_start_path =
+                qualifier_start.syntax().ancestors().find_map(ast::Path::cast)?;
+            if let Some(qualifier_start_resolution) =
+                source_analyzer.resolve_path(db, &qualifier_start_path)
+            {
+                let qualifier_resolution = if qualifier_start_path == qualifier {
+                    qualifier_start_resolution
+                } else {
+                    source_analyzer.resolve_path(db, &qualifier)?
+                };
+                if let PathResolution::Def(ModuleDef::Adt(assoc_item_path)) = qualifier_resolution {
+                    Some(ImportCandidate::TraitAssocItem(
+                        assoc_item_path.ty(db),
+                        segment.syntax().to_string(),
+                    ))
+                } else {
+                    None
+                }
+            } else {
+                Some(ImportCandidate::QualifierStart(qualifier_start.syntax().to_string()))
+            }
+        } else {
+            Some(ImportCandidate::UnqualifiedName(
+                segment.syntax().descendants().find_map(ast::NameRef::cast)?.syntax().to_string(),
+            ))
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};
-
     use super::*;
+    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};
 
     #[test]
     fn applicable_when_found_an_import() {
@@ -290,4 +491,303 @@ mod tests {
             ",
         );
     }
+
+    #[test]
+    fn not_applicable_for_imported_function() {
+        check_assist_not_applicable(
+            auto_import,
+            r"
+            pub mod test_mod {
+                pub fn test_function() {}
+            }
+
+            use test_mod::test_function;
+            fn main() {
+                test_function<|>
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn associated_struct_function() {
+        check_assist(
+            auto_import,
+            r"
+            mod test_mod {
+                pub struct TestStruct {}
+                impl TestStruct {
+                    pub fn test_function() {}
+                }
+            }
+
+            fn main() {
+                TestStruct::test_function<|>
+            }
+            ",
+            r"
+            use test_mod::TestStruct;
+
+            mod test_mod {
+                pub struct TestStruct {}
+                impl TestStruct {
+                    pub fn test_function() {}
+                }
+            }
+
+            fn main() {
+                TestStruct::test_function<|>
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn associated_struct_const() {
+        check_assist(
+            auto_import,
+            r"
+            mod test_mod {
+                pub struct TestStruct {}
+                impl TestStruct {
+                    const TEST_CONST: u8 = 42;
+                }
+            }
+
+            fn main() {
+                TestStruct::TEST_CONST<|>
+            }
+            ",
+            r"
+            use test_mod::TestStruct;
+
+            mod test_mod {
+                pub struct TestStruct {}
+                impl TestStruct {
+                    const TEST_CONST: u8 = 42;
+                }
+            }
+
+            fn main() {
+                TestStruct::TEST_CONST<|>
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn associated_trait_function() {
+        check_assist(
+            auto_import,
+            r"
+            mod test_mod {
+                pub trait TestTrait {
+                    fn test_function();
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    fn test_function() {}
+                }
+            }
+
+            fn main() {
+                test_mod::TestStruct::test_function<|>
+            }
+            ",
+            r"
+            use test_mod::TestTrait;
+
+            mod test_mod {
+                pub trait TestTrait {
+                    fn test_function();
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    fn test_function() {}
+                }
+            }
+
+            fn main() {
+                test_mod::TestStruct::test_function<|>
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn not_applicable_for_imported_trait_for_function() {
+        check_assist_not_applicable(
+            auto_import,
+            r"
+            mod test_mod {
+                pub trait TestTrait {
+                    fn test_function();
+                }
+                pub trait TestTrait2 {
+                    fn test_function();
+                }
+                pub enum TestEnum {
+                    One,
+                    Two,
+                }
+                impl TestTrait2 for TestEnum {
+                    fn test_function() {}
+                }
+                impl TestTrait for TestEnum {
+                    fn test_function() {}
+                }
+            }
+
+            use test_mod::TestTrait2;
+            fn main() {
+                test_mod::TestEnum::test_function<|>;
+            }
+            ",
+        )
+    }
+
+    #[test]
+    fn associated_trait_const() {
+        check_assist(
+            auto_import,
+            r"
+            mod test_mod {
+                pub trait TestTrait {
+                    const TEST_CONST: u8;
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    const TEST_CONST: u8 = 42;
+                }
+            }
+
+            fn main() {
+                test_mod::TestStruct::TEST_CONST<|>
+            }
+            ",
+            r"
+            use test_mod::TestTrait;
+
+            mod test_mod {
+                pub trait TestTrait {
+                    const TEST_CONST: u8;
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    const TEST_CONST: u8 = 42;
+                }
+            }
+
+            fn main() {
+                test_mod::TestStruct::TEST_CONST<|>
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn not_applicable_for_imported_trait_for_const() {
+        check_assist_not_applicable(
+            auto_import,
+            r"
+            mod test_mod {
+                pub trait TestTrait {
+                    const TEST_CONST: u8;
+                }
+                pub trait TestTrait2 {
+                    const TEST_CONST: f64;
+                }
+                pub enum TestEnum {
+                    One,
+                    Two,
+                }
+                impl TestTrait2 for TestEnum {
+                    const TEST_CONST: f64 = 42.0;
+                }
+                impl TestTrait for TestEnum {
+                    const TEST_CONST: u8 = 42;
+                }
+            }
+
+            use test_mod::TestTrait2;
+            fn main() {
+                test_mod::TestEnum::TEST_CONST<|>;
+            }
+            ",
+        )
+    }
+
+    #[test]
+    fn trait_method() {
+        check_assist(
+            auto_import,
+            r"
+            mod test_mod {
+                pub trait TestTrait {
+                    fn test_method(&self);
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    fn test_method(&self) {}
+                }
+            }
+
+            fn main() {
+                let test_struct = test_mod::TestStruct {};
+                test_struct.test_meth<|>od()
+            }
+            ",
+            r"
+            use test_mod::TestTrait;
+
+            mod test_mod {
+                pub trait TestTrait {
+                    fn test_method(&self);
+                }
+                pub struct TestStruct {}
+                impl TestTrait for TestStruct {
+                    fn test_method(&self) {}
+                }
+            }
+
+            fn main() {
+                let test_struct = test_mod::TestStruct {};
+                test_struct.test_meth<|>od()
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn not_applicable_for_imported_trait_for_method() {
+        check_assist_not_applicable(
+            auto_import,
+            r"
+            mod test_mod {
+                pub trait TestTrait {
+                    fn test_method(&self);
+                }
+                pub trait TestTrait2 {
+                    fn test_method(&self);
+                }
+                pub enum TestEnum {
+                    One,
+                    Two,
+                }
+                impl TestTrait2 for TestEnum {
+                    fn test_method(&self) {}
+                }
+                impl TestTrait for TestEnum {
+                    fn test_method(&self) {}
+                }
+            }
+
+            use test_mod::TestTrait2;
+            fn main() {
+                let one = test_mod::TestEnum::One;
+                one.test<|>_method();
+            }
+            ",
+        )
+    }
 }
diff --git a/crates/ra_hir/src/code_model.rs b/crates/ra_hir/src/code_model.rs
index 4fb679f6d..a56b8ab04 100644
--- a/crates/ra_hir/src/code_model.rs
+++ b/crates/ra_hir/src/code_model.rs
@@ -818,7 +818,7 @@ impl TypeParam {
     }
 }
 
-// FIXME: rename to `ImplBlock`
+// FIXME: rename from `ImplBlock` to `Impl`
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct ImplBlock {
     pub(crate) id: ImplId,