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

Author: roife

crates/ide/src/goto_definition.rs

@@ -420,8 +420,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());
@@ -434,8 +436,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)?;
@@ -451,8 +455,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

@@ -335,8 +335,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;
                 };
@@ -356,11 +358,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);
@@ -369,9 +374,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;
                 };
@@ -2418,6 +2425,27 @@ 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)
+        }
+    }
+}
+            "#,
+        )
+    }
+
     #[test]
     fn labeled_block_tail_expr() {
         check(