Generalise syntax rewriting infrastructure to allow removal of nodes

Author: matklad

crates/ra_assists/src/assist_ctx.rs

@@ -11,6 +11,7 @@ use ra_syntax::{
 use ra_text_edit::TextEditBuilder;
 
 use crate::{AssistAction, AssistId, AssistLabel, GroupLabel, ResolvedAssist};
+use algo::SyntaxRewriter;
 
 #[derive(Clone, Debug)]
 pub(crate) struct Assist(pub(crate) Vec<AssistInfo>);
@@ -234,6 +235,11 @@ impl ActionBuilder {
     pub(crate) fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
         algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
     }
+    pub(crate) fn rewrite(&mut self, rewriter: SyntaxRewriter) {
+        let node = rewriter.rewrite_root().unwrap();
+        let new = rewriter.rewrite(&node);
+        algo::diff(&node, &new).into_text_edit(&mut self.edit)
+    }
 
     fn build(self) -> AssistAction {
         AssistAction {

crates/ra_assists/src/ast_transform.rs

@@ -3,7 +3,10 @@ use rustc_hash::FxHashMap;
 
 use hir::{PathResolution, SemanticsScope};
 use ra_ide_db::RootDatabase;
-use ra_syntax::ast::{self, AstNode};
+use ra_syntax::{
+    algo::SyntaxRewriter,
+    ast::{self, AstNode},
+};
 
 pub trait AstTransform<'a> {
     fn get_substitution(&self, node: &ra_syntax::SyntaxNode) -> Option<ra_syntax::SyntaxNode>;
@@ -153,15 +156,14 @@ impl<'a> QualifyPaths<'a> {
 }
 
 pub fn apply<'a, N: AstNode>(transformer: &dyn AstTransform<'a>, node: N) -> N {
-    let syntax = node.syntax();
-    let result = ra_syntax::algo::replace_descendants(syntax, |element| match element {
+    SyntaxRewriter::from_fn(|element| match element {
         ra_syntax::SyntaxElement::Node(n) => {
             let replacement = transformer.get_substitution(&n)?;
             Some(replacement.into())
         }
         _ => None,
-    });
-    N::cast(result).unwrap()
+    })
+    .rewrite_ast(&node)
 }
 
 impl<'a> AstTransform<'a> for QualifyPaths<'a> {

crates/ra_assists/src/handlers/merge_imports.rs

@@ -1,9 +1,9 @@
 use std::iter::successors;
 
 use ra_syntax::{
-    algo::neighbor,
+    algo::{neighbor, SyntaxRewriter},
     ast::{self, edit::AstNodeEdit, make},
-    AstNode, AstToken, Direction, InsertPosition, SyntaxElement, TextRange, T,
+    AstNode, Direction, InsertPosition, SyntaxElement, T,
 };
 
 use crate::{Assist, AssistCtx, AssistId};
@@ -22,52 +22,39 @@ use crate::{Assist, AssistCtx, AssistId};
 // ```
 pub(crate) fn merge_imports(ctx: AssistCtx) -> Option<Assist> {
     let tree: ast::UseTree = ctx.find_node_at_offset()?;
-    let (new_tree, to_delete) = if let Some(use_item) =
-        tree.syntax().parent().and_then(ast::UseItem::cast)
-    {
+    let mut rewriter = SyntaxRewriter::default();
+    let mut offset = ctx.frange.range.start();
+
+    if let Some(use_item) = tree.syntax().parent().and_then(ast::UseItem::cast) {
         let (merged, to_delete) = next_prev()
             .filter_map(|dir| neighbor(&use_item, dir))
             .filter_map(|it| Some((it.clone(), it.use_tree()?)))
             .find_map(|(use_item, use_tree)| {
                 Some((try_merge_trees(&tree, &use_tree)?, use_item.clone()))
             })?;
 
-        let mut range = to_delete.syntax().text_range();
-        let next_ws = to_delete
-            .syntax()
-            .next_sibling_or_token()
-            .and_then(|it| it.into_token())
-            .and_then(ast::Whitespace::cast);
-        if let Some(ws) = next_ws {
-            range = range.extend_to(&ws.syntax().text_range())
+        rewriter.replace_ast(&tree, &merged);
+        rewriter += to_delete.remove();
+
+        if to_delete.syntax().text_range().end() < offset {
+            offset -= to_delete.syntax().text_range().len();
         }
-        (merged, range)
     } else {
         let (merged, to_delete) = next_prev()
             .filter_map(|dir| neighbor(&tree, dir))
             .find_map(|use_tree| Some((try_merge_trees(&tree, &use_tree)?, use_tree.clone())))?;
 
-        let mut range = to_delete.syntax().text_range();
-        if let Some((dir, nb)) = next_prev().find_map(|dir| Some((dir, neighbor(&to_delete, dir)?)))
-        {
-            let nb_range = nb.syntax().text_range();
-            if dir == Direction::Prev {
-                range = TextRange::from_to(nb_range.end(), range.end());
-            } else {
-                range = TextRange::from_to(range.start(), nb_range.start());
-            }
+        rewriter.replace_ast(&tree, &merged);
+        rewriter += to_delete.remove();
+
+        if to_delete.syntax().text_range().end() < offset {
+            offset -= to_delete.syntax().text_range().len();
         }
-        (merged, range)
     };
 
-    let mut offset = ctx.frange.range.start();
     ctx.add_assist(AssistId("merge_imports"), "Merge imports", |edit| {
-        edit.replace_ast(tree, new_tree);
-        edit.delete(to_delete);
-
-        if to_delete.end() <= offset {
-            offset -= to_delete.len();
-        }
+        edit.rewrite(rewriter);
+        // FIXME: we only need because our diff is imprecise
         edit.set_cursor(offset);
     })
 }
@@ -156,7 +143,7 @@ use std::fmt::Debug;
 use std::fmt<|>::Display;
 ",
             r"
-use std::fmt<|>::{Display, Debug};
+use std::fmt:<|>:{Display, Debug};
 ",
         );
     }
@@ -178,7 +165,7 @@ use std::{fmt<|>::{Debug, Display}};
 use std::{fmt::Debug, fmt<|>::Display};
 ",
             r"
-use std::{fmt<|>::{Display, Debug}};
+use std::{fmt::<|>{Display, Debug}};
 ",
         );
     }
@@ -197,6 +184,38 @@ use foo::baz;
 use foo<|>::{bar, baz};
 
 /// Doc comment
+",
+        );
+    }
+
+    #[test]
+    fn works_with_trailing_comma() {
+        check_assist(
+            merge_imports,
+            r"
+use {
+    foo<|>::bar,
+    foo::baz,
+};
+",
+            r"
+use {
+    foo<|>::{bar, baz},
+};
+",
+        );
+        check_assist(
+            merge_imports,
+            r"
+use {
+    foo::baz,
+    foo<|>::bar,
+};
+",
+            r"
+use {
+    foo::{bar<|>, baz},
+};
 ",
         );
     }

