fix: Fix search for associated trait items being inconsistent

4 files changed, 125 insertions, 43 deletions

diff --git a/crates/ide-db/src/search.rs b/crates/ide-db/src/search.rs
index 692fce06b0f..f02e0963505 100644
--- a/crates/ide-db/src/search.rs
+++ b/crates/ide-db/src/search.rs
@@ -7,16 +7,14 @@
 use std::{convert::TryInto, mem, sync::Arc};
 
 use base_db::{FileId, FileRange, SourceDatabase, SourceDatabaseExt};
-use hir::{
-    AsAssocItem, DefWithBody, HasAttrs, HasSource, InFile, ModuleSource, Semantics, Visibility,
-};
+use hir::{DefWithBody, HasAttrs, HasSource, InFile, ModuleSource, Semantics, Visibility};
 use once_cell::unsync::Lazy;
 use rustc_hash::FxHashMap;
 use syntax::{ast, match_ast, AstNode, TextRange, TextSize};
 
 use crate::{
     defs::{Definition, NameClass, NameRefClass},
-    traits::convert_to_def_in_trait,
+    traits::{as_trait_assoc_def, convert_to_def_in_trait},
     RootDatabase,
 };
 
@@ -314,6 +312,7 @@ impl Definition {
                 _ => None,
             },
             def: self,
+            trait_assoc_def: as_trait_assoc_def(sema.db, self),
             sema,
             scope: None,
             include_self_kw_refs: None,
