Item search now respects trait impl items

1 files changed, 37 insertions, 10 deletions

diff --git a/crates/ide_db/src/search.rs b/crates/ide_db/src/search.rs
index 67840602b..8152630f5 100644
--- a/crates/ide_db/src/search.rs
+++ b/crates/ide_db/src/search.rs
@@ -8,7 +8,8 @@ use std::{convert::TryInto, mem};
 
 use base_db::{FileId, FileRange, SourceDatabase, SourceDatabaseExt};
 use hir::{
-    DefWithBody, HasAttrs, HasSource, InFile, ModuleDef, ModuleSource, Semantics, Visibility,
+    AsAssocItem, DefWithBody, HasAttrs, HasSource, InFile, ModuleDef, ModuleSource, Semantics,
+    Visibility,
 };
 use once_cell::unsync::Lazy;
 use rustc_hash::FxHashMap;
@@ -303,13 +304,13 @@ impl Definition {
         }
     }
 
-    pub fn usages<'a>(&'a self, sema: &'a Semantics<RootDatabase>) -> FindUsages<'a> {
+    pub fn usages<'a>(self, sema: &'a Semantics<RootDatabase>) -> FindUsages<'a> {
         FindUsages { def: self, sema, scope: None, include_self_kw_refs: None }
     }
 }
 
 pub struct FindUsages<'a> {
-    def: &'a Definition,
+    def: Definition,
     sema: &'a Semantics<'a, RootDatabase>,
     scope: Option<SearchScope>,
     include_self_kw_refs: Option<hir::Type>,
@@ -318,7 +319,7 @@ pub struct FindUsages<'a> {
 impl<'a> FindUsages<'a> {
     /// Enable searching for `Self` when the definition is a type.
     pub fn include_self_refs(mut self) -> FindUsages<'a> {
-        self.include_self_kw_refs = def_to_ty(self.sema, self.def);
+        self.include_self_kw_refs = def_to_ty(self.sema, &self.def);
         self
     }
 
@@ -445,7 +446,7 @@ impl<'a> FindUsages<'a> {
         sink: &mut dyn FnMut(FileId, FileReference) -> bool,
     ) -> bool {
         match NameRefClass::classify_lifetime(self.sema, lifetime) {
-            Some(NameRefClass::Definition(def)) if &def == self.def => {
+            Some(NameRefClass::Definition(def)) if def == self.def => {
                 let FileRange { file_id, range } = self.sema.original_range(lifetime.syntax());
                 let reference = FileReference {
                     range,
@@ -464,7 +465,7 @@ impl<'a> FindUsages<'a> {
         sink: &mut dyn FnMut(FileId, FileReference) -> bool,
     ) -> bool {
         match NameRefClass::classify(self.sema, &name_ref) {
-            Some(NameRefClass::Definition(def)) if &def == self.def => {
+            Some(NameRefClass::Definition(def)) if def == self.def => {
                 let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
                 let reference = FileReference {
                     range,
@@ -489,10 +490,10 @@ impl<'a> FindUsages<'a> {
             Some(NameRefClass::FieldShorthand { local_ref: local, field_ref: field }) => {
                 let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
                 let access = match self.def {
-                    Definition::Field(_) if &field == self.def => {
+                    Definition::Field(_) if field == self.def => {
                         reference_access(&field, &name_ref)
                     }
-                    Definition::Local(l) if &local == l => {
+                    Definition::Local(l) if local == l => {
                         reference_access(&Definition::Local(local), &name_ref)
                     }
                     _ => return false,
@@ -513,7 +514,7 @@ impl<'a> FindUsages<'a> {
         match NameClass::classify(self.sema, name) {
             Some(NameClass::PatFieldShorthand { local_def: _, field_ref })
                 if matches!(
-                    self.def, Definition::Field(_) if &field_ref == self.def
+                    self.def, Definition::Field(_) if field_ref == self.def
                 ) =>
             {
                 let FileRange { file_id, range } = self.sema.original_range(name.syntax());
@@ -525,12 +526,38 @@ impl<'a> FindUsages<'a> {
                 };
                 sink(file_id, reference)
             }
-            Some(NameClass::ConstReference(def)) if *self.def == def => {
+            Some(NameClass::ConstReference(def)) if self.def == def => {
                 let FileRange { file_id, range } = self.sema.original_range(name.syntax());
                 let reference =
                     FileReference { range, name: ast::NameLike::Name(name.clone()), access: None };
                 sink(file_id, reference)
             }
+            // Resolve trait impl function definitions to the trait definition's version if self.def is the trait definition's
+            Some(NameClass::Definition(Definition::ModuleDef(mod_def))) => {
+                /* poor man's try block */
+                (|| {
+                    let this = match self.def {
+                        Definition::ModuleDef(this) if this != mod_def => this,
+                        _ => return None,
+                    };
+                    let this_trait = this
+                        .as_assoc_item(self.sema.db)?
+                        .containing_trait_or_trait_impl(self.sema.db)?;
+                    let trait_ = mod_def
+                        .as_assoc_item(self.sema.db)?
+                        .containing_trait_or_trait_impl(self.sema.db)?;
+                    (trait_ == this_trait).then(|| {
+                        let FileRange { file_id, range } = self.sema.original_range(name.syntax());
+                        let reference = FileReference {
+                            range,
+                            name: ast::NameLike::Name(name.clone()),
+                            access: None,
+                        };
+                        sink(file_id, reference)
+                    })
+                })()
+                .unwrap_or(false)
+            }
             _ => false,
         }
     }