Highlight all related tokens in macro inputs

Author: Veykril

crates/hir/src/semantics.rs

@@ -166,10 +166,12 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.speculative_expand(actual_macro_call, speculative_args, token_to_map)
     }
 
+    // FIXME: Rename to descend_into_macros_single
     pub fn descend_into_macros(&self, token: SyntaxToken) -> SyntaxToken {
         self.imp.descend_into_macros(token).pop().unwrap()
     }
 
+    // FIXME: Rename to descend_into_macros
     pub fn descend_into_macros_many(&self, token: SyntaxToken) -> SmallVec<[SyntaxToken; 1]> {
         self.imp.descend_into_macros(token)
     }
@@ -236,6 +238,16 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.descend_node_at_offset(node, offset).find_map(N::cast)
     }
 
+    /// Find an AstNode by offset inside SyntaxNode, if it is inside *MacroCall*,
+    /// descend it and find again
+    pub fn find_nodes_at_offset_with_descend<'slf, N: AstNode + 'slf>(
+        &'slf self,
+        node: &SyntaxNode,
+        offset: TextSize,
+    ) -> impl Iterator<Item = N> + 'slf {
+        self.imp.descend_node_at_offset(node, offset).flat_map(N::cast)
+    }
+
     pub fn resolve_lifetime_param(&self, lifetime: &ast::Lifetime) -> Option<LifetimeParam> {
         self.imp.resolve_lifetime_param(lifetime)
     }
