2 files changed, 127 insertions, 20 deletions

diff --git a/crates/ra_assists/src/handlers/auto_import.rs b/crates/ra_assists/src/handlers/auto_import.rs
index a25f0650d..a9778fab7 100644
--- a/crates/ra_assists/src/handlers/auto_import.rs
+++ b/crates/ra_assists/src/handlers/auto_import.rs
@@ -5,7 +5,13 @@ use crate::{
     assist_ctx::{Assist, AssistCtx},
     insert_use_statement, AssistId,
 };
-use hir::{db::HirDatabase, Adt, ModPath, Module, ModuleDef, PathResolution, SourceAnalyzer};
+use ast::{FnDefOwner, ModuleItem, ModuleItemOwner};
+use hir::{
+    db::{DefDatabase, HirDatabase},
+    Adt, AssocContainerId, Crate, Function, HasSource, InFile, ModPath, Module, ModuleDef,
+    PathResolution, SourceAnalyzer, SourceBinder, Trait,
+};
+use rustc_hash::FxHashSet;
 use std::collections::BTreeSet;
 
 // Assist: auto_import
@@ -135,21 +141,88 @@ impl ImportCandidate {
         ImportsLocator::new(db)
             .find_imports(&self.get_search_query())
             .into_iter()
-            .filter_map(|module_def| match self {
+            .map(|module_def| match self {
                 ImportCandidate::TraitFunction(function_callee, _) => {
-                    if let ModuleDef::Function(function) = module_def {
-                        dbg!(function);
-                        todo!()
-                    } else {
-                        None
+                    let mut applicable_traits = Vec::new();
+                    if let ModuleDef::Function(located_function) = module_def {
+                        let trait_candidates = Self::get_trait_candidates(
+                            db,
+                            located_function,
+                            module_with_name_to_import.krate(),
+                        )
+                        .into_iter()
+                        .map(|trait_candidate| trait_candidate.into())
+                        .collect();
+
+                        function_callee.ty(db).iterate_path_candidates(
+                            db,
+                            module_with_name_to_import.krate(),
+                            &trait_candidates,
+                            None,
+                            |_, assoc| {
+                                if let AssocContainerId::TraitId(trait_id) = assoc.container(db) {
+                                    applicable_traits.push(
+                                        module_with_name_to_import
+                                            .find_use_path(db, ModuleDef::Trait(trait_id.into())),
+                                    );
+                                };
+                                None::<()>
+                            },
+                        );
                     }
+                    applicable_traits
                 }
-                _ => module_with_name_to_import.find_use_path(db, module_def),
+                _ => vec![module_with_name_to_import.find_use_path(db, module_def)],
             })
+            .flatten()
+            .filter_map(std::convert::identity)
             .filter(|use_path| !use_path.segments.is_empty())
             .take(20)
             .collect::<BTreeSet<_>>()
     }
+
+    fn get_trait_candidates(
+        db: &RootDatabase,
+        called_function: Function,
+        root_crate: Crate,
+    ) -> FxHashSet<Trait> {
+        let mut source_binder = SourceBinder::new(db);
+        root_crate
+            .dependencies(db)
+            .into_iter()
+            .map(|dependency| db.crate_def_map(dependency.krate.into()))
+            .chain(std::iter::once(db.crate_def_map(root_crate.into())))
+            .map(|crate_def_map| {
+                crate_def_map
+                    .modules
+                    .iter()
+                    .filter_map(|(_, module_data)| module_data.declaration_source(db))
+                    .filter_map(|in_file_module| {
+                        Some((in_file_module.file_id, in_file_module.value.item_list()?.items()))
+                    })
+                    .map(|(file_id, item_list)| {
+                        let mut if_file_trait_defs = Vec::new();
+                        for module_item in item_list {
+                            if let ModuleItem::TraitDef(trait_def) = module_item {
+                                if let Some(item_list) = trait_def.item_list() {
+                                    if item_list
+                                        .functions()
+                                        .any(|fn_def| fn_def == called_function.source(db).value)
+                                    {
+                                        if_file_trait_defs.push(InFile::new(file_id, trait_def))
+                                    }
+                                }
+                            }
+                        }
+                        if_file_trait_defs
+                    })
+                    .flatten()
+                    .filter_map(|in_file_trait_def| source_binder.to_def(in_file_trait_def))
+                    .collect::<FxHashSet<_>>()
+            })
+            .flatten()
+            .collect()
+    }
 }
 
 #[cfg(test)]
@@ -452,7 +525,45 @@ mod tests {
     }
 
     #[test]
-    #[ignore] // TODO kb
+    fn not_applicable_for_imported_trait() {
+        check_assist_not_applicable(
+            auto_import,
+            r"
+            mod test_mod {
+                pub trait TestTrait {
+                    fn test_method(&self);
+                    fn test_function();
+                }
+
+                pub trait TestTrait2 {
+                    fn test_method(&self);
+                    fn test_function();
+                }
+                pub enum TestEnum {
+                    One,
+                    Two,
+                }
+
+                impl TestTrait2 for TestEnum {
+                    fn test_method(&self) {}
+                    fn test_function() {}
+                }
+
+                impl TestTrait for TestEnum {
+                    fn test_method(&self) {}
+                    fn test_function() {}
+                }
+            }
+
+            use test_mod::TestTrait2;
+            fn main() {
+                test_mod::TestEnum::test_function<|>;
+            }
+            ",
+        )
+    }
+
+    #[test]
     fn trait_method() {
         check_assist(
             auto_import,
diff --git a/crates/ra_hir/src/code_model.rs b/crates/ra_hir/src/code_model.rs
index 4fb679f6d..73158b8bd 100644
--- a/crates/ra_hir/src/code_model.rs
+++ b/crates/ra_hir/src/code_model.rs
@@ -123,7 +123,7 @@ impl_froms!(
 );
 
 pub use hir_def::{
-    attr::Attrs, item_scope::ItemInNs, visibility::Visibility, AssocItemId, AssocItemLoc,
+    attr::Attrs, item_scope::ItemInNs, visibility::Visibility, AssocContainerId, AssocItemId,
 };
 use rustc_hash::FxHashSet;
 
@@ -696,16 +696,12 @@ impl AssocItem {
             AssocItem::TypeAlias(t) => t.module(db),
         }
     }
-    pub fn container(self, db: &impl DefDatabase) -> AssocItemContainer {
-        let container = match self {
-            AssocItem::Function(it) => it.id.lookup(db).container,
-            AssocItem::Const(it) => it.id.lookup(db).container,
-            AssocItem::TypeAlias(it) => it.id.lookup(db).container,
-        };
-        match container {
-            AssocContainerId::TraitId(id) => AssocItemContainer::Trait(id.into()),
-            AssocContainerId::ImplId(id) => AssocItemContainer::ImplBlock(id.into()),
-            AssocContainerId::ContainerId(_) => panic!("invalid AssocItem"),
+
+    pub fn container(self, db: &impl DefDatabase) -> AssocContainerId {
+        match self {
+            AssocItem::Function(f) => f.id.lookup(db).container,
+            AssocItem::Const(c) => c.id.lookup(db).container,
+            AssocItem::TypeAlias(t) => t.id.lookup(db).container,
         }
     }
 }