Auto merge of #13933 - Veykril:assoc-search, r=Veykril

Refine search for const and function assoc items

This changes our searching behavior, before we always associated all usages and definitions of associated items with all implementations of a trait and the trait itself. Now, when searching for references of a an associated trait item, we still do the same and consider all implementations, but when searching for an associated item of an implementation we now only consider the uses of that specific implementations associated item.

This does not affect associated type aliases as we unfortunately are missing information in the IDE layer here still.

7 files changed, 319 insertions, 36 deletions

diff --git a/crates/hir-ty/src/method_resolution.rs b/crates/hir-ty/src/method_resolution.rs
index ae25704f204..64376e10bcc 100644
--- a/crates/hir-ty/src/method_resolution.rs
+++ b/crates/hir-ty/src/method_resolution.rs
@@ -712,17 +712,17 @@ fn lookup_impl_assoc_item_for_trait_ref(
     let table = InferenceTable::new(db, env);
 
     let impl_data = find_matching_impl(impls, table, trait_ref)?;
-    impl_data.items.iter().find_map(|it| match it {
+    impl_data.items.iter().find_map(|&it| match it {
         AssocItemId::FunctionId(f) => {
-            (db.function_data(*f).name == *name).then_some(AssocItemId::FunctionId(*f))
+            (db.function_data(f).name == *name).then_some(AssocItemId::FunctionId(f))
         }
         AssocItemId::ConstId(c) => db
-            .const_data(*c)
+            .const_data(c)
             .name
             .as_ref()
-            .map(|n| *n == *name)
-            .and_then(|result| if result { Some(AssocItemId::ConstId(*c)) } else { None }),
-        _ => None,
+            .map(|n| n == name)
+            .and_then(|result| if result { Some(AssocItemId::ConstId(c)) } else { None }),
+        AssocItemId::TypeAliasId(_) => None,
     })
 }
 
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index a1d2ec02f23..2fc2673bd22 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -2129,7 +2129,7 @@ pub enum AssocItem {
     Const(Const),
     TypeAlias(TypeAlias),
 }
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum AssocItemContainer {
     Trait(Trait),
     Impl(Impl),
diff --git a/crates/hir/src/source_analyzer.rs b/crates/hir/src/source_analyzer.rs
index 059b80bcf13..2354eb2c9cc 100644
--- a/crates/hir/src/source_analyzer.rs
+++ b/crates/hir/src/source_analyzer.rs
@@ -504,7 +504,7 @@ impl SourceAnalyzer {
                         AssocItemId::ConstId(const_id) => {
                             self.resolve_impl_const_or_trait_def(db, const_id, subs).into()
                         }
-                        _ => assoc,
+                        assoc => assoc,
                     };
 
                     return Some(PathResolution::Def(AssocItem::from(assoc).into()));
@@ -517,7 +517,13 @@ impl SourceAnalyzer {
                 prefer_value_ns = true;
             } else if let Some(path_pat) = parent().and_then(ast::PathPat::cast) {
                 let pat_id = self.pat_id(&path_pat.into())?;
-                if let Some((assoc, _)) = infer.assoc_resolutions_for_pat(pat_id) {
+                if let Some((assoc, subs)) = infer.assoc_resolutions_for_pat(pat_id) {
+                    let assoc = match assoc {
+                        AssocItemId::ConstId(const_id) => {
+                            self.resolve_impl_const_or_trait_def(db, const_id, subs).into()
+                        }
+                        assoc => assoc,
+                    };
                     return Some(PathResolution::Def(AssocItem::from(assoc).into()));
                 }
                 if let Some(VariantId::EnumVariantId(variant)) =
diff --git a/crates/ide-db/src/search.rs b/crates/ide-db/src/search.rs
index f6d4ccc3ce5..fd09fdeb0bd 100644
--- a/crates/ide-db/src/search.rs
+++ b/crates/ide-db/src/search.rs
@@ -7,7 +7,9 @@
 use std::{mem, sync::Arc};
 
 use base_db::{FileId, FileRange, SourceDatabase, SourceDatabaseExt};
-use hir::{DefWithBody, HasAttrs, HasSource, InFile, ModuleSource, Semantics, Visibility};
+use hir::{
+    AsAssocItem, DefWithBody, HasAttrs, HasSource, InFile, ModuleSource, Semantics, Visibility,
+};
 use memchr::memmem::Finder;
 use once_cell::unsync::Lazy;
 use parser::SyntaxKind;
@@ -311,15 +313,15 @@ impl Definition {
 
     pub fn usages<'a>(self, sema: &'a Semantics<'_, RootDatabase>) -> FindUsages<'a> {
         FindUsages {
-            local_repr: match self {
-                Definition::Local(local) => Some(local.representative(sema.db)),
-                _ => None,
-            },
             def: self,
-            trait_assoc_def: as_trait_assoc_def(sema.db, self),
+            assoc_item_container: self.as_assoc_item(sema.db).map(|a| a.container(sema.db)),
             sema,
             scope: None,
             include_self_kw_refs: None,
+            local_repr: match self {
+                Definition::Local(local) => Some(local.representative(sema.db)),
+                _ => None,
+            },
             search_self_mod: false,
         }
     }
@@ -328,8 +330,7 @@ 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>,
+    assoc_item_container: Option<hir::AssocItemContainer>,
     sema: &'a Semantics<'a, RootDatabase>,
     scope: Option<SearchScope>,
     include_self_kw_refs: Option<hir::Type>,
@@ -380,7 +381,9 @@ impl<'a> FindUsages<'a> {
         let sema = self.sema;
 
         let search_scope = {
-            let base = self.trait_assoc_def.unwrap_or(self.def).search_scope(sema.db);
+            // FIXME: Is the trait scope needed for trait impl assoc items?
+            let base =
+                as_trait_assoc_def(sema.db, self.def).unwrap_or(self.def).search_scope(sema.db);
             match &self.scope {
                 None => base,
                 Some(scope) => base.intersection(scope),
@@ -651,13 +654,26 @@ impl<'a> FindUsages<'a> {
                 sink(file_id, reference)
             }
             Some(NameRefClass::Definition(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,
-                } =>
+                if self.def == def
+                    // is our def a trait assoc item? then we want to find everything
+                    || matches!(self.assoc_item_container, Some(hir::AssocItemContainer::Trait(_)))
+                        && convert_to_def_in_trait(self.sema.db, def) == self.def =>
+            {
+                let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
+                let reference = FileReference {
+                    range,
+                    name: ast::NameLike::NameRef(name_ref.clone()),
+                    category: ReferenceCategory::new(&def, name_ref),
+                };
+                sink(file_id, reference)
+            }
+            // FIXME: special case type aliases, we can't filter between impl and trait defs here as we lack the substitutions
+            // so we always resolve all assoc type aliases to both their trait def and impl defs
+            Some(NameRefClass::Definition(def))
+                if self.assoc_item_container.is_some()
+                    && matches!(self.def, Definition::TypeAlias(_))
+                    && convert_to_def_in_trait(self.sema.db, def)
+                        == convert_to_def_in_trait(self.sema.db, self.def) =>
             {
                 let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
                 let reference = FileReference {
@@ -748,12 +764,14 @@ impl<'a> FindUsages<'a> {
                 false
             }
             Some(NameClass::Definition(def)) if def != self.def => {
-                // 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
-                ) {
+                // only when looking for trait assoc items, we want to find other assoc items
+                if !matches!(self.assoc_item_container, Some(hir::AssocItemContainer::Trait(_)))
+                    // FIXME: special case type aliases, we can't filter between impl and trait defs here as we lack the substitutions
+                    // so we always resolve all assoc type aliases to both their trait def and impl defs
+                    && !(matches!(self.def, Definition::TypeAlias(_))
+                        && convert_to_def_in_trait(self.sema.db, def)
+                            == convert_to_def_in_trait(self.sema.db, self.def))
+                {
                     return false;
                 }
                 let FileRange { file_id, range } = self.sema.original_range(name.syntax());
diff --git a/crates/ide/src/highlight_related.rs b/crates/ide/src/highlight_related.rs
index 55f8779eed7..c889eb930f3 100644
--- a/crates/ide/src/highlight_related.rs
+++ b/crates/ide/src/highlight_related.rs
@@ -1356,7 +1356,6 @@ fn main() {
             r#"
 trait Trait {
     fn func(self) {}
-     //^^^^
 }
 
 impl Trait for () {
@@ -1376,7 +1375,6 @@ fn main() {
             r#"
 trait Trait {
     fn func(self) {}
-     //^^^^
 }
 
 impl Trait for () {
diff --git a/crates/ide/src/references.rs b/crates/ide/src/references.rs
index 0f758cfa2d3..e49f68c57ba 100644
--- a/crates/ide/src/references.rs
+++ b/crates/ide/src/references.rs
@@ -1636,4 +1636,265 @@ pub fn deri$0ve(_stream: TokenStream) -> TokenStream {}
             "#]],
         );
     }
+
+    #[test]
+    fn assoc_items_trait_def() {
+        check(
+            r#"
+trait Trait {
+    const CONST$0: usize;
+}
+
+impl Trait for () {
+    const CONST: usize = 0;
+}
+
+impl Trait for ((),) {
+    const CONST: usize = 0;
+}
+
+fn f<T: Trait>() {
+    let _ = <()>::CONST;
+
+    let _ = T::CONST;
+}
+"#,
+            expect![[r#"
+                CONST Const FileId(0) 18..37 24..29
+
+                FileId(0) 71..76
+                FileId(0) 125..130
+                FileId(0) 183..188
+                FileId(0) 206..211
+            "#]],
+        );
+        check(
+            r#"
+trait Trait {
+    type TypeAlias$0;
+}
+
+impl Trait for () {
+    type TypeAlias = ();
+}
+
+impl Trait for ((),) {
+    type TypeAlias = ();
+}
+
+fn f<T: Trait>() {
+    let _: <() as Trait>::TypeAlias;
+
+    let _: T::TypeAlias;
+}
+"#,
+            expect![[r#"
+                TypeAlias TypeAlias FileId(0) 18..33 23..32
+
+                FileId(0) 66..75
+                FileId(0) 117..126
+                FileId(0) 181..190
+                FileId(0) 207..216
+            "#]],
+        );
+        check(
+            r#"
+trait Trait {
+    fn function$0() {}
+}
+
+impl Trait for () {
+    fn function() {}
+}
+
+impl Trait for ((),) {
+    fn function() {}
+}
+
+fn f<T: Trait>() {
+    let _ = <()>::function;
+
+    let _ = T::function;
+}
+"#,
+            expect![[r#"
+                function Function FileId(0) 18..34 21..29
+
+                FileId(0) 65..73
+                FileId(0) 112..120
+                FileId(0) 166..174
+                FileId(0) 192..200
+            "#]],
+        );
+    }
+
+    #[test]
+    fn assoc_items_trait_impl_def() {
+        check(
+            r#"
+trait Trait {
+    const CONST: usize;
+}
+
+impl Trait for () {
+    const CONST$0: usize = 0;
+}
+
+impl Trait for ((),) {
+    const CONST: usize = 0;
+}
+
+fn f<T: Trait>() {
+    let _ = <()>::CONST;
+
+    let _ = T::CONST;
+}
+"#,
+            expect![[r#"
+                CONST Const FileId(0) 65..88 71..76
+
+                FileId(0) 183..188
+            "#]],
+        );
+        check(
+            r#"
+trait Trait {
+    type TypeAlias;
+}
+
+impl Trait for () {
+    type TypeAlias$0 = ();
+}
+
+impl Trait for ((),) {
+    type TypeAlias = ();
+}
+
+fn f<T: Trait>() {
+    let _: <() as Trait>::TypeAlias;
+
+    let _: T::TypeAlias;
+}
+"#,
+            expect![[r#"
+                TypeAlias TypeAlias FileId(0) 61..81 66..75
+
+                FileId(0) 23..32
+                FileId(0) 117..126
+                FileId(0) 181..190
+                FileId(0) 207..216
+            "#]],
+        );
+        check(
+            r#"
+trait Trait {
+    fn function() {}
+}
+
+impl Trait for () {
+    fn function$0() {}
+}
+
+impl Trait for ((),) {
+    fn function() {}
+}
+
+fn f<T: Trait>() {
+    let _ = <()>::function;
+
+    let _ = T::function;
+}
+"#,
+            expect![[r#"
+                function Function FileId(0) 62..78 65..73
+
+                FileId(0) 166..174
+            "#]],
+        );
+    }
+
+    #[test]
+    fn assoc_items_ref() {
+        check(
+            r#"
+trait Trait {
+    const CONST: usize;
+}
+
+impl Trait for () {
+    const CONST: usize = 0;
+}
+
+impl Trait for ((),) {
+    const CONST: usize = 0;
+}
+
+fn f<T: Trait>() {
+    let _ = <()>::CONST$0;
+
+    let _ = T::CONST;
+}
+"#,
+            expect![[r#"
+                CONST Const FileId(0) 65..88 71..76
+
+                FileId(0) 183..188
+            "#]],
+        );
+        check(
+            r#"
+trait Trait {
+    type TypeAlias;
+}
+
+impl Trait for () {
+    type TypeAlias = ();
+}
+
+impl Trait for ((),) {
+    type TypeAlias = ();
+}
+
+fn f<T: Trait>() {
+    let _: <() as Trait>::TypeAlias$0;
+
+    let _: T::TypeAlias;
+}
+"#,
+            expect![[r#"
+                TypeAlias TypeAlias FileId(0) 18..33 23..32
+
+                FileId(0) 66..75
+                FileId(0) 117..126
+                FileId(0) 181..190
+                FileId(0) 207..216
+            "#]],
+        );
+        check(
+            r#"
+trait Trait {
+    fn function() {}
+}
+
+impl Trait for () {
+    fn function() {}
+}
+
+impl Trait for ((),) {
+    fn function() {}
+}
+
+fn f<T: Trait>() {
+    let _ = <()>::function$0;
+
+    let _ = T::function;
+}
+"#,
+            expect![[r#"
+                function Function FileId(0) 62..78 65..73
+
+                FileId(0) 166..174
+            "#]],
+        );
+    }
 }
diff --git a/crates/rust-analyzer/src/config.rs b/crates/rust-analyzer/src/config.rs
index b0afbdc9a42..6d6e367e374 100644
--- a/crates/rust-analyzer/src/config.rs
+++ b/crates/rust-analyzer/src/config.rs
@@ -1044,7 +1044,7 @@ impl Config {
         &self.data.cargo_extraEnv
     }
 
-    pub fn check_on_save_extra_env(&self) -> FxHashMap<String, String> {
+    pub fn check_extra_env(&self) -> FxHashMap<String, String> {
         let mut extra_env = self.data.cargo_extraEnv.clone();
         extra_env.extend(self.data.check_extraEnv.clone());
         extra_env
@@ -1165,7 +1165,7 @@ impl Config {
                 FlycheckConfig::CustomCommand {
                     command,
                     args,
-                    extra_env: self.check_on_save_extra_env(),
+                    extra_env: self.check_extra_env(),
                     invocation_strategy: match self.data.check_invocationStrategy {
                         InvocationStrategy::Once => flycheck::InvocationStrategy::Once,
                         InvocationStrategy::PerWorkspace => {
@@ -1210,7 +1210,7 @@ impl Config {
                     CargoFeaturesDef::Selected(it) => it,
                 },
                 extra_args: self.data.check_extraArgs.clone(),
-                extra_env: self.check_on_save_extra_env(),
+                extra_env: self.check_extra_env(),
                 ansi_color_output: self.color_diagnostic_output(),
             },
         }