@@ -325,6 +324,8 @@ impl Definition {
 #[derive(Clone)]
 pub struct FindUsages<'a> {
     def: Definition,
+    /// If def is an assoc item from a trait or trait impl, this is the corresponding item of the trait definition
+    trait_assoc_def: Option<Definition>,
     sema: &'a Semantics<'a, RootDatabase>,
     scope: Option<SearchScope>,
     include_self_kw_refs: Option<hir::Type>,
@@ -375,7 +376,7 @@ impl<'a> FindUsages<'a> {
         let sema = self.sema;
 
         let search_scope = {
-            let base = self.def.search_scope(sema.db);
+            let base = self.trait_assoc_def.unwrap_or(self.def).search_scope(sema.db);
             match &self.scope {
                 None => base,
                 Some(scope) => base.intersection(scope),
@@ -621,7 +622,13 @@ impl<'a> FindUsages<'a> {
                 sink(file_id, reference)
             }
             Some(NameRefClass::Definition(def))
-                if convert_to_def_in_trait(self.sema.db, def) == self.def =>
+                if match self.trait_assoc_def {
+                    Some(trait_assoc_def) => {
+                        // we have a trait assoc item, so force resolve all assoc items to their trait version
+                        convert_to_def_in_trait(self.sema.db, def) == trait_assoc_def
+                    }
+                    None => self.def == def,
+                } =>
             {
                 let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
                 let reference = FileReference {
@@ -711,30 +718,22 @@ impl<'a> FindUsages<'a> {
                 }
                 false
             }
-            // Resolve trait impl function definitions to the trait definition's version if self.def is the trait definition's
             Some(NameClass::Definition(def)) if def != self.def => {
-                /* poor man's try block */
-                (|| {
-                    let this_trait = self
-                        .def
-                        .as_assoc_item(self.sema.db)?
-                        .containing_trait_or_trait_impl(self.sema.db)?;
-                    let trait_ = def
-                        .as_assoc_item(self.sema.db)?
-                        .containing_trait_or_trait_impl(self.sema.db)?;
-                    (trait_ == this_trait && self.def.name(self.sema.db) == def.name(self.sema.db))
-                        .then(|| {
-                            let FileRange { file_id, range } =
-                                self.sema.original_range(name.syntax());
-                            let reference = FileReference {
-                                range,
-                                name: ast::NameLike::Name(name.clone()),
-                                category: None,
-                            };
-                            sink(file_id, reference)
-                        })
-                })()
-                .unwrap_or(false)
+                // if the def we are looking for is a trait (impl) assoc item, we'll have to resolve the items to trait definition assoc item
+                if !matches!(
+                    self.trait_assoc_def,
+                    Some(trait_assoc_def)
+                        if convert_to_def_in_trait(self.sema.db, def) == trait_assoc_def
+                ) {
+                    return false;
+                }
+                let FileRange { file_id, range } = self.sema.original_range(name.syntax());
+                let reference = FileReference {
+                    range,
+                    name: ast::NameLike::Name(name.clone()),
+                    category: None,
+                };
+                sink(file_id, reference)
             }
             _ => false,
         }
diff --git a/crates/ide-db/src/traits.rs b/crates/ide-db/src/traits.rs
index 666499ed7a5..b607cdfee38 100644
--- a/crates/ide-db/src/traits.rs
+++ b/crates/ide-db/src/traits.rs
@@ -71,26 +71,44 @@ pub fn get_missing_assoc_items(
 
 /// Converts associated trait impl items to their trait definition counterpart
 pub(crate) fn convert_to_def_in_trait(db: &dyn HirDatabase, def: Definition) -> Definition {
-    use hir::AssocItem::*;
     (|| {
         let assoc = def.as_assoc_item(db)?;
         let trait_ = assoc.containing_trait_impl(db)?;
-        let name = match assoc {
-            Function(it) => it.name(db),
-            Const(it) => it.name(db)?,
-            TypeAlias(it) => it.name(db),
-        };
-        let item = trait_.items(db).into_iter().find(|it| match (it, assoc) {
-            (Function(trait_func), Function(_)) => trait_func.name(db) == name,
-            (Const(trait_konst), Const(_)) => trait_konst.name(db).map_or(false, |it| it == name),
-            (TypeAlias(trait_type_alias), TypeAlias(_)) => trait_type_alias.name(db) == name,
-            _ => false,
-        })?;
-        Some(Definition::from(item))
+        assoc_item_of_trait(db, assoc, trait_)
     })()
     .unwrap_or(def)
 }
 
+/// If this is an trait (impl) assoc item, returns the assoc item of the corresponding trait definition.
+pub(crate) fn as_trait_assoc_def(db: &dyn HirDatabase, def: Definition) -> Option<Definition> {
+    let assoc = def.as_assoc_item(db)?;
+    let trait_ = match assoc.container(db) {
+        hir::AssocItemContainer::Trait(_) => return Some(def),
+        hir::AssocItemContainer::Impl(i) => i.trait_(db),
+    }?;
+    assoc_item_of_trait(db, assoc, trait_)
+}
+
+fn assoc_item_of_trait(
+    db: &dyn HirDatabase,
+    assoc: hir::AssocItem,
+    trait_: hir::Trait,
+) -> Option<Definition> {
+    use hir::AssocItem::*;
+    let name = match assoc {
+        Function(it) => it.name(db),
+        Const(it) => it.name(db)?,
+        TypeAlias(it) => it.name(db),
+    };
+    let item = trait_.items(db).into_iter().find(|it| match (it, assoc) {
+        (Function(trait_func), Function(_)) => trait_func.name(db) == name,
+        (Const(trait_konst), Const(_)) => trait_konst.name(db).map_or(false, |it| it == name),
+        (TypeAlias(trait_type_alias), TypeAlias(_)) => trait_type_alias.name(db) == name,
+        _ => false,
+    })?;
+    Some(Definition::from(item))
+}
+
 #[cfg(test)]
 mod tests {
     use base_db::{fixture::ChangeFixture, FilePosition};
diff --git a/crates/ide/src/highlight_related.rs b/crates/ide/src/highlight_related.rs
index 7c0a7fa1e72..9b4bfc41075 100644
--- a/crates/ide/src/highlight_related.rs
+++ b/crates/ide/src/highlight_related.rs
@@ -1310,4 +1310,68 @@ fn foo((
 "#,
         );
     }
+
+    #[test]
+    fn test_hl_trait_impl_methods() {
+        check(
+            r#"
+trait Trait {
+    fn func$0(self) {}
+     //^^^^
+}
+
+impl Trait for () {
+    fn func(self) {}
+     //^^^^
+}
+
+fn main() {
+    <()>::func(());
+        //^^^^
+    ().func();
+     //^^^^
+}
+"#,
+        );
+        check(
+            r#"
+trait Trait {
+    fn func(self) {}
+     //^^^^
+}
+
+impl Trait for () {
+    fn func$0(self) {}
+     //^^^^
+}
+
+fn main() {
+    <()>::func(());
+        //^^^^
+    ().func();
+     //^^^^
+}
+"#,
+        );
+        check(
+            r#"
+trait Trait {
+    fn func(self) {}
+     //^^^^
+}
+
+impl Trait for () {
+    fn func(self) {}
+     //^^^^
+}
+
+fn main() {
+    <()>::func(());
+        //^^^^
+    ().func$0();
+     //^^^^
+}
+"#,
+        );
+    }
 }
diff --git a/crates/ide/src/references.rs b/crates/ide/src/references.rs
index 626e8fe34a9..8a00d6f1441 100644
--- a/crates/ide/src/references.rs
+++ b/crates/ide/src/references.rs
@@ -73,6 +73,7 @@ pub(crate) fn find_all_refs(
             });
             let mut usages =
                 def.usages(sema).set_scope(search_scope.clone()).include_self_refs().all();
+
             if literal_search {
                 retain_adt_literal_usages(&mut usages, def, sema);
             }
@@ -105,7 +106,7 @@ pub(crate) fn find_all_refs(
         }
         None => {
             let search = make_searcher(false);
-            Some(find_defs(sema, &syntax, position.offset)?.into_iter().map(search).collect())
+            Some(find_defs(sema, &syntax, position.offset)?.map(search).collect())
         }
     }
 }