Resolve methods and functions better

2 files changed, 74 insertions, 9 deletions

diff --git a/crates/ra_assists/src/handlers/auto_import.rs b/crates/ra_assists/src/handlers/auto_import.rs
index 1fb701da5..10984d8ad 100644
--- a/crates/ra_assists/src/handlers/auto_import.rs
+++ b/crates/ra_assists/src/handlers/auto_import.rs
@@ -46,9 +46,9 @@ pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
 
     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()
+    if dbg!(source_analyzer
+        .resolve_path(ctx.db, &name_ref_to_import.syntax().ancestors().find_map(ast::Path::cast)?))
+    .is_some()
     {
         return None;
     }
@@ -290,4 +290,21 @@ 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<|>
+            }
+            ",
+        );
+    }
 }
diff --git a/crates/ra_hir/src/source_analyzer.rs b/crates/ra_hir/src/source_analyzer.rs
index bb9a35c5d..49e1a10e0 100644
--- a/crates/ra_hir/src/source_analyzer.rs
+++ b/crates/ra_hir/src/source_analyzer.rs
@@ -20,7 +20,10 @@ use hir_def::{
 use hir_expand::{
     hygiene::Hygiene, name::AsName, AstId, HirFileId, InFile, MacroCallId, MacroCallKind,
 };
-use hir_ty::{InEnvironment, InferenceResult, TraitEnvironment};
+use hir_ty::{
+    method_resolution::{iterate_method_candidates, LookupMode},
+    Canonical, InEnvironment, InferenceResult, TraitEnvironment,
+};
 use ra_syntax::{
     ast::{self, AstNode},
     AstPtr, SyntaxNode, SyntaxNodePtr, SyntaxToken, TextRange, TextUnit,
@@ -28,8 +31,8 @@ use ra_syntax::{
 use rustc_hash::FxHashSet;
 
 use crate::{
-    db::HirDatabase, Adt, Const, DefWithBody, EnumVariant, Function, Local, MacroDef, Name, Path,
-    ScopeDef, Static, Struct, Trait, Type, TypeAlias, TypeParam,
+    db::HirDatabase, Adt, AssocItem, Const, DefWithBody, EnumVariant, Function, Local, MacroDef,
+    ModuleDef, Name, Path, ScopeDef, Static, Struct, Trait, Type, TypeAlias, TypeParam,
 };
 
 /// `SourceAnalyzer` is a convenience wrapper which exposes HIR API in terms of
@@ -289,9 +292,11 @@ impl SourceAnalyzer {
 
     pub fn resolve_path(&self, db: &impl HirDatabase, path: &ast::Path) -> Option<PathResolution> {
         if let Some(path_expr) = path.syntax().parent().and_then(ast::PathExpr::cast) {
-            let expr_id = self.expr_id(&path_expr.into())?;
-            if let Some(assoc) = self.infer.as_ref()?.assoc_resolutions_for_expr(expr_id) {
-                return Some(PathResolution::AssocItem(assoc.into()));
+            let path_resolution = self
+                .resolve_as_full_path(path_expr.clone())
+                .or_else(|| self.resolve_as_path_to_method(db, &path_expr));
+            if path_resolution.is_some() {
+                return path_resolution;
             }
         }
         if let Some(path_pat) = path.syntax().parent().and_then(ast::PathPat::cast) {
@@ -305,6 +310,49 @@ impl SourceAnalyzer {
         self.resolve_hir_path(db, &hir_path)
     }
 
+    fn resolve_as_full_path(&self, path_expr: ast::PathExpr) -> Option<PathResolution> {
+        let expr_id = self.expr_id(&path_expr.into())?;
+        self.infer
+            .as_ref()?
+            .assoc_resolutions_for_expr(expr_id)
+            .map(|assoc| PathResolution::AssocItem(assoc.into()))
+    }
+
+    fn resolve_as_path_to_method(
+        &self,
+        db: &impl HirDatabase,
+        path_expr: &ast::PathExpr,
+    ) -> Option<PathResolution> {
+        let full_path = path_expr.path()?;
+        let path_to_method = full_path.qualifier()?;
+        let method_name = full_path.segment()?.syntax().to_string();
+        match self.resolve_path(db, &path_to_method)? {
+            PathResolution::Def(ModuleDef::Adt(adt)) => {
+                let ty = adt.ty(db);
+                iterate_method_candidates(
+                    &Canonical { value: ty.ty.value, num_vars: 0 },
+                    db,
+                    ty.ty.environment,
+                    self.resolver.krate()?,
+                    &self.resolver.traits_in_scope(db),
+                    None,
+                    LookupMode::Path,
+                    |_, assoc_item_id| {
+                        let assoc = assoc_item_id.into();
+                        if let AssocItem::Function(function) = assoc {
+                            if function.name(db).to_string() == method_name {
+                                return Some(assoc);
+                            }
+                        }
+                        None
+                    },
+                )
+            }
+            _ => None,
+        }
+        .map(PathResolution::AssocItem)
+    }
+
     fn resolve_local_name(&self, name_ref: &ast::NameRef) -> Option<ScopeEntryWithSyntax> {
         let name = name_ref.as_name();
         let source_map = self.body_source_map.as_ref()?;