Auto merge of #12844 - Veykril:highlight-attr, r=Veykril

fix: Improve syntax highlighting in attributes

Fixes https://github.com/rust-lang/rust-analyzer/issues/12842

2 files changed, 64 insertions, 7 deletions

diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index edcb2fc6a72..043f2b7c24d 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -22,7 +22,7 @@ use smallvec::{smallvec, SmallVec};
 use syntax::{
     algo::skip_trivia_token,
     ast::{self, HasAttrs as _, HasGenericParams, HasLoopBody},
-    match_ast, AstNode, Direction, SyntaxNode, SyntaxNodePtr, SyntaxToken, TextSize,
+    match_ast, AstNode, Direction, SyntaxKind, SyntaxNode, SyntaxNodePtr, SyntaxToken, TextSize,
 };
 
 use crate::{
@@ -217,6 +217,10 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.descend_into_macros_with_same_text(token)
     }
 
+    pub fn descend_into_macros_with_kind_preference(&self, token: SyntaxToken) -> SyntaxToken {
+        self.imp.descend_into_macros_with_kind_preference(token)
+    }
+
     /// Maps a node down by mapping its first and last token down.
     pub fn descend_node_into_attributes<N: AstNode>(&self, node: N) -> SmallVec<[N; 1]> {
         self.imp.descend_node_into_attributes(node)
@@ -680,6 +684,32 @@ impl<'db> SemanticsImpl<'db> {
         res
     }
 
+    fn descend_into_macros_with_kind_preference(&self, token: SyntaxToken) -> SyntaxToken {
+        let fetch_kind = |token: &SyntaxToken| match token.parent() {
+            Some(node) => match node.kind() {
+                kind @ (SyntaxKind::NAME | SyntaxKind::NAME_REF) => {
+                    node.parent().map_or(kind, |it| it.kind())
+                }
+                _ => token.kind(),
+            },
+            None => token.kind(),
+        };
+        let preferred_kind = fetch_kind(&token);
+        let mut res = None;
+        self.descend_into_macros_impl(token.clone(), &mut |InFile { value, .. }| {
+            if fetch_kind(&value) == preferred_kind {
+                res = Some(value);
+                true
+            } else {
+                if let None = res {
+                    res = Some(value)
+                }
+                false
+            }
+        });
+        res.unwrap_or(token)
+    }
+
     fn descend_into_macros_single(&self, token: SyntaxToken) -> SyntaxToken {
         let mut res = token.clone();
         self.descend_into_macros_impl(token, &mut |InFile { value, .. }| {
diff --git a/crates/ide/src/syntax_highlighting.rs b/crates/ide/src/syntax_highlighting.rs
index 9fb6a302632..d7ad6a75799 100644
--- a/crates/ide/src/syntax_highlighting.rs
+++ b/crates/ide/src/syntax_highlighting.rs
@@ -206,6 +206,19 @@ fn traverse(
     let is_unlinked = sema.to_module_def(file_id).is_none();
     let mut bindings_shadow_count: FxHashMap<Name, u32> = FxHashMap::default();
 
+    enum AttrOrDerive {
+        Attr(ast::Item),
+        Derive(ast::Item),
+    }
+
+    impl AttrOrDerive {
+        fn item(&self) -> &ast::Item {
+            match self {
+                AttrOrDerive::Attr(item) | AttrOrDerive::Derive(item) => item,
+            }
+        }
+    }
+
     let mut tt_level = 0;
     let mut attr_or_derive_item = None;
     let mut current_macro: Option<ast::Macro> = None;
@@ -260,7 +273,7 @@ fn traverse(
 
                         if attr_or_derive_item.is_none() {
                             if sema.is_attr_macro_call(&item) {
-                                attr_or_derive_item = Some(item);
+                                attr_or_derive_item = Some(AttrOrDerive::Attr(item));
                             } else {
                                 let adt = match item {
                                     ast::Item::Enum(it) => Some(ast::Adt::Enum(it)),
@@ -270,7 +283,8 @@ fn traverse(
                                 };
                                 match adt {
                                     Some(adt) if sema.is_derive_annotated(&adt) => {
-                                        attr_or_derive_item = Some(ast::Item::from(adt));
+                                        attr_or_derive_item =
+                                            Some(AttrOrDerive::Derive(ast::Item::from(adt)));
                                     }
                                     _ => (),
                                 }
@@ -292,7 +306,9 @@ fn traverse(
                         current_macro = None;
                         macro_highlighter = MacroHighlighter::default();
                     }
-                    Some(item) if attr_or_derive_item.as_ref().map_or(false, |it| *it == item) => {
+                    Some(item)
+                        if attr_or_derive_item.as_ref().map_or(false, |it| *it.item() == item) =>
+                    {
                         attr_or_derive_item = None;
                     }
                     _ => (),
@@ -330,15 +346,26 @@ fn traverse(
 
         // Descending tokens into macros is expensive even if no descending occurs, so make sure
         // that we actually are in a position where descending is possible.
-        let in_macro = tt_level > 0 || attr_or_derive_item.is_some();
+        let in_macro = tt_level > 0
+            || match attr_or_derive_item {
+                Some(AttrOrDerive::Attr(_)) => true,
+                Some(AttrOrDerive::Derive(_)) => inside_attribute,
+                None => false,
+            };
         let descended_element = if in_macro {
             // Attempt to descend tokens into macro-calls.
             match element {
                 NodeOrToken::Token(token) if token.kind() != COMMENT => {
-                    let token = sema.descend_into_macros_single(token);
+                    let token = match attr_or_derive_item {
+                        Some(AttrOrDerive::Attr(_)) => {
+                            sema.descend_into_macros_with_kind_preference(token)
+                        }
+                        Some(AttrOrDerive::Derive(_)) | None => {
+                            sema.descend_into_macros_single(token)
+                        }
+                    };
                     match token.parent().and_then(ast::NameLike::cast) {
                         // Remap the token into the wrapping single token nodes
-                        // FIXME: if the node doesn't resolve, we also won't do token based highlighting!
                         Some(parent) => match (token.kind(), parent.syntax().kind()) {
                             (T![self] | T![ident], NAME | NAME_REF) => NodeOrToken::Node(parent),
                             (T![self] | T![super] | T![crate] | T![Self], NAME_REF) => {