1 files changed, 93 insertions, 96 deletions

diff --git a/crates/ra_assists/src/handlers/auto_import.rs b/crates/ra_assists/src/handlers/auto_import.rs
index c564f5027..86615d659 100644
--- a/crates/ra_assists/src/handlers/auto_import.rs
+++ b/crates/ra_assists/src/handlers/auto_import.rs
@@ -36,15 +36,8 @@ use std::collections::BTreeSet;
 // # pub mod std { pub mod collections { pub struct HashMap { } } }
 // ```
 pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
-    let auto_import_assets = if let Some(path_under_caret) = ctx.find_node_at_offset::<ast::Path>()
-    {
-        AutoImportAssets::for_regular_path(path_under_caret, &ctx)?
-    } else {
-        AutoImportAssets::for_method_call(ctx.find_node_at_offset()?, &ctx)?
-    };
-
-    let proposed_imports = auto_import_assets
-        .search_for_imports(ctx.db, auto_import_assets.module_with_name_to_import);
+    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;
     }
@@ -54,7 +47,6 @@ pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
     } 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| {
@@ -77,6 +69,14 @@ struct AutoImportAssets {
 }
 
 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);
@@ -111,36 +111,33 @@ impl AutoImportAssets {
         })
     }
 
-    fn get_search_query(&self) -> String {
+    fn get_search_query(&self) -> &str {
         match &self.import_candidate {
-            ImportCandidate::UnqualifiedName(name_ref)
-            | ImportCandidate::QualifierStart(name_ref) => name_ref.syntax().to_string(),
-            ImportCandidate::TraitFunction(_, trait_function) => {
-                trait_function.syntax().to_string()
-            }
-            ImportCandidate::TraitMethod(_, trait_method) => trait_method.syntax().to_string(),
+            ImportCandidate::UnqualifiedName(name) => name,
+            ImportCandidate::QualifierStart(qualifier_start) => qualifier_start,
+            ImportCandidate::TraitFunction(_, trait_function_name) => trait_function_name,
+            ImportCandidate::TraitMethod(_, trait_method_name) => trait_method_name,
         }
     }
 
     fn get_import_group_message(&self) -> String {
         match &self.import_candidate {
-            ImportCandidate::UnqualifiedName(name_ref)
-            | ImportCandidate::QualifierStart(name_ref) => format!("Import {}", name_ref.syntax()),
-            ImportCandidate::TraitFunction(_, trait_function) => {
-                format!("Import a trait for function {}", trait_function.syntax())
+            ImportCandidate::UnqualifiedName(name) => format!("Import {}", name),
+            ImportCandidate::QualifierStart(qualifier_start) => {
+                format!("Import {}", qualifier_start)
             }
-            ImportCandidate::TraitMethod(_, trait_method) => {
-                format!("Import a trait for method {}", trait_method.syntax())
+            ImportCandidate::TraitFunction(_, trait_function_name) => {
+                format!("Import a trait for function {}", trait_function_name)
+            }
+            ImportCandidate::TraitMethod(_, trait_method_name) => {
+                format!("Import a trait for method {}", trait_method_name)
             }
         }
     }
 
-    fn search_for_imports(
-        &self,
-        db: &RootDatabase,
-        module_with_name_to_import: Module,
-    ) -> BTreeSet<ModPath> {
+    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()
@@ -148,49 +145,46 @@ impl AutoImportAssets {
                 ImportCandidate::TraitFunction(function_callee, _) => {
                     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.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::<()>
-                            },
-                        );
+                        let trait_candidates: FxHashSet<_> =
+                            Self::get_trait_candidates(db, located_function, current_crate)
+                                .into_iter()
+                                .map(|trait_candidate| trait_candidate.into())
+                                .collect();
+                        if !trait_candidates.is_empty() {
+                            function_callee.iterate_path_candidates(
+                                db,
+                                current_crate,
+                                &trait_candidates,
+                                None,
+                                |_, assoc| {
+                                    if let AssocContainerId::TraitId(trait_id) = assoc.container(db)
+                                    {
+                                        applicable_traits.push(
+                                            self.module_with_name_to_import.find_use_path(
+                                                db,
+                                                ModuleDef::Trait(trait_id.into()),
+                                            ),
+                                        );
+                                    };
+                                    None::<()>
+                                },
+                            );
+                        };
                     }
                     applicable_traits
                 }
                 ImportCandidate::TraitMethod(function_callee, _) => {
                     let mut applicable_traits = Vec::new();
                     if let ModuleDef::Function(located_function) = module_def {
-                        let trait_candidates: FxHashSet<_> = Self::get_trait_candidates(
-                            db,
-                            located_function,
-                            module_with_name_to_import.krate(),
-                        )
-                        .into_iter()
-                        .map(|trait_candidate| trait_candidate.into())
-                        .collect();
-
+                        let trait_candidates: FxHashSet<_> =
+                            Self::get_trait_candidates(db, located_function, current_crate)
+                                .into_iter()
+                                .map(|trait_candidate| trait_candidate.into())
+                                .collect();
                         if !trait_candidates.is_empty() {
                             function_callee.iterate_method_candidates(
                                 db,
-                                module_with_name_to_import.krate(),
+                                current_crate,
                                 &trait_candidates,
                                 None,
                                 |_, funciton| {
@@ -198,7 +192,7 @@ impl AutoImportAssets {
                                         funciton.container(db)
                                     {
                                         applicable_traits.push(
-                                            module_with_name_to_import.find_use_path(
+                                            self.module_with_name_to_import.find_use_path(
                                                 db,
                                                 ModuleDef::Trait(trait_id.into()),
                                             ),
@@ -211,7 +205,7 @@ impl AutoImportAssets {
                     }
                     applicable_traits
                 }
-                _ => vec![module_with_name_to_import.find_use_path(db, module_def)],
+                _ => vec![self.module_with_name_to_import.find_use_path(db, module_def)],
             })
             .flatten()
             .filter_map(std::convert::identity)
@@ -235,22 +229,19 @@ impl AutoImportAssets {
                 crate_def_map
                     .modules
                     .iter()
-                    .map(|(_, module_data)| {
-                        let mut traits = Vec::new();
-                        for module_def_id in module_data.scope.declarations() {
-                            if let ModuleDef::Trait(trait_candidate) = module_def_id.into() {
-                                if trait_candidate
-                                    .items(db)
-                                    .into_iter()
-                                    .any(|item| item == AssocItem::Function(called_function))
-                                {
-                                    traits.push(trait_candidate)
-                                }
-                            }
+                    .map(|(_, module_data)| module_data.scope.declarations())
+                    .flatten()
+                    .filter_map(|module_def_id| match module_def_id.into() {
+                        ModuleDef::Trait(trait_candidate)
+                            if trait_candidate
+                                .items(db)
+                                .into_iter()
+                                .any(|item| item == AssocItem::Function(called_function)) =>
+                        {
+                            Some(trait_candidate)
                         }
-                        traits
+                        _ => None,
                     })
-                    .flatten()
                     .collect::<FxHashSet<_>>()
             })
             .flatten()
@@ -259,12 +250,20 @@ impl AutoImportAssets {
 }
 
 #[derive(Debug)]
-// TODO kb rustdocs
 enum ImportCandidate {
-    UnqualifiedName(ast::NameRef),
-    QualifierStart(ast::NameRef),
-    TraitFunction(Type, ast::PathSegment),
-    TraitMethod(Type, ast::NameRef),
+    /// Simple name like 'HashMap'
+    UnqualifiedName(String),
+    /// First part of the qualified name.
+    /// For 'std::collections::HashMap', that will be 'std'.
+    QualifierStart(String),
+    /// A trait function that has no self parameter.
+    /// For 'test_mod::TestEnum::test_function', `Type` is the `test_mod::TestEnum` expression type
+    /// and `String`is the `test_function`
+    TraitFunction(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 {
@@ -278,7 +277,7 @@ impl ImportCandidate {
         }
         Some(Self::TraitMethod(
             source_analyzer.type_of(db, &method_call.expr()?)?,
-            method_call.name_ref()?,
+            method_call.name_ref()?.syntax().to_string(),
         ))
     }
 
@@ -299,36 +298,34 @@ impl ImportCandidate {
             if let Some(qualifier_start_resolution) =
                 source_analyzer.resolve_path(db, &qualifier_start_path)
             {
-                let qualifier_resolution = if &qualifier_start_path == path_under_caret {
+                let qualifier_resolution = if qualifier_start_path == qualifier {
                     qualifier_start_resolution
                 } else {
                     source_analyzer.resolve_path(db, &qualifier)?
                 };
                 if let PathResolution::Def(ModuleDef::Adt(function_callee)) = qualifier_resolution {
-                    Some(ImportCandidate::TraitFunction(function_callee.ty(db), segment))
+                    Some(ImportCandidate::TraitFunction(
+                        function_callee.ty(db),
+                        segment.syntax().to_string(),
+                    ))
                 } else {
                     None
                 }
             } else {
-                Some(ImportCandidate::QualifierStart(qualifier_start))
+                Some(ImportCandidate::QualifierStart(qualifier_start.syntax().to_string()))
             }
         } else {
-            if source_analyzer.resolve_path(db, path_under_caret).is_none() {
-                Some(ImportCandidate::UnqualifiedName(
-                    segment.syntax().descendants().find_map(ast::NameRef::cast)?,
-                ))
-            } else {
-                None
-            }
+            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() {