refactor: improve macro handling in navigation for control-flow kws

Author: roife

crates/ide/src/goto_definition.rs

@@ -416,8 +416,10 @@ fn nav_for_branches(
             .descend_into_macros(token.clone())
             .into_iter()
             .filter_map(|token| {
-                let match_expr =
-                    sema.token_ancestors_with_macros(token).find_map(ast::MatchExpr::cast)?;
+                let match_expr = sema
+                    .token_ancestors_with_macros(token)
+                    .take_while(|node| !ast::MacroCall::can_cast(node.kind()))
+                    .find_map(ast::MatchExpr::cast)?;
                 let file_id = sema.hir_file_for(match_expr.syntax());
                 let focus_range = match_expr.match_token()?.text_range();
                 let match_expr_in_file = InFile::new(file_id, match_expr.into());
@@ -430,8 +432,10 @@ fn nav_for_branches(
             .descend_into_macros(token.clone())
             .into_iter()
             .filter_map(|token| {
-                let match_arm =
-                    sema.token_ancestors_with_macros(token).find_map(ast::MatchArm::cast)?;
+                let match_arm = sema
+                    .token_ancestors_with_macros(token)
+                    .take_while(|node| !ast::MacroCall::can_cast(node.kind()))
+                    .find_map(ast::MatchArm::cast)?;
                 let match_expr = sema
                     .ancestors_with_macros(match_arm.syntax().clone())
                     .find_map(ast::MatchExpr::cast)?;
@@ -447,8 +451,10 @@ fn nav_for_branches(
             .descend_into_macros(token.clone())
             .into_iter()
             .filter_map(|token| {
-                let if_expr =
-                    sema.token_ancestors_with_macros(token).find_map(ast::IfExpr::cast)?;
+                let if_expr = sema
+                    .token_ancestors_with_macros(token)
+                    .take_while(|node| !ast::MacroCall::can_cast(node.kind()))
+                    .find_map(ast::IfExpr::cast)?;
                 let file_id = sema.hir_file_for(if_expr.syntax());
                 let focus_range = if_expr.if_token()?.text_range();
                 let if_expr_in_file = InFile::new(file_id, if_expr.into());

crates/ide/src/highlight_related.rs

@@ -318,8 +318,10 @@ pub(crate) fn highlight_branches(
     match token.kind() {
         T![match] => {
             for token in sema.descend_into_macros(token.clone()) {
-                let Some(match_expr) =
-                    sema.token_ancestors_with_macros(token).find_map(ast::MatchExpr::cast)
+                let Some(match_expr) = sema
+                    .token_ancestors_with_macros(token)
+                    .take_while(|node| !ast::MacroCall::can_cast(node.kind()))
+                    .find_map(ast::MatchExpr::cast)
                 else {
                     continue;
                 };
@@ -339,11 +341,14 @@ pub(crate) fn highlight_branches(
         }
         T![=>] => {
             for token in sema.descend_into_macros(token.clone()) {
-                let Some(arm) =
-                    sema.token_ancestors_with_macros(token).find_map(ast::MatchArm::cast)
+                let Some(arm) = sema
+                    .token_ancestors_with_macros(token)
+                    .take_while(|node| !ast::MacroCall::can_cast(node.kind()))
+                    .find_map(ast::MatchArm::cast)
                 else {
                     continue;
                 };
+
                 let file_id = sema.hir_file_for(arm.syntax());
                 let range = arm.fat_arrow_token().map(|token| token.text_range());
                 push_to_highlights(file_id, range, &mut highlights);
@@ -352,9 +357,11 @@ pub(crate) fn highlight_branches(
             }
         }
         T![if] => {
-            for tok in sema.descend_into_macros(token.clone()) {
-                let Some(if_expr) =
-                    sema.token_ancestors_with_macros(tok).find_map(ast::IfExpr::cast)
+            for token in sema.descend_into_macros(token.clone()) {
+                let Some(if_expr) = sema
+                    .token_ancestors_with_macros(token)
+                    .take_while(|node| !ast::MacroCall::can_cast(node.kind()))
+                    .find_map(ast::IfExpr::cast)
                 else {
                     continue;
                 };
@@ -2388,4 +2395,25 @@ fn main() {
 "#,
         )
     }
+
+    #[test]
+    fn match_in_macro() {
+        // We should not highlight the outer `match` expression.
+        check(
+            r#"
+macro_rules! M {
+    (match) => { 1 };
+}
+
+fn main() {
+    match Some(1) {
+        Some(x) => x,
+        None => {
+            M!(match$0)
+        }
+    }
+}
+            "#,
+        )
+    }
 }