From b87ee914fa3f74e762fa3dd24bd1c903b9330abe Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Mon, 15 May 2023 20:13:05 +0200
Subject: [PATCH] feat: Highlight used trait assoc items when cursor is on
 trait import or trait bound

---
 crates/ide/src/highlight_related.rs | 104 +++++++++++++++++++++++++++-
 1 file changed, 101 insertions(+), 3 deletions(-)

diff --git a/crates/ide/src/highlight_related.rs b/crates/ide/src/highlight_related.rs
index 3a519fe65a1..3c4a1fe1e74 100644
--- a/crates/ide/src/highlight_related.rs
+++ b/crates/ide/src/highlight_related.rs
@@ -4,7 +4,9 @@ use ide_db::{
     defs::{Definition, IdentClass},
     helpers::pick_best_token,
     search::{FileReference, ReferenceCategory, SearchScope},
-    syntax_helpers::node_ext::{for_each_break_and_continue_expr, for_each_tail_expr, walk_expr},
+    syntax_helpers::node_ext::{
+        for_each_break_and_continue_expr, for_each_tail_expr, full_path_of_name_ref, walk_expr,
+    },
     FxHashSet, RootDatabase,
 };
 use syntax::{
@@ -39,11 +41,13 @@ pub struct HighlightRelatedConfig {
 // Highlights constructs related to the thing under the cursor:
 //
 // . if on an identifier, highlights all references to that identifier in the current file
+// .. additionally, if the identifier is a trait in a where clause, type parameter trait bound or use item, highlights all references to that trait's assoc items in the corresponding scope
 // . if on an `async` or `await token, highlights all yield points for that async context
 // . if on a `return` or `fn` keyword, `?` character or `->` return type arrow, highlights all exit points for that context
 // . if on a `break`, `loop`, `while` or `for` token, highlights all break points for that loop or block context
+// . if on a `move` or `|` token that belongs to a closure, highlights all captures of the closure.
 //
-// Note: `?` and `->` do not currently trigger this behavior in the VSCode editor.
+// Note: `?`, `|` and `->` do not currently trigger this behavior in the VSCode editor.
 pub(crate) fn highlight_related(
     sema: &Semantics<'_, RootDatabase>,
     config: HighlightRelatedConfig,
@@ -129,7 +133,7 @@ fn highlight_references(
     token: SyntaxToken,
     file_id: FileId,
 ) -> Option<Vec<HighlightedRange>> {
-    let defs = find_defs(sema, token);
+    let defs = find_defs(sema, token.clone());
     let usages = defs
         .iter()
         .filter_map(|&d| {
@@ -144,6 +148,59 @@ fn highlight_references(
         .map(|FileReference { category, range, .. }| HighlightedRange { range, category });
     let mut res = FxHashSet::default();
     for &def in &defs {
+        // highlight trait usages
+        if let Definition::Trait(t) = def {
+            let trait_item_use_scope = (|| {
+                let name_ref = token.parent().and_then(ast::NameRef::cast)?;
+                let path = full_path_of_name_ref(&name_ref)?;
+                let parent = path.syntax().parent()?;
+                match_ast! {
+                    match parent {
+                        ast::UseTree(it) => it.syntax().ancestors().find(|it| {
+                            ast::SourceFile::can_cast(it.kind()) || ast::Module::can_cast(it.kind())
+                        }),
+                        ast::PathType(it) => it
+                            .syntax()
+                            .ancestors()
+                            .nth(2)
+                            .and_then(ast::TypeBoundList::cast)?
+                            .syntax()
+                            .parent()
+                            .filter(|it| ast::WhereClause::can_cast(it.kind()) || ast::TypeParam::can_cast(it.kind()))?
+                            .ancestors()
+                            .find(|it| {
+                                ast::Item::can_cast(it.kind())
+                            }),
+                        _ => None,
+                    }
+                }
+            })();
+            if let Some(trait_item_use_scope) = trait_item_use_scope {
+                res.extend(
+                    t.items_with_supertraits(sema.db)
+                        .into_iter()
+                        .filter_map(|item| {
+                            Definition::from(item)
+                                .usages(sema)
+                                .set_scope(Some(SearchScope::file_range(FileRange {
+                                    file_id,
+                                    range: trait_item_use_scope.text_range(),
+                                })))
+                                .include_self_refs()
+                                .all()
+                                .references
+                                .remove(&file_id)
+                        })
+                        .flatten()
+                        .map(|FileReference { category, range, .. }| HighlightedRange {
+                            range,
+                            category,
+                        }),
+                );
+            }
+        }
+
+        // highlight the defs themselves
         match def {
             Definition::Local(local) => {
                 let category = local.is_mut(sema.db).then_some(ReferenceCategory::Write);
@@ -1476,6 +1533,47 @@ fn f() {
     let c = move$0 |y| x + y;
     //               ^ read
 }
+"#,
+        );
+    }
+
+    #[test]
+    fn test_trait_highlights_assoc_item_uses() {
+        check(
+            r#"
+trait Foo {
+    //^^^
+    type T;
+    const C: usize;
+    fn f() {}
+    fn m(&self) {}
+}
+impl Foo for i32 {
+   //^^^
+    type T = i32;
+    const C: usize = 0;
+    fn f() {}
+    fn m(&self) {}
+}
+fn f<T: Foo$0>(t: T) {
+      //^^^
+    let _: T::T;
+            //^
+    t.m();
+    //^
+    T::C;
+     //^
+    T::f();
+     //^
+}
+
+fn f2<T: Foo>(t: T) {
+       //^^^
+    let _: T::T;
+    t.m();
+    T::C;
+    T::f();
+}
 "#,
         );
     }