crates/ra_hir_expand/src/eager.rs

@@ -26,8 +26,8 @@ use crate::{
 };
 
 use ra_parser::FragmentKind;
-use ra_syntax::{algo::replace_descendants, SyntaxElement, SyntaxNode};
-use std::{collections::HashMap, sync::Arc};
+use ra_syntax::{algo::SyntaxRewriter, SyntaxNode};
+use std::sync::Arc;
 
 pub fn expand_eager_macro(
     db: &dyn AstDatabase,
@@ -95,10 +95,10 @@ fn eager_macro_recur(
     curr: InFile<SyntaxNode>,
     macro_resolver: &dyn Fn(ast::Path) -> Option<MacroDefId>,
 ) -> Option<SyntaxNode> {
-    let mut original = curr.value.clone();
+    let original = curr.value.clone();
 
     let children = curr.value.descendants().filter_map(ast::MacroCall::cast);
-    let mut replaces: HashMap<SyntaxElement, SyntaxElement> = HashMap::default();
+    let mut rewriter = SyntaxRewriter::default();
 
     // Collect replacement
     for child in children {
@@ -119,12 +119,9 @@ fn eager_macro_recur(
             }
         };
 
-        replaces.insert(child.syntax().clone().into(), insert.into());
+        rewriter.replace(child.syntax(), &insert);
     }
 
-    if !replaces.is_empty() {
-        original = replace_descendants(&original, |n| replaces.get(n).cloned());
-    }
-
-    Some(original)
+    let res = rewriter.rewrite(&original);
+    Some(res)
 }

crates/ra_ide/src/expand_macro.rs

@@ -3,10 +3,9 @@
 use hir::Semantics;
 use ra_ide_db::RootDatabase;
 use ra_syntax::{
-    algo::{find_node_at_offset, replace_descendants},
-    ast, AstNode, NodeOrToken, SyntaxElement, SyntaxKind, SyntaxNode, WalkEvent, T,
+    algo::{find_node_at_offset, SyntaxRewriter},
+    ast, AstNode, NodeOrToken, SyntaxKind, SyntaxNode, WalkEvent, T,
 };
-use rustc_hash::FxHashMap;
 
 use crate::FilePosition;
 
@@ -37,7 +36,7 @@ fn expand_macro_recur(
     let mut expanded = sema.expand(macro_call)?;
 
     let children = expanded.descendants().filter_map(ast::MacroCall::cast);
-    let mut replaces: FxHashMap<SyntaxElement, SyntaxElement> = FxHashMap::default();
+    let mut rewriter = SyntaxRewriter::default();
 
     for child in children.into_iter() {
         if let Some(new_node) = expand_macro_recur(sema, &child) {
@@ -47,12 +46,13 @@ fn expand_macro_recur(
             if expanded == *child.syntax() {
                 expanded = new_node;
             } else {
-                replaces.insert(child.syntax().clone().into(), new_node.into());
+                rewriter.replace(child.syntax(), &new_node)
             }
         }
     }
 
-    Some(replace_descendants(&expanded, |n| replaces.get(n).cloned()))
+    let res = rewriter.rewrite(&expanded);
+    Some(res)
 }
 
 // FIXME: It would also be cool to share logic here and in the mbe tests,