Deduplicate ast expression walking logic

crates/syntax/src/ast/node_ext.rs

@@ -5,7 +5,7 @@ use std::{borrow::Cow, fmt, iter::successors};
 
 use itertools::Itertools;
 use parser::SyntaxKind;
-use rowan::{GreenNodeData, GreenTokenData};
+use rowan::{GreenNodeData, GreenTokenData, WalkEvent};
 
 use crate::{
     ast::{self, support, AstChildren, AstNode, AstToken, AttrsOwner, NameOwner, SyntaxNode},
@@ -51,6 +51,56 @@ impl ast::BlockExpr {
     }
 }
 
+impl ast::Expr {
+    /// Preorder walk all the expression's child expressions.
+    pub fn walk(&self, cb: &mut dyn FnMut(ast::Expr)) {
+        let mut preorder = self.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
+                // let 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() {
+                        expr.walk(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 => {
+                    // skip const args, those expressions are a different context
+                    if ast::GenericArg::can_cast(node.kind()) {
+                        preorder.skip_subtree();
+                    } else 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();
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub enum Macro {
     MacroRules(ast::MacroRules),