Simplify extract_function assist

2 files changed, 118 insertions, 101 deletions

diff --git a/crates/syntax/src/ast/expr_ext.rs b/crates/syntax/src/ast/expr_ext.rs
index 6317d84ba3b..5307516ecb7 100644
--- a/crates/syntax/src/ast/expr_ext.rs
+++ b/crates/syntax/src/ast/expr_ext.rs
@@ -1,5 +1,7 @@
 //! Various extension methods to ast Expr Nodes, which are hard to code-generate.
 
+use rowan::WalkEvent;
+
 use crate::{
     ast::{self, support, AstChildren, AstNode},
     AstToken,
@@ -34,6 +36,122 @@ impl ast::Expr {
         }
         None
     }
+
+    /// Preorder walk all the expression's child expressions.
+    pub fn walk(&self, cb: &mut dyn FnMut(ast::Expr)) {
+        self.preorder(&mut |ev| {
+            if let WalkEvent::Enter(expr) = ev {
+                cb(expr);
+            }
+            false
+        })
+    }
+
+    /// Preorder walk all the expression's child expressions preserving events.
+    /// If the callback returns true the subtree of the expression will be skipped.
+    /// Note that the subtree may already be skipped due to the context analysis this function does.
+    pub fn preorder(&self, cb: &mut dyn FnMut(WalkEvent<ast::Expr>) -> bool) {
+        let mut preorder = self.syntax().preorder();
+        while let Some(event) = preorder.next() {
+            let node = match event {
+                WalkEvent::Enter(node) => node,
+                WalkEvent::Leave(node) => {
+                    if let Some(expr) = ast::Expr::cast(node) {
+                        cb(WalkEvent::Leave(expr));
+                    }
+                    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.preorder(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,
+                        };
+                        let skip = cb(WalkEvent::Enter(expr));
+                        if skip || is_different_context {
+                            preorder.skip_subtree();
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    /// Preorder walk all the expression's child patterns.
+    pub fn walk_patterns(&self, cb: &mut dyn FnMut(ast::Pat)) {
+        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()) {
+                Some(ast::Stmt::LetStmt(l)) => {
+                    if let Some(pat) = l.pat() {
+                        pat.walk(cb);
+                    }
+                    if let Some(expr) = l.initializer() {
+                        expr.walk_patterns(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 patterns
+                Some(ast::Stmt::Item(_)) => preorder.skip_subtree(),
+                None => {
+                    // skip const args, those are a different context
+                    if ast::GenericArg::can_cast(node.kind()) {
+                        preorder.skip_subtree();
+                    } else if let Some(expr) = ast::Expr::cast(node.clone()) {
+                        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,
+                        };
+                        if is_different_context {
+                            preorder.skip_subtree();
+                        }
+                    } else if let Some(pat) = ast::Pat::cast(node) {
+                        preorder.skip_subtree();
+                        pat.walk(cb);
+                    }
+                }
+            }
+        }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
diff --git a/crates/syntax/src/ast/node_ext.rs b/crates/syntax/src/ast/node_ext.rs
index e9465536c8d..ea156087fff 100644
--- a/crates/syntax/src/ast/node_ext.rs
+++ b/crates/syntax/src/ast/node_ext.rs
@@ -55,107 +55,6 @@ 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();
-                        }
-                    }
-                }
-            }
-        }
-    }
-
-    /// Preorder walk all the expression's child patterns.
-    pub fn walk_patterns(&self, cb: &mut dyn FnMut(ast::Pat)) {
-        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()) {
-                Some(ast::Stmt::LetStmt(l)) => {
-                    if let Some(pat) = l.pat() {
-                        pat.walk(cb);
-                    }
-                    if let Some(expr) = l.initializer() {
-                        expr.walk_patterns(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 patterns
-                Some(ast::Stmt::Item(_)) => preorder.skip_subtree(),
-                None => {
-                    // skip const args, those are a different context
-                    if ast::GenericArg::can_cast(node.kind()) {
-                        preorder.skip_subtree();
-                    } else if let Some(expr) = ast::Expr::cast(node.clone()) {
-                        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,
-                        };
-                        if is_different_context {
-                            preorder.skip_subtree();
-                        }
-                    } else if let Some(pat) = ast::Pat::cast(node) {
-                        preorder.skip_subtree();
-                        pat.walk(cb);
-                    }
-                }
-            }
-        }
-    }
-}
-
 impl ast::Pat {
     /// Preorder walk all the pattern's sub patterns.
     pub fn walk(&self, cb: &mut dyn FnMut(ast::Pat)) {