goto defenition works for type-inferred methods

2 files changed, 71 insertions, 5 deletions

diff --git a/crates/ra_hir/src/ty.rs b/crates/ra_hir/src/ty.rs
index 0692d3b2a..fa46ddfe9 100644
--- a/crates/ra_hir/src/ty.rs
+++ b/crates/ra_hir/src/ty.rs
@@ -28,6 +28,7 @@ use log;
 use ena::unify::{InPlaceUnificationTable, UnifyKey, UnifyValue, NoError};
 use ra_arena::map::ArenaMap;
 use join_to_string::join;
+use rustc_hash::FxHashMap;
 
 use ra_db::Cancelable;
 
@@ -448,14 +449,14 @@ fn type_for_struct(db: &impl HirDatabase, s: Struct) -> Cancelable<Ty> {
     })
 }
 
-pub fn type_for_enum(db: &impl HirDatabase, s: Enum) -> Cancelable<Ty> {
+pub(crate) fn type_for_enum(db: &impl HirDatabase, s: Enum) -> Cancelable<Ty> {
     Ok(Ty::Adt {
         def_id: s.def_id(),
         name: s.name(db)?.unwrap_or_else(Name::missing),
     })
 }
 
-pub fn type_for_enum_variant(db: &impl HirDatabase, ev: EnumVariant) -> Cancelable<Ty> {
+pub(crate) fn type_for_enum_variant(db: &impl HirDatabase, ev: EnumVariant) -> Cancelable<Ty> {
     let enum_parent = ev.parent_enum(db)?;
 
     type_for_enum(db, enum_parent)
@@ -512,10 +513,18 @@ pub(super) fn type_for_field(
 /// The result of type inference: A mapping from expressions and patterns to types.
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct InferenceResult {
+    /// For each method call expr, record the function it resolved to.
+    method_resolutions: FxHashMap<ExprId, DefId>,
     type_of_expr: ArenaMap<ExprId, Ty>,
     type_of_pat: ArenaMap<PatId, Ty>,
 }
 
+impl InferenceResult {
+    pub fn method_resolution(&self, expr: ExprId) -> Option<DefId> {
+        self.method_resolutions.get(&expr).map(|it| *it)
+    }
+}
+
 impl Index<ExprId> for InferenceResult {
     type Output = Ty;
 
@@ -541,6 +550,7 @@ struct InferenceContext<'a, D: HirDatabase> {
     module: Module,
     impl_block: Option<ImplBlock>,
     var_unification_table: InPlaceUnificationTable<TypeVarId>,
+    method_resolutions: FxHashMap<ExprId, DefId>,
     type_of_expr: ArenaMap<ExprId, Ty>,
     type_of_pat: ArenaMap<PatId, Ty>,
     /// The return type of the function being inferred.
@@ -631,6 +641,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
         impl_block: Option<ImplBlock>,
     ) -> Self {
         InferenceContext {
+            method_resolutions: FxHashMap::default(),
             type_of_expr: ArenaMap::default(),
             type_of_pat: ArenaMap::default(),
             var_unification_table: InPlaceUnificationTable::new(),
@@ -655,6 +666,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
             *ty = resolved;
         }
         InferenceResult {
+            method_resolutions: mem::replace(&mut self.method_resolutions, Default::default()),
             type_of_expr: expr_types,
             type_of_pat: pat_types,
         }
@@ -664,6 +676,10 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
         self.type_of_expr.insert(expr, ty);
     }
 
+    fn write_method_resolution(&mut self, expr: ExprId, def_id: DefId) {
+        self.method_resolutions.insert(expr, def_id);
+    }
+
     fn write_pat_ty(&mut self, pat: PatId, ty: Ty) {
         self.type_of_pat.insert(pat, ty);
     }
@@ -900,7 +916,10 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 let receiver_ty = self.infer_expr(*receiver, &Expectation::none())?;
                 let resolved = receiver_ty.clone().lookup_method(self.db, method_name)?;
                 let method_ty = match resolved {
-                    Some(def_id) => self.db.type_for_def(def_id)?,
+                    Some(def_id) => {
+                        self.write_method_resolution(expr, def_id);
+                        self.db.type_for_def(def_id)?
+                    }
                     None => Ty::Unknown,
                 };
                 let method_ty = self.insert_type_vars(method_ty);
diff --git a/crates/ra_ide_api/src/goto_definition.rs b/crates/ra_ide_api/src/goto_definition.rs
index e2537758d..332a2fb8d 100644
--- a/crates/ra_ide_api/src/goto_definition.rs
+++ b/crates/ra_ide_api/src/goto_definition.rs
@@ -47,15 +47,34 @@ pub(crate) fn reference_definition(
     name_ref: &ast::NameRef,
 ) -> Cancelable<ReferenceResult> {
     use self::ReferenceResult::*;
-    if let Some(fn_descr) =
+    if let Some(function) =
         hir::source_binder::function_from_child_node(db, file_id, name_ref.syntax())?
     {
-        let scope = fn_descr.scopes(db)?;
+        let scope = function.scopes(db)?;
         // First try to resolve the symbol locally
         if let Some(entry) = scope.resolve_local_name(name_ref) {
             let nav = NavigationTarget::from_scope_entry(file_id, &entry);
             return Ok(Exact(nav));
         };
+
+        // Next check if it is a method
+        if let Some(method_call) = name_ref
+            .syntax()
+            .parent()
+            .and_then(ast::MethodCallExpr::cast)
+        {
+            let infer_result = function.infer(db)?;
+            let syntax_mapping = function.body_syntax_mapping(db)?;
+            let expr = ast::Expr::cast(method_call.syntax()).unwrap();
+            if let Some(def_id) = syntax_mapping
+                .node_expr(expr)
+                .and_then(|it| infer_result.method_resolution(it))
+            {
+                if let Some(target) = NavigationTarget::from_def(db, def_id.resolve(db)?)? {
+                    return Ok(Exact(target));
+                }
+            };
+        }
     }
     // Then try module name resolution
     if let Some(module) =
@@ -167,4 +186,32 @@ mod tests {
             "foo SOURCE_FILE FileId(2) [0; 10)",
         );
     }
+
+    #[test]
+    fn goto_definition_works_for_methods() {
+        check_goto(
+            "
+            //- /lib.rs
+            struct Foo;
+            impl Foo {
+                fn frobnicate(&self) {  }
+            }
+
+            fn bar(foo: &Foo) {
+                foo.frobnicate<|>();
+            }
+            ",
+            "frobnicate FN_DEF FileId(1) [27; 52) [30; 40)",
+        );
+
+        check_goto(
+            "
+            //- /lib.rs
+            mod <|>foo;
+            //- /foo/mod.rs
+            // empty
+            ",
+            "foo SOURCE_FILE FileId(2) [0; 10)",
+        );
+    }
 }