Auto merge of rust-lang#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.

Author: bors

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

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

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)) =

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());

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 () {