From bff9f4cdc5ed96fc977ee69f0f4a9be636b022b8 Mon Sep 17 00:00:00 2001 From: roifewu Date: Tue, 15 Apr 2025 11:04:46 +0800 Subject: [PATCH] refactor: simplify functions related to branch_exit_points in highlight_related --- .../crates/ide/src/goto_definition.rs | 90 +++++++++++++------ .../crates/ide/src/highlight_related.rs | 46 ++-------- .../crates/ide/src/references.rs | 10 ++- 3 files changed, 74 insertions(+), 72 deletions(-) diff --git a/src/tools/rust-analyzer/crates/ide/src/goto_definition.rs b/src/tools/rust-analyzer/crates/ide/src/goto_definition.rs index 35eb18e3c720..c90a544a718a 100644 --- a/src/tools/rust-analyzer/crates/ide/src/goto_definition.rs +++ b/src/tools/rust-analyzer/crates/ide/src/goto_definition.rs @@ -291,14 +291,14 @@ fn handle_control_flow_keywords( token: &SyntaxToken, ) -> Option> { match token.kind() { - // For `fn` / `loop` / `while` / `for` / `async`, return the keyword it self, + // For `fn` / `loop` / `while` / `for` / `async` / `match`, return the keyword it self, // so that VSCode will find the references when using `ctrl + click` T![fn] | T![async] | T![try] | T![return] => nav_for_exit_points(sema, token), T![loop] | T![while] | T![break] | T![continue] => nav_for_break_points(sema, token), T![for] if token.parent().and_then(ast::ForExpr::cast).is_some() => { nav_for_break_points(sema, token) } - T![match] | T![=>] | T![if] => nav_for_branches(sema, token), + T![match] | T![=>] | T![if] => nav_for_branch_exit_points(sema, token), _ => None, } } @@ -408,22 +408,66 @@ fn nav_for_exit_points( Some(navs) } -fn nav_for_branches( +pub(crate) fn find_branch_root( + sema: &Semantics<'_, RootDatabase>, + token: &SyntaxToken, +) -> Vec { + fn find_root( + sema: &Semantics<'_, RootDatabase>, + token: &SyntaxToken, + pred: impl Fn(SyntaxNode) -> Option, + ) -> Vec { + let mut result = Vec::new(); + for token in sema.descend_into_macros(token.clone()) { + for node in sema.token_ancestors_with_macros(token) { + if ast::MacroCall::can_cast(node.kind()) { + break; + } + + if let Some(node) = pred(node) { + result.push(node); + break; + } + } + } + result + } + + match token.kind() { + T![match] => { + find_root(sema, token, |node| Some(ast::MatchExpr::cast(node)?.syntax().clone())) + } + T![=>] => find_root(sema, token, |node| Some(ast::MatchArm::cast(node)?.syntax().clone())), + T![if] => find_root(sema, token, |node| { + let if_expr = ast::IfExpr::cast(node)?; + + iter::successors(Some(if_expr.clone()), |if_expr| { + let parent_if = if_expr.syntax().parent().and_then(ast::IfExpr::cast)?; + if let ast::ElseBranch::IfExpr(nested_if) = parent_if.else_branch()? { + (nested_if.syntax() == if_expr.syntax()).then_some(parent_if) + } else { + None + } + }) + .last() + .map(|if_expr| if_expr.syntax().clone()) + }), + _ => vec![], + } +} + +fn nav_for_branch_exit_points( sema: &Semantics<'_, RootDatabase>, token: &SyntaxToken, ) -> Option> { let db = sema.db; let navs = match token.kind() { - T![match] => sema - .descend_into_macros(token.clone()) + T![match] => find_branch_root(sema, token) .into_iter() - .filter_map(|token| { - 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()); + .filter_map(|node| { + let file_id = sema.hir_file_for(&node); + let match_expr = ast::MatchExpr::cast(node)?; let focus_range = match_expr.match_token()?.text_range(); let match_expr_in_file = InFile::new(file_id, match_expr.into()); Some(expr_to_nav(db, match_expr_in_file, Some(focus_range))) @@ -431,14 +475,10 @@ fn nav_for_branches( .flatten() .collect_vec(), - T![=>] => sema - .descend_into_macros(token.clone()) + T![=>] => find_branch_root(sema, token) .into_iter() - .filter_map(|token| { - let match_arm = sema - .token_ancestors_with_macros(token) - .take_while(|node| !ast::MacroCall::can_cast(node.kind())) - .find_map(ast::MatchArm::cast)?; + .filter_map(|node| { + let match_arm = ast::MatchArm::cast(node)?; let match_expr = sema .ancestors_with_macros(match_arm.syntax().clone()) .find_map(ast::MatchExpr::cast)?; @@ -450,15 +490,11 @@ fn nav_for_branches( .flatten() .collect_vec(), - T![if] => sema - .descend_into_macros(token.clone()) + T![if] => find_branch_root(sema, token) .into_iter() - .filter_map(|token| { - 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()); + .filter_map(|node| { + let file_id = sema.hir_file_for(&node); + let if_expr = ast::IfExpr::cast(node)?; let focus_range = if_expr.if_token()?.text_range(); let if_expr_in_file = InFile::new(file_id, if_expr.into()); Some(expr_to_nav(db, if_expr_in_file, Some(focus_range))) @@ -3785,9 +3821,9 @@ fn main() { r#" fn main() { if true { + // ^^ () } else if$0 false { - // ^^ () } else { () diff --git a/src/tools/rust-analyzer/crates/ide/src/highlight_related.rs b/src/tools/rust-analyzer/crates/ide/src/highlight_related.rs index 77b8599fd021..6aa17f90ef05 100644 --- a/src/tools/rust-analyzer/crates/ide/src/highlight_related.rs +++ b/src/tools/rust-analyzer/crates/ide/src/highlight_related.rs @@ -331,17 +331,10 @@ pub(crate) fn highlight_branch_exit_points( }); }; + let nodes = goto_definition::find_branch_root(sema, &token).into_iter(); match token.kind() { T![match] => { - for token in sema.descend_into_macros(token.clone()) { - 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; - }; - + for match_expr in nodes.filter_map(ast::MatchExpr::cast) { let file_id = sema.hir_file_for(match_expr.syntax()); let range = match_expr.match_token().map(|token| token.text_range()); push_to_highlights(file_id, range, &mut highlights); @@ -349,22 +342,13 @@ pub(crate) fn highlight_branch_exit_points( let Some(arm_list) = match_expr.match_arm_list() else { continue; }; - for arm in arm_list.arms() { push_tail_expr(arm.expr(), &mut highlights); } } } T![=>] => { - for token in sema.descend_into_macros(token.clone()) { - 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; - }; - + for arm in nodes.filter_map(ast::MatchArm::cast) { 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); @@ -373,27 +357,7 @@ pub(crate) fn highlight_branch_exit_points( } } T![if] => { - 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; - }; - - // Find the root of the if expression - let mut if_to_process = iter::successors(Some(if_expr.clone()), |if_expr| { - let parent_if = if_expr.syntax().parent().and_then(ast::IfExpr::cast)?; - if let ast::ElseBranch::IfExpr(nested_if) = parent_if.else_branch()? { - (nested_if.syntax() == if_expr.syntax()).then_some(parent_if) - } else { - None - } - }) - .last() - .or(Some(if_expr)); - + for mut if_to_process in nodes.map(ast::IfExpr::cast) { while let Some(cur_if) = if_to_process.take() { let file_id = sema.hir_file_for(cur_if.syntax()); @@ -415,7 +379,7 @@ pub(crate) fn highlight_branch_exit_points( } } } - _ => unreachable!(), + _ => {} } highlights diff --git a/src/tools/rust-analyzer/crates/ide/src/references.rs b/src/tools/rust-analyzer/crates/ide/src/references.rs index 205e93d7a24c..fe874bc99b40 100644 --- a/src/tools/rust-analyzer/crates/ide/src/references.rs +++ b/src/tools/rust-analyzer/crates/ide/src/references.rs @@ -21,6 +21,7 @@ use hir::{PathResolution, Semantics}; use ide_db::{ FileId, RootDatabase, defs::{Definition, NameClass, NameRefClass}, + helpers::pick_best_token, search::{ReferenceCategory, SearchScope, UsageSearchResult}, }; use itertools::Itertools; @@ -397,10 +398,11 @@ fn handle_control_flow_keywords( .attach_first_edition(file_id) .map(|it| it.edition(sema.db)) .unwrap_or(Edition::CURRENT); - let token = file - .syntax() - .token_at_offset(offset) - .find(|t| t.kind().is_keyword(edition) || t.kind() == T![=>])?; + let token = pick_best_token(file.syntax().token_at_offset(offset), |kind| match kind { + _ if kind.is_keyword(edition) => 4, + T![=>] => 3, + _ => 1, + })?; let references = match token.kind() { T![fn] | T![return] | T![try] => highlight_related::highlight_exit_points(sema, token),