1 files changed, 92 insertions, 20 deletions

diff --git a/crates/ra_assists/src/handlers/auto_import.rs b/crates/ra_assists/src/handlers/auto_import.rs
index 27d96b941..a25f0650d 100644
--- a/crates/ra_assists/src/handlers/auto_import.rs
+++ b/crates/ra_assists/src/handlers/auto_import.rs
@@ -1,10 +1,11 @@
-use ra_ide_db::imports_locator::ImportsLocator;
+use ra_ide_db::{imports_locator::ImportsLocator, RootDatabase};
 use ra_syntax::ast::{self, AstNode};
 
 use crate::{
     assist_ctx::{Assist, AssistCtx},
     insert_use_statement, AssistId,
 };
+use hir::{db::HirDatabase, Adt, ModPath, Module, ModuleDef, PathResolution, SourceAnalyzer};
 use std::collections::BTreeSet;
 
 // Assist: auto_import
@@ -44,29 +45,13 @@ pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
     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 import_candidate = ImportCandidate::new(&path_under_caret, &source_analyzer, ctx.db)?;
+    let proposed_imports = import_candidate.search_for_imports(ctx.db, module_with_name_to_import);
     if proposed_imports.is_empty() {
         return None;
     }
 
-    let mut group = ctx.add_assist_group(format!("Import {}", name_to_import));
+    let mut group = ctx.add_assist_group(format!("Import {}", import_candidate.get_search_query()));
     for import in proposed_imports {
         group.add_assist(AssistId("auto_import"), format!("Import `{}`", &import), |edit| {
             edit.target(path_under_caret.syntax().text_range());
@@ -81,6 +66,92 @@ pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
     group.finish()
 }
 
+#[derive(Debug)]
+// TODO kb rustdocs
+enum ImportCandidate {
+    UnqualifiedName(ast::NameRef),
+    QualifierStart(ast::NameRef),
+    TraitFunction(Adt, ast::PathSegment),
+}
+
+impl ImportCandidate {
+    // TODO kb refactor this mess
+    fn new(
+        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 == path_under_caret {
+                    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, segment))
+                } else {
+                    None
+                }
+            } else {
+                Some(ImportCandidate::QualifierStart(qualifier_start))
+            }
+        } 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
+            }
+        }
+    }
+
+    fn get_search_query(&self) -> String {
+        match self {
+            ImportCandidate::UnqualifiedName(name_ref)
+            | ImportCandidate::QualifierStart(name_ref) => name_ref.syntax().to_string(),
+            ImportCandidate::TraitFunction(_, trait_function) => {
+                trait_function.syntax().to_string()
+            }
+        }
+    }
+
+    fn search_for_imports(
+        &self,
+        db: &RootDatabase,
+        module_with_name_to_import: Module,
+    ) -> BTreeSet<ModPath> {
+        ImportsLocator::new(db)
+            .find_imports(&self.get_search_query())
+            .into_iter()
+            .filter_map(|module_def| match self {
+                ImportCandidate::TraitFunction(function_callee, _) => {
+                    if let ModuleDef::Function(function) = module_def {
+                        dbg!(function);
+                        todo!()
+                    } else {
+                        None
+                    }
+                }
+                _ => module_with_name_to_import.find_use_path(db, module_def),
+            })
+            .filter(|use_path| !use_path.segments.is_empty())
+            .take(20)
+            .collect::<BTreeSet<_>>()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};
@@ -381,6 +452,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore] // TODO kb
     fn trait_method() {
         check_assist(
             auto_import,