Properly qualify trait methods in qualify_path assist

5 files changed, 118 insertions, 70 deletions

diff --git a/crates/assists/src/handlers/auto_import.rs b/crates/assists/src/handlers/auto_import.rs
index e595b5b93..13e390a1f 100644
--- a/crates/assists/src/handlers/auto_import.rs
+++ b/crates/assists/src/handlers/auto_import.rs
@@ -45,7 +45,7 @@ pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
     let group = import_group_message(import_assets.import_candidate());
     let scope = ImportScope::find_insert_use_container(import_assets.syntax_under_caret(), ctx)?;
     let syntax = scope.as_syntax_node();
-    for import in proposed_imports {
+    for (import, _) in proposed_imports {
         acc.add_group(
             &group,
             AssistId("auto_import", AssistKind::QuickFix),
diff --git a/crates/assists/src/handlers/qualify_path.rs b/crates/assists/src/handlers/qualify_path.rs
index bbbf47bef..479ff498c 100644
--- a/crates/assists/src/handlers/qualify_path.rs
+++ b/crates/assists/src/handlers/qualify_path.rs
@@ -1,6 +1,12 @@
-use std::collections::BTreeSet;
-
-use syntax::{ast, AstNode, TextRange};
+use std::iter;
+
+use hir::AsName;
+use ide_db::RootDatabase;
+use syntax::{
+    ast,
+    ast::{make, ArgListOwner},
+    AstNode, TextRange,
+};
 use test_utils::mark;
 
 use crate::{
@@ -10,6 +16,8 @@ use crate::{
     AssistId, AssistKind, GroupLabel,
 };
 
+const ASSIST_ID: AssistId = AssistId("qualify_path", AssistKind::QuickFix);
+
 // Assist: qualify_path
 //
 // If the name is unresolved, provides all possible qualified paths for it.
@@ -53,30 +61,14 @@ pub(crate) fn qualify_path(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
         ImportCandidate::UnqualifiedName(candidate) => {
             qualify_path_unqualified_name(acc, proposed_imports, range, &candidate.name)
         }
-        ImportCandidate::TraitAssocItem(candidate) => {
+        ImportCandidate::TraitAssocItem(_) => {
             let path = ast::Path::cast(import_assets.syntax_under_caret().clone())?;
             let (qualifier, segment) = (path.qualifier()?, path.segment()?);
-            qualify_path_trait_assoc_item(
-                acc,
-                proposed_imports,
-                range,
-                qualifier,
-                segment,
-                &candidate.name,
-            )
+            qualify_path_trait_assoc_item(acc, proposed_imports, range, qualifier, segment)
         }
-        ImportCandidate::TraitMethod(candidate) => {
+        ImportCandidate::TraitMethod(_) => {
             let mcall_expr = ast::MethodCallExpr::cast(import_assets.syntax_under_caret().clone())?;
-            let receiver = mcall_expr.receiver()?;
-            let name_ref = mcall_expr.name_ref()?;
-            qualify_path_trait_method(
-                acc,
-                proposed_imports,
-                range,
-                receiver,
-                name_ref,
-                &candidate.name,
-            )
+            qualify_path_trait_method(acc, ctx.sema.db, proposed_imports, range, mcall_expr)?;
         }
     };
     Some(())
@@ -85,17 +77,17 @@ pub(crate) fn qualify_path(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
 // a test that covers this -> `associated_struct_const`
 fn qualify_path_qualifier_start(
     acc: &mut Assists,
-    proposed_imports: BTreeSet<hir::ModPath>,
+    proposed_imports: Vec<(hir::ModPath, hir::ItemInNs)>,
     range: TextRange,
     segment: ast::PathSegment,
-    qualifier_start: &str,
+    qualifier_start: &ast::NameRef,
 ) {
     mark::hit!(qualify_path_qualifier_start);
     let group_label = GroupLabel(format!("Qualify {}", qualifier_start));
-    for import in proposed_imports {
+    for (import, _) in proposed_imports {
         acc.add_group(
             &group_label,
-            AssistId("qualify_path", AssistKind::QuickFix),
+            ASSIST_ID,
             format!("Qualify with `{}`", &import),
             range,
             |builder| {
@@ -109,16 +101,16 @@ fn qualify_path_qualifier_start(
 // a test that covers this -> `applicable_when_found_an_import_partial`
 fn qualify_path_unqualified_name(
     acc: &mut Assists,
-    proposed_imports: BTreeSet<hir::ModPath>,
+    proposed_imports: Vec<(hir::ModPath, hir::ItemInNs)>,
     range: TextRange,
-    name: &str,
+    name: &ast::NameRef,
 ) {
     mark::hit!(qualify_path_unqualified_name);
     let group_label = GroupLabel(format!("Qualify {}", name));
-    for import in proposed_imports {
+    for (import, _) in proposed_imports {
         acc.add_group(
             &group_label,
-            AssistId("qualify_path", AssistKind::QuickFix),
+            ASSIST_ID,
             format!("Qualify as `{}`", &import),
             range,
             |builder| builder.replace(range, mod_path_to_ast(&import).to_string()),
@@ -129,18 +121,17 @@ fn qualify_path_unqualified_name(
 // a test that covers this -> `associated_trait_const`
 fn qualify_path_trait_assoc_item(
     acc: &mut Assists,
-    proposed_imports: BTreeSet<hir::ModPath>,
+    proposed_imports: Vec<(hir::ModPath, hir::ItemInNs)>,
     range: TextRange,
     qualifier: ast::Path,
     segment: ast::PathSegment,
-    trait_assoc_item_name: &str,
 ) {
     mark::hit!(qualify_path_trait_assoc_item);
-    let group_label = GroupLabel(format!("Qualify {}", trait_assoc_item_name));
-    for import in proposed_imports {
+    let group_label = GroupLabel(format!("Qualify {}", &segment));
+    for (import, _) in proposed_imports {
         acc.add_group(
             &group_label,
-            AssistId("qualify_path", AssistKind::QuickFix),
+            ASSIST_ID,
             format!("Qualify with cast as `{}`", &import),
             range,
             |builder| {
@@ -154,33 +145,74 @@ fn qualify_path_trait_assoc_item(
 // a test that covers this -> `trait_method`
 fn qualify_path_trait_method(
     acc: &mut Assists,
-    proposed_imports: BTreeSet<hir::ModPath>,
+    db: &RootDatabase,
+    proposed_imports: Vec<(hir::ModPath, hir::ItemInNs)>,
     range: TextRange,
-    receiver: ast::Expr,
-    name_ref: ast::NameRef,
-    trait_method_name: &str,
-) {
+    mcall_expr: ast::MethodCallExpr,
+) -> Option<()> {
     mark::hit!(qualify_path_trait_method);
+
+    let receiver = mcall_expr.receiver()?;
+    let trait_method_name = mcall_expr.name_ref()?;
+    let arg_list = mcall_expr.arg_list().map(|arg_list| arg_list.args());
     let group_label = GroupLabel(format!("Qualify {}", trait_method_name));
-    for import in proposed_imports {
+    let find_method = |item: &hir::AssocItem| {
+        item.name(db).map(|name| name == trait_method_name.as_name()).unwrap_or(false)
+    };
+    for (import, trait_) in proposed_imports.into_iter().filter_map(filter_trait) {
         acc.add_group(
             &group_label,
-            AssistId("qualify_path", AssistKind::QuickFix), // < Does this still count as quickfix?
+            ASSIST_ID,
             format!("Qualify `{}`", &import),
             range,
             |builder| {
                 let import = mod_path_to_ast(&import);
-                // TODO: check the receiver self type and emit refs accordingly, don't discard other function parameters
-                builder.replace(range, format!("{}::{}(&{})", import, name_ref, receiver));
+                if let Some(hir::AssocItem::Function(method)) =
+                    trait_.items(db).into_iter().find(find_method)
+                {
+                    if let Some(self_access) = method.self_param(db).map(|sp| sp.access(db)) {
+                        let receiver = receiver.clone();
+                        let receiver = match self_access {
+                            hir::Access::Shared => make::expr_ref(receiver, false),
+                            hir::Access::Exclusive => make::expr_ref(receiver, true),
+                            hir::Access::Owned => receiver,
+                        };
+                        builder.replace(
+                            range,
+                            format!(
+                                "{}::{}{}",
+                                import,
+                                trait_method_name,
+                                match arg_list.clone() {
+                                    Some(args) => make::arg_list(iter::once(receiver).chain(args)),
+                                    None => make::arg_list(iter::once(receiver)),
+                                }
+                            ),
+                        );
+                    }
+                }
             },
         );
     }
+    Some(())
+}
+
+fn filter_trait(
+    (import, trait_): (hir::ModPath, hir::ItemInNs),
+) -> Option<(hir::ModPath, hir::Trait)> {
+    if let hir::ModuleDef::Trait(trait_) = hir::ModuleDef::from(trait_.as_module_def_id()?) {
+        Some((import, trait_))
+    } else {
+        None
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};
+
+    use super::*;
+
     #[test]
     fn applicable_when_found_an_import_partial() {
         mark::check!(qualify_path_unqualified_name);
diff --git a/crates/assists/src/tests/generated.rs b/crates/assists/src/tests/generated.rs
index 41f536574..7d5618263 100644
--- a/crates/assists/src/tests/generated.rs
+++ b/crates/assists/src/tests/generated.rs
@@ -713,6 +713,25 @@ fn handle(action: Action) {
 }
 
 #[test]
+fn doctest_qualify_path() {
+    check_doc_test(
+        "qualify_path",
+        r#####"
+fn main() {
+    let map = HashMap<|>::new();
+}
+pub mod std { pub mod collections { pub struct HashMap { } } }
+"#####,
+        r#####"
+fn main() {
+    let map = std::collections::HashMap::new();
+}
+pub mod std { pub mod collections { pub struct HashMap { } } }
+"#####,
+    )
+}
+
+#[test]
 fn doctest_remove_dbg() {
     check_doc_test(
         "remove_dbg",
diff --git a/crates/assists/src/utils/import_assets.rs b/crates/assists/src/utils/import_assets.rs
index 601f51098..23db3a74b 100644
--- a/crates/assists/src/utils/import_assets.rs
+++ b/crates/assists/src/utils/import_assets.rs
@@ -1,6 +1,4 @@
 //! Look up accessible paths for items.
-use std::collections::BTreeSet;
-
 use either::Either;
 use hir::{AsAssocItem, AssocItemContainer, ModuleDef, Semantics};
 use ide_db::{imports_locator, RootDatabase};
@@ -29,12 +27,12 @@ pub(crate) enum ImportCandidate {
 #[derive(Debug)]
 pub(crate) struct TraitImportCandidate {
     pub ty: hir::Type,
-    pub name: String,
+    pub name: ast::NameRef,
 }
 
 #[derive(Debug)]
 pub(crate) struct PathImportCandidate {
-    pub name: String,
+    pub name: ast::NameRef,
 }
 
 #[derive(Debug)]
@@ -86,9 +84,9 @@ impl ImportAssets {
     fn get_search_query(&self) -> &str {
         match &self.import_candidate {
             ImportCandidate::UnqualifiedName(candidate)
-            | ImportCandidate::QualifierStart(candidate) => &candidate.name,
+            | ImportCandidate::QualifierStart(candidate) => candidate.name.text(),
             ImportCandidate::TraitAssocItem(candidate)
-            | ImportCandidate::TraitMethod(candidate) => &candidate.name,
+            | ImportCandidate::TraitMethod(candidate) => candidate.name.text(),
         }
     }
 
@@ -96,7 +94,7 @@ impl ImportAssets {
         &self,
         sema: &Semantics<RootDatabase>,
         config: &InsertUseConfig,
-    ) -> BTreeSet<hir::ModPath> {
+    ) -> Vec<(hir::ModPath, hir::ItemInNs)> {
         let _p = profile::span("import_assists::search_for_imports");
         self.search_for(sema, Some(config.prefix_kind))
     }
@@ -106,7 +104,7 @@ impl ImportAssets {
     pub(crate) fn search_for_relative_paths(
         &self,
         sema: &Semantics<RootDatabase>,
-    ) -> BTreeSet<hir::ModPath> {
+    ) -> Vec<(hir::ModPath, hir::ItemInNs)> {
         let _p = profile::span("import_assists::search_for_relative_paths");
         self.search_for(sema, None)
     }
@@ -115,7 +113,7 @@ impl ImportAssets {
         &self,
         sema: &Semantics<RootDatabase>,
         prefixed: Option<hir::PrefixKind>,
-    ) -> BTreeSet<hir::ModPath> {
+    ) -> Vec<(hir::ModPath, hir::ItemInNs)> {
         let db = sema.db;
         let mut trait_candidates = FxHashSet::default();
         let current_crate = self.module_with_name_to_import.krate();
@@ -181,7 +179,7 @@ impl ImportAssets {
             }
         };
 
-        imports_locator::find_imports(sema, current_crate, &self.get_search_query())
+        let mut res = imports_locator::find_imports(sema, current_crate, &self.get_search_query())
             .into_iter()
             .filter_map(filter)
             .filter_map(|candidate| {
@@ -191,10 +189,13 @@ impl ImportAssets {
                 } else {
                     self.module_with_name_to_import.find_use_path(db, item)
                 }
+                .map(|path| (path, item))
             })
-            .filter(|use_path| !use_path.segments.is_empty())
+            .filter(|(use_path, _)| !use_path.segments.is_empty())
             .take(20)
-            .collect::<BTreeSet<_>>()
+            .collect::<Vec<_>>();
+        res.sort_by_key(|(path, _)| path.clone());
+        res
     }
 
     fn assoc_to_trait(assoc: AssocItemContainer) -> Option<hir::Trait> {
@@ -215,7 +216,7 @@ impl ImportCandidate {
             Some(_) => None,
             None => Some(Self::TraitMethod(TraitImportCandidate {
                 ty: sema.type_of_expr(&method_call.receiver()?)?,
-                name: method_call.name_ref()?.syntax().to_string(),
+                name: method_call.name_ref()?,
             })),
         }
     }
@@ -243,24 +244,17 @@ impl ImportCandidate {
                     hir::PathResolution::Def(hir::ModuleDef::Adt(assoc_item_path)) => {
                         ImportCandidate::TraitAssocItem(TraitImportCandidate {
                             ty: assoc_item_path.ty(sema.db),
-                            name: segment.syntax().to_string(),
+                            name: segment.name_ref()?,
                         })
                     }
                     _ => return None,
                 }
             } else {
-                ImportCandidate::QualifierStart(PathImportCandidate {
-                    name: qualifier_start.syntax().to_string(),
-                })
+                ImportCandidate::QualifierStart(PathImportCandidate { name: qualifier_start })
             }
         } else {
             ImportCandidate::UnqualifiedName(PathImportCandidate {
-                name: segment
-                    .syntax()
-                    .descendants()
-                    .find_map(ast::NameRef::cast)?
-                    .syntax()
-                    .to_string(),
+                name: segment.syntax().descendants().find_map(ast::NameRef::cast)?,
             })
         };
         Some(candidate)
diff --git a/crates/syntax/src/ast/make.rs b/crates/syntax/src/ast/make.rs
index 74dbdfaf7..5b06cb767 100644
--- a/crates/syntax/src/ast/make.rs
+++ b/crates/syntax/src/ast/make.rs
@@ -172,6 +172,9 @@ pub fn expr_call(f: ast::Expr, arg_list: ast::ArgList) -> ast::Expr {
 pub fn expr_method_call(receiver: ast::Expr, method: &str, arg_list: ast::ArgList) -> ast::Expr {
     expr_from_text(&format!("{}.{}{}", receiver, method, arg_list))
 }
+pub fn expr_ref(expr: ast::Expr, exclusive: bool) -> ast::Expr {
+    expr_from_text(&if exclusive { format!("&mut {}", expr) } else { format!("&{}", expr) })
+}
 fn expr_from_text(text: &str) -> ast::Expr {
     ast_from_text(&format!("const C: () = {};", text))
 }