From e6e4417bbbcc7e135d1b372e4e278cd3efa2d4b8 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Fri, 23 Apr 2021 18:36:43 +0200
Subject: [PATCH] Remove SyntaxRewriter::from_fn

---
 crates/ide_assists/src/ast_transform.rs | 33 ++++++++++++++-----------
 crates/ide_assists/src/utils.rs         |  3 ++-
 crates/syntax/src/algo.rs               | 19 +++-----------
 3 files changed, 24 insertions(+), 31 deletions(-)

diff --git a/crates/ide_assists/src/ast_transform.rs b/crates/ide_assists/src/ast_transform.rs
index 4a3ed7783e..e5ae718c9f 100644
--- a/crates/ide_assists/src/ast_transform.rs
+++ b/crates/ide_assists/src/ast_transform.rs
@@ -3,20 +3,27 @@ use hir::{HirDisplay, PathResolution, SemanticsScope};
 use ide_db::helpers::mod_path_to_ast;
 use rustc_hash::FxHashMap;
 use syntax::{
-    algo::SyntaxRewriter,
     ast::{self, AstNode},
-    SyntaxNode,
+    ted, SyntaxNode,
 };
 
-pub fn apply<'a, N: AstNode>(transformer: &dyn AstTransform<'a>, node: N) -> N {
-    SyntaxRewriter::from_fn(|element| match element {
-        syntax::SyntaxElement::Node(n) => {
-            let replacement = transformer.get_substitution(&n, transformer)?;
-            Some(replacement.into())
+pub fn apply<'a, N: AstNode>(transformer: &dyn AstTransform<'a>, node: &N) {
+    let mut skip_to = None;
+    for event in node.syntax().preorder() {
+        match event {
+            syntax::WalkEvent::Enter(node) if skip_to.is_none() => {
+                skip_to = transformer.get_substitution(&node, transformer).zip(Some(node));
+            }
+            syntax::WalkEvent::Enter(_) => (),
+            syntax::WalkEvent::Leave(node) => match &skip_to {
+                Some((replacement, skip_target)) if *skip_target == node => {
+                    ted::replace(node, replacement.clone_for_update());
+                    skip_to.take();
+                }
+                _ => (),
+            },
         }
-        _ => None,
-    })
-    .rewrite_ast(&node)
+    }
 }
 
 /// `AstTransform` helps with applying bulk transformations to syntax nodes.
@@ -191,11 +198,9 @@ impl<'a> AstTransform<'a> for QualifyPaths<'a> {
                 let found_path = from.find_use_path(self.source_scope.db.upcast(), def)?;
                 let mut path = mod_path_to_ast(&found_path);
 
-                let type_args = p
-                    .segment()
-                    .and_then(|s| s.generic_arg_list())
-                    .map(|arg_list| apply(recur, arg_list));
+                let type_args = p.segment().and_then(|s| s.generic_arg_list());
                 if let Some(type_args) = type_args {
+                    apply(recur, &type_args);
                     let last_segment = path.segment().unwrap();
                     path = path.with_segment(last_segment.with_generic_args(type_args))
                 }
diff --git a/crates/ide_assists/src/utils.rs b/crates/ide_assists/src/utils.rs
index d67524937e..5a90ad715b 100644
--- a/crates/ide_assists/src/utils.rs
+++ b/crates/ide_assists/src/utils.rs
@@ -140,7 +140,8 @@ pub fn add_trait_assoc_items_to_impl(
 
     let items = items
         .into_iter()
-        .map(|it| ast_transform::apply(&*ast_transform, it))
+        .map(|it| it.clone_for_update())
+        .inspect(|it| ast_transform::apply(&*ast_transform, it))
         .map(|it| match it {
             ast::AssocItem::Fn(def) => ast::AssocItem::Fn(add_body(def)),
             ast::AssocItem::TypeAlias(def) => ast::AssocItem::TypeAlias(def.remove_bounds()),
diff --git a/crates/syntax/src/algo.rs b/crates/syntax/src/algo.rs
index a153a9e1c3..c9229c4e07 100644
--- a/crates/syntax/src/algo.rs
+++ b/crates/syntax/src/algo.rs
@@ -342,10 +342,10 @@ enum InsertPos {
 
 #[derive(Default)]
 pub struct SyntaxRewriter<'a> {
-    f: Option<Box<dyn Fn(&SyntaxElement) -> Option<SyntaxElement> + 'a>>,
     //FIXME: add debug_assertions that all elements are in fact from the same file.
     replacements: FxHashMap<SyntaxElement, Replacement>,
     insertions: IndexMap<InsertPos, Vec<SyntaxElement>>,
+    _pd: std::marker::PhantomData<&'a ()>,
 }
 
 impl fmt::Debug for SyntaxRewriter<'_> {
@@ -357,14 +357,7 @@ impl fmt::Debug for SyntaxRewriter<'_> {
     }
 }
 
-impl<'a> SyntaxRewriter<'a> {
-    pub fn from_fn(f: impl Fn(&SyntaxElement) -> Option<SyntaxElement> + 'a) -> SyntaxRewriter<'a> {
-        SyntaxRewriter {
-            f: Some(Box::new(f)),
-            replacements: FxHashMap::default(),
-            insertions: IndexMap::default(),
-        }
-    }
+impl SyntaxRewriter<'_> {
     pub fn delete<T: Clone + Into<SyntaxElement>>(&mut self, what: &T) {
         let what = what.clone().into();
         let replacement = Replacement::Delete;
@@ -470,7 +463,7 @@ impl<'a> SyntaxRewriter<'a> {
     pub fn rewrite(&self, node: &SyntaxNode) -> SyntaxNode {
         let _p = profile::span("rewrite");
 
-        if self.f.is_none() && self.replacements.is_empty() && self.insertions.is_empty() {
+        if self.replacements.is_empty() && self.insertions.is_empty() {
             return node.clone();
         }
         let green = self.rewrite_children(node);
@@ -495,7 +488,6 @@ impl<'a> SyntaxRewriter<'a> {
             }
         }
 
-        assert!(self.f.is_none());
         self.replacements
             .keys()
             .filter_map(element_to_node_or_parent)
@@ -510,10 +502,6 @@ impl<'a> SyntaxRewriter<'a> {
     }
 
     fn replacement(&self, element: &SyntaxElement) -> Option<Replacement> {
-        if let Some(f) = &self.f {
-            assert!(self.replacements.is_empty());
-            return f(element).map(Replacement::Single);
-        }
         self.replacements.get(element).cloned()
     }
 
@@ -574,7 +562,6 @@ fn element_to_green(element: SyntaxElement) -> NodeOrToken<rowan::GreenNode, row
 
 impl ops::AddAssign for SyntaxRewriter<'_> {
     fn add_assign(&mut self, rhs: SyntaxRewriter) {
-        assert!(rhs.f.is_none());
         self.replacements.extend(rhs.replacements);
         for (pos, insertions) in rhs.insertions.into_iter() {
             match self.insertions.entry(pos) {