Refine search for const and function assoc items

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(),
             },
         }