Merge #9413

9413: internal: Deduplicate ast expression walking logic r=Veykril a=Veykril

Deduplicates the duplication introduced in #9375 and #9396 while also fixing a few bugs in both the assist and related highlighting.

Co-authored-by: Lukas Wirth <[email protected]>

Author: bors[bot]

crates/ide/src/highlight_related.rs

@@ -2,13 +2,13 @@ use hir::Semantics;
 use ide_db::{
     base_db::FilePosition,
     defs::Definition,
-    helpers::pick_best_token,
+    helpers::{for_each_break_expr, for_each_tail_expr, pick_best_token},
     search::{FileReference, ReferenceAccess, SearchScope},
     RootDatabase,
 };
 use syntax::{
     ast::{self, LoopBodyOwner},
-    match_ast, AstNode, SyntaxNode, SyntaxToken, TextRange, WalkEvent, T,
+    match_ast, AstNode, SyntaxNode, SyntaxToken, TextRange, T,
 };
 
 use crate::{display::TryToNav, references, NavigationTarget};
@@ -95,7 +95,7 @@ fn highlight_exit_points(
     ) -> Option<Vec<HighlightedRange>> {
         let mut highlights = Vec::new();
         let body = body?;
-        walk(&body, &mut |expr| match expr {
+        body.walk(&mut |expr| match expr {
             ast::Expr::ReturnExpr(expr) => {
                 if let Some(token) = expr.return_token() {
                     highlights.push(HighlightedRange { access: None, range: token.text_range() });
@@ -120,7 +120,7 @@ fn highlight_exit_points(
         };
 
         if let Some(tail) = tail {
-            for_each_inner_tail(&tail, &mut |tail| {
+            for_each_tail_expr(&tail, &mut |tail| {
                 let range = match tail {
                     ast::Expr::BreakExpr(b) => b
                         .break_token()
@@ -161,7 +161,7 @@ fn highlight_break_points(token: SyntaxToken) -> Option<Vec<HighlightedRange>> {
             label.as_ref().map(|it| it.syntax().text_range()),
         );
         highlights.extend(range.map(|range| HighlightedRange { access: None, range }));
-        for_each_break(label, body, &mut |break_| {
+        for_each_break_expr(label, body, &mut |break_| {
             let range = cover_range(
                 break_.break_token().map(|it| it.text_range()),
                 break_.lifetime().map(|it| it.syntax().text_range()),
@@ -216,7 +216,7 @@ fn highlight_yield_points(token: SyntaxToken) -> Option<Vec<HighlightedRange>> {
         let mut highlights = Vec::new();
         highlights.push(HighlightedRange { access: None, range: async_token?.text_range() });
         if let Some(body) = body {
-            walk(&body, &mut |expr| {
+            body.walk(&mut |expr| {
                 if let ast::Expr::AwaitExpr(expr) = expr {
                     if let Some(token) = expr.await_token() {
                         highlights
@@ -240,156 +240,6 @@ fn highlight_yield_points(token: SyntaxToken) -> Option<Vec<HighlightedRange>> {
     None
 }
 
-/// Preorder walk all the expression's child expressions
-fn walk(expr: &ast::Expr, cb: &mut dyn FnMut(ast::Expr)) {
-    let mut preorder = expr.syntax().preorder();
-    while let Some(event) = preorder.next() {
-        let node = match event {
-            WalkEvent::Enter(node) => node,
-            WalkEvent::Leave(_) => continue,
-        };
-        match ast::Stmt::cast(node.clone()) {
-            // recursively walk the initializer, skipping potential const pat expressions
-            // lets statements aren't usually nested too deeply so this is fine to recurse on
-            Some(ast::Stmt::LetStmt(l)) => {
-                if let Some(expr) = l.initializer() {
-                    walk(&expr, cb);
-                }
-                preorder.skip_subtree();
-            }
-            // Don't skip subtree since we want to process the expression child next
-            Some(ast::Stmt::ExprStmt(_)) => (),
-            // skip inner items which might have their own expressions
-            Some(ast::Stmt::Item(_)) => preorder.skip_subtree(),
-            None => {
-                if let Some(expr) = ast::Expr::cast(node) {
-                    let is_different_context = match &expr {
-                        ast::Expr::EffectExpr(effect) => {
-                            matches!(
-                                effect.effect(),
-                                ast::Effect::Async(_) | ast::Effect::Try(_) | ast::Effect::Const(_)
-                            )
-                        }
-                        ast::Expr::ClosureExpr(__) => true,
-                        _ => false,
-                    };
-                    cb(expr);
-                    if is_different_context {
-                        preorder.skip_subtree();
-                    }
-                } else {
-                    preorder.skip_subtree();
-                }
-            }
-        }
-    }
-}
-
-// FIXME: doesn't account for labeled breaks in labeled blocks
-fn for_each_inner_tail(expr: &ast::Expr, cb: &mut dyn FnMut(&ast::Expr)) {
-    match expr {
-        ast::Expr::BlockExpr(b) => {
-            if let Some(e) = b.tail_expr() {
-                for_each_inner_tail(&e, cb);
-            }
-        }
-        ast::Expr::EffectExpr(e) => match e.effect() {
-            ast::Effect::Label(label) => {
-                for_each_break(Some(label), e.block_expr(), &mut |b| cb(&ast::Expr::BreakExpr(b)));
-                if let Some(b) = e.block_expr() {
-                    for_each_inner_tail(&ast::Expr::BlockExpr(b), cb);
-                }
-            }
-            ast::Effect::Unsafe(_) => {
-                if let Some(e) = e.block_expr().and_then(|b| b.tail_expr()) {
-                    for_each_inner_tail(&e, cb);
-                }
-            }
-            ast::Effect::Async(_) | ast::Effect::Try(_) | ast::Effect::Const(_) => cb(expr),
-        },
-        ast::Expr::IfExpr(if_) => {
-            if_.blocks().for_each(|block| for_each_inner_tail(&ast::Expr::BlockExpr(block), cb))
-        }
-        ast::Expr::LoopExpr(l) => {
-            for_each_break(l.label(), l.loop_body(), &mut |b| cb(&ast::Expr::BreakExpr(b)))
-        }
-        ast::Expr::MatchExpr(m) => {
-            if let Some(arms) = m.match_arm_list() {
-                arms.arms().filter_map(|arm| arm.expr()).for_each(|e| for_each_inner_tail(&e, cb));
-            }
-        }
-        ast::Expr::ArrayExpr(_)
-        | ast::Expr::AwaitExpr(_)
-        | ast::Expr::BinExpr(_)
-        | ast::Expr::BoxExpr(_)
-        | ast::Expr::BreakExpr(_)
-        | ast::Expr::CallExpr(_)
-        | ast::Expr::CastExpr(_)
-        | ast::Expr::ClosureExpr(_)
-        | ast::Expr::ContinueExpr(_)
-        | ast::Expr::FieldExpr(_)
-        | ast::Expr::ForExpr(_)
-        | ast::Expr::IndexExpr(_)
-        | ast::Expr::Literal(_)
-        | ast::Expr::MacroCall(_)
-        | ast::Expr::MacroStmts(_)
-        | ast::Expr::MethodCallExpr(_)
-        | ast::Expr::ParenExpr(_)
-        | ast::Expr::PathExpr(_)
-        | ast::Expr::PrefixExpr(_)
-        | ast::Expr::RangeExpr(_)
-        | ast::Expr::RecordExpr(_)
-        | ast::Expr::RefExpr(_)
-        | ast::Expr::ReturnExpr(_)
-        | ast::Expr::TryExpr(_)
-        | ast::Expr::TupleExpr(_)
-        | ast::Expr::WhileExpr(_)
-        | ast::Expr::YieldExpr(_) => cb(expr),
-    }
-}
-
-fn for_each_break(
-    label: Option<ast::Label>,
-    body: Option<ast::BlockExpr>,
-    cb: &mut dyn FnMut(ast::BreakExpr),
-) {
-    let label = label.and_then(|lbl| lbl.lifetime());
-    let mut depth = 0;
-    if let Some(b) = body {
-        let preorder = &mut b.syntax().preorder();
-        let ev_as_expr = |ev| match ev {
-            WalkEvent::Enter(it) => Some(WalkEvent::Enter(ast::Expr::cast(it)?)),
-            WalkEvent::Leave(it) => Some(WalkEvent::Leave(ast::Expr::cast(it)?)),
-        };
-        let eq_label = |lt: Option<ast::Lifetime>| {
-            lt.zip(label.as_ref()).map_or(false, |(lt, lbl)| lt.text() == lbl.text())
-        };
-        while let Some(node) = preorder.find_map(ev_as_expr) {
-            match node {
-                WalkEvent::Enter(expr) => match expr {
-                    ast::Expr::LoopExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::ForExpr(_) => {
-                        depth += 1
-                    }
-                    ast::Expr::EffectExpr(e) if e.label().is_some() => depth += 1,
-                    ast::Expr::BreakExpr(b)
-                        if (depth == 0 && b.lifetime().is_none()) || eq_label(b.lifetime()) =>
-                    {
-                        cb(b);
-                    }
-                    _ => (),
-                },
-                WalkEvent::Leave(expr) => match expr {
-                    ast::Expr::LoopExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::ForExpr(_) => {
-                        depth -= 1
-                    }
-                    ast::Expr::EffectExpr(e) if e.label().is_some() => depth -= 1,
-                    _ => (),
-                },
-            }
-        }
-    }
-}
-
 fn cover_range(r0: Option<TextRange>, r1: Option<TextRange>) -> Option<TextRange> {
     match (r0, r1) {
         (Some(r0), Some(r1)) => Some(r0.cover(r1)),