@@ -482,12 +494,13 @@ impl<'db> SemanticsImpl<'db> {
                                     .as_ref()?
                                     .map_token_down(self.db.upcast(), None, token.as_ref())?;
 
+                                let len = queue.len();
                                 queue.extend(tokens.inspect(|token| {
                                     if let Some(parent) = token.value.parent() {
                                         self.cache(find_root(&parent), token.file_id);
                                     }
                                 }));
-                                return Some(());
+                                return (queue.len() != len).then(|| ());
                             },
                             ast::Item(item) => {
                                 match self.with_ctx(|ctx| ctx.item_to_macro_call(token.with_value(item))) {
@@ -500,12 +513,13 @@ impl<'db> SemanticsImpl<'db> {
                                             .as_ref()?
                                             .map_token_down(self.db.upcast(), None, token.as_ref())?;
 
+                                        let len = queue.len();
                                         queue.extend(tokens.inspect(|token| {
                                             if let Some(parent) = token.value.parent() {
                                                 self.cache(find_root(&parent), token.file_id);
                                             }
                                         }));
-                                        return Some(());
+                                        return (queue.len() != len).then(|| ());
                                     }
                                     None => {}
                                 }

crates/ide/src/highlight_related.rs

@@ -6,6 +6,7 @@ use ide_db::{
     search::{FileReference, ReferenceAccess, SearchScope},
     RootDatabase,
 };
+use itertools::Itertools;
 use syntax::{
     ast::{self, LoopBodyOwner},
     match_ast, AstNode, SyntaxNode, SyntaxToken, TextRange, TextSize, T,
@@ -70,7 +71,7 @@ fn highlight_references(
     syntax: &SyntaxNode,
     FilePosition { offset, file_id }: FilePosition,
 ) -> Option<Vec<HighlightedRange>> {
-    let defs = find_defs(sema, syntax, offset)?;
+    let defs = find_defs(sema, syntax, offset);
     let usages = defs
         .iter()
         .flat_map(|&d| {
@@ -99,7 +100,12 @@ fn highlight_references(
         })
     });
 
-    Some(declarations.chain(usages).collect())
+    let res: Vec<_> = declarations.chain(usages).collect();
+    if res.is_empty() {
+        None
+    } else {
+        Some(res)
+    }
 }
 
 fn highlight_exit_points(
@@ -270,29 +276,41 @@ fn find_defs(
     sema: &Semantics<RootDatabase>,
     syntax: &SyntaxNode,
     offset: TextSize,
-) -> Option<Vec<Definition>> {
-    let defs = match sema.find_node_at_offset_with_descend(syntax, offset)? {
-        ast::NameLike::NameRef(name_ref) => match NameRefClass::classify(sema, &name_ref)? {
-            NameRefClass::Definition(def) => vec![def],
-            NameRefClass::FieldShorthand { local_ref, field_ref } => {
-                vec![Definition::Local(local_ref), Definition::Field(field_ref)]
-            }
-        },
-        ast::NameLike::Name(name) => match NameClass::classify(sema, &name)? {
-            NameClass::Definition(it) | NameClass::ConstReference(it) => vec![it],
-            NameClass::PatFieldShorthand { local_def, field_ref } => {
-                vec![Definition::Local(local_def), Definition::Field(field_ref)]
-            }
-        },
-        ast::NameLike::Lifetime(lifetime) => NameRefClass::classify_lifetime(sema, &lifetime)
-            .and_then(|class| match class {
-                NameRefClass::Definition(it) => Some(it),
-                _ => None,
+) -> Vec<Definition> {
+    sema.find_nodes_at_offset_with_descend(syntax, offset)
+        .flat_map(|name_like| {
+            Some(match name_like {
+                ast::NameLike::NameRef(name_ref) => {
+                    match NameRefClass::classify(sema, &name_ref)? {
+                        NameRefClass::Definition(def) => vec![def],
+                        NameRefClass::FieldShorthand { local_ref, field_ref } => {
+                            vec![Definition::Local(local_ref), Definition::Field(field_ref)]
+                        }
+                    }
+                }
+                ast::NameLike::Name(name) => match NameClass::classify(sema, &name)? {
+                    NameClass::Definition(it) | NameClass::ConstReference(it) => vec![it],
+                    NameClass::PatFieldShorthand { local_def, field_ref } => {
+                        vec![Definition::Local(local_def), Definition::Field(field_ref)]
+                    }
+                },
+                ast::NameLike::Lifetime(lifetime) => {
+                    NameRefClass::classify_lifetime(sema, &lifetime)
+                        .and_then(|class| match class {
+                            NameRefClass::Definition(it) => Some(it),
+                            _ => None,
+                        })
+                        .or_else(|| {
+                            NameClass::classify_lifetime(sema, &lifetime)
+                                .and_then(NameClass::defined)
+                        })
+                        .map(|it| vec![it])?
+                }
             })
-            .or_else(|| NameClass::classify_lifetime(sema, &lifetime).and_then(NameClass::defined))
-            .map(|it| vec![it])?,
-    };
-    Some(defs)
+        })
+        .flatten()
+        .unique()
+        .collect()
 }
 
 #[cfg(test)]
@@ -392,6 +410,46 @@ fn foo() {
         );
     }
 
+    #[test]
+    fn test_multi_macro_usage() {
+        check(
+            r#"
+macro_rules! foo {
+    ($ident:ident) => {
+        fn $ident() -> $ident { loop {} }
+        struct $ident;
+    }
+}
+
+foo!(bar$0);
+  // ^^^
+  // ^^^
+fn foo() {
+    let bar: bar = bar();
+          // ^^^
+                // ^^^
+}
+"#,
+        );
+        check(
+            r#"
+macro_rules! foo {
+    ($ident:ident) => {
+        fn $ident() -> $ident { loop {} }
+        struct $ident;
+    }
+}
+
+foo!(bar);
+  // ^^^
+fn foo() {
+    let bar: bar$0 = bar();
+          // ^^^
+}
+"#,
+        );
+    }
+
     #[test]
     fn test_hl_yield_points() {
         check(

crates/ide_db/src/defs.rs

@@ -17,7 +17,7 @@ use syntax::{
 use crate::RootDatabase;
 
 // FIXME: a more precise name would probably be `Symbol`?
-#[derive(Debug, PartialEq, Eq, Copy, Clone)]
+#[derive(Debug, PartialEq, Eq, Copy, Clone, Hash)]
 pub enum Definition {
     Macro(MacroDef),
     Field(Field),