Attempt to track attr macros during highlighting

2 files changed, 40 insertions, 0 deletions

diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index 920e18208..2d08a7704 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -123,6 +123,10 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.expand_attr_macro(item)
     }
 
+    pub fn is_attr_macro_call(&self, item: &ast::Item) -> bool {
+        self.imp.is_attr_macro_call(item)
+    }
+
     pub fn speculative_expand(
         &self,
         actual_macro_call: &ast::MacroCall,
@@ -348,6 +352,12 @@ impl<'db> SemanticsImpl<'db> {
         Some(node)
     }
 
+    fn is_attr_macro_call(&self, item: &ast::Item) -> bool {
+        let sa = self.analyze(item.syntax());
+        let src = InFile::new(sa.file_id, item.clone());
+        self.with_ctx(|ctx| ctx.item_to_macro_call(src).is_some())
+    }
+
     fn speculative_expand(
         &self,
         actual_macro_call: &ast::MacroCall,
diff --git a/crates/ide/src/syntax_highlighting.rs b/crates/ide/src/syntax_highlighting.rs
index 79c2f4a1e..b03f1c71f 100644
--- a/crates/ide/src/syntax_highlighting.rs
+++ b/crates/ide/src/syntax_highlighting.rs
@@ -192,6 +192,7 @@ fn traverse(
     let mut bindings_shadow_count: FxHashMap<Name, u32> = FxHashMap::default();
 
     let mut current_macro_call: Option<ast::MacroCall> = None;
+    let mut current_attr_macro_call = None;
     let mut current_macro: Option<ast::Macro> = None;
     let mut macro_highlighter = MacroHighlighter::default();
     let mut inside_attribute = false;
@@ -227,6 +228,19 @@ fn traverse(
             }
             _ => (),
         }
+        match event.clone().map(|it| it.into_node().and_then(ast::Item::cast)) {
+            WalkEvent::Enter(Some(item)) => {
+                if sema.is_attr_macro_call(&item) {
+                    current_attr_macro_call = Some(item);
+                }
+            }
+            WalkEvent::Leave(Some(item)) => {
+                if current_attr_macro_call == Some(item) {
+                    current_attr_macro_call = None;
+                }
+            }
+            _ => (),
+        }
 
         match event.clone().map(|it| it.into_node().and_then(ast::Macro::cast)) {
             WalkEvent::Enter(Some(mac)) => {
@@ -286,6 +300,22 @@ fn traverse(
                 }
                 None => token.into(),
             }
+        } else if current_attr_macro_call.is_some() {
+            let token = match element.clone().into_token() {
+                Some(it) => it,
+                _ => continue,
+            };
+            let token = sema.descend_into_macros(token.clone());
+            match token.parent() {
+                Some(parent) => {
+                    // We only care Name and Name_ref
+                    match (token.kind(), parent.kind()) {
+                        (IDENT, NAME) | (IDENT, NAME_REF) => parent.into(),
+                        _ => token.into(),
+                    }
+                }
+                None => token.into(),
+            }
         } else {
             element.clone()
         };