3 files changed, 74 insertions, 31 deletions

diff --git a/crates/ra_syntax/src/algo.rs b/crates/ra_syntax/src/algo.rs
index 2a8dac757..664894d1f 100644
--- a/crates/ra_syntax/src/algo.rs
+++ b/crates/ra_syntax/src/algo.rs
@@ -266,6 +266,15 @@ impl<'a> SyntaxRewriter<'a> {
         let replacement = Replacement::Single(with.clone().into());
         self.replacements.insert(what, replacement);
     }
+    pub fn replace_with_many<T: Clone + Into<SyntaxElement>>(
+        &mut self,
+        what: &T,
+        with: Vec<SyntaxElement>,
+    ) {
+        let what = what.clone().into();
+        let replacement = Replacement::Many(with);
+        self.replacements.insert(what, replacement);
+    }
     pub fn replace_ast<T: AstNode>(&mut self, what: &T, with: &T) {
         self.replace(what.syntax(), with.syntax())
     }
@@ -302,31 +311,41 @@ impl<'a> SyntaxRewriter<'a> {
 
     fn rewrite_children(&self, node: &SyntaxNode) -> SyntaxNode {
         //  FIXME: this could be made much faster.
-        let new_children =
-            node.children_with_tokens().flat_map(|it| self.rewrite_self(&it)).collect::<Vec<_>>();
+        let mut new_children = Vec::new();
+        for child in node.children_with_tokens() {
+            self.rewrite_self(&mut new_children, &child);
+        }
         with_children(node, new_children)
     }
 
     fn rewrite_self(
         &self,
+        acc: &mut Vec<NodeOrToken<rowan::GreenNode, rowan::GreenToken>>,
         element: &SyntaxElement,
-    ) -> Option<NodeOrToken<rowan::GreenNode, rowan::GreenToken>> {
+    ) {
         if let Some(replacement) = self.replacement(&element) {
-            return match replacement {
+            match replacement {
                 Replacement::Single(NodeOrToken::Node(it)) => {
-                    Some(NodeOrToken::Node(it.green().clone()))
+                    acc.push(NodeOrToken::Node(it.green().clone()))
                 }
                 Replacement::Single(NodeOrToken::Token(it)) => {
-                    Some(NodeOrToken::Token(it.green().clone()))
+                    acc.push(NodeOrToken::Token(it.green().clone()))
                 }
-                Replacement::Delete => None,
+                Replacement::Many(replacements) => {
+                    acc.extend(replacements.iter().map(|it| match it {
+                        NodeOrToken::Node(it) => NodeOrToken::Node(it.green().clone()),
+                        NodeOrToken::Token(it) => NodeOrToken::Token(it.green().clone()),
+                    }))
+                }
+                Replacement::Delete => (),
             };
+            return;
         }
         let res = match element {
             NodeOrToken::Token(it) => NodeOrToken::Token(it.green().clone()),
             NodeOrToken::Node(it) => NodeOrToken::Node(self.rewrite_children(it).green().clone()),
         };
-        Some(res)
+        acc.push(res)
     }
 }
 
@@ -341,6 +360,7 @@ impl ops::AddAssign for SyntaxRewriter<'_> {
 enum Replacement {
     Delete,
     Single(SyntaxElement),
+    Many(Vec<SyntaxElement>),
 }
 
 fn with_children(
diff --git a/crates/ra_syntax/src/ast/edit.rs b/crates/ra_syntax/src/ast/edit.rs
index 24a1e1d91..29eb3fcb9 100644
--- a/crates/ra_syntax/src/ast/edit.rs
+++ b/crates/ra_syntax/src/ast/edit.rs
@@ -1,7 +1,10 @@
 //! This module contains functions for editing syntax trees. As the trees are
 //! immutable, all function here return a fresh copy of the tree, instead of
 //! doing an in-place modification.
-use std::{iter, ops::RangeInclusive};
+use std::{
+    fmt, iter,
+    ops::{self, RangeInclusive},
+};
 
 use arrayvec::ArrayVec;
 
@@ -437,6 +440,28 @@ impl From<u8> for IndentLevel {
     }
 }
 
+impl fmt::Display for IndentLevel {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let spaces = "                                        ";
+        let buf;
+        let len = self.0 as usize * 4;
+        let indent = if len <= spaces.len() {
+            &spaces[..len]
+        } else {
+            buf = iter::repeat(' ').take(len).collect::<String>();
+            &buf
+        };
+        fmt::Display::fmt(indent, f)
+    }
+}
+
+impl ops::Add<u8> for IndentLevel {
+    type Output = IndentLevel;
+    fn add(self, rhs: u8) -> IndentLevel {
+        IndentLevel(self.0 + rhs)
+    }
+}
+
 impl IndentLevel {
     pub fn from_node(node: &SyntaxNode) -> IndentLevel {
         let first_token = match node.first_token() {
@@ -453,6 +478,14 @@ impl IndentLevel {
         IndentLevel(0)
     }
 
+    /// XXX: this intentionally doesn't change the indent of the very first token.
+    /// Ie, in something like
+    /// ```
+    /// fn foo() {
+    ///    92
+    /// }
+    /// ```
+    /// if you indent the block, the `{` token would stay put.
     fn increase_indent(self, node: SyntaxNode) -> SyntaxNode {
         let mut rewriter = SyntaxRewriter::default();
         node.descendants_with_tokens()
@@ -463,12 +496,7 @@ impl IndentLevel {
                 text.contains('\n')
             })
             .for_each(|ws| {
-                let new_ws = make::tokens::whitespace(&format!(
-                    "{}{:width$}",
-                    ws.syntax().text(),
-                    "",
-                    width = self.0 as usize * 4
-                ));
+                let new_ws = make::tokens::whitespace(&format!("{}{}", ws.syntax(), self,));
                 rewriter.replace(ws.syntax(), &new_ws)
             });
         rewriter.rewrite(&node)
@@ -485,7 +513,7 @@ impl IndentLevel {
             })
             .for_each(|ws| {
                 let new_ws = make::tokens::whitespace(
-                    &ws.syntax().text().replace(&format!("\n{:1$}", "", self.0 as usize * 4), "\n"),
+                    &ws.syntax().text().replace(&format!("\n{}", self), "\n"),
                 );
                 rewriter.replace(ws.syntax(), &new_ws)
             });
diff --git a/crates/ra_syntax/src/ast/make.rs b/crates/ra_syntax/src/ast/make.rs
index 2db017038..da0eb0926 100644
--- a/crates/ra_syntax/src/ast/make.rs
+++ b/crates/ra_syntax/src/ast/make.rs
@@ -280,7 +280,12 @@ pub fn param_list(pats: impl IntoIterator<Item = ast::Param>) -> ast::ParamList
     ast_from_text(&format!("fn f({}) {{ }}", args))
 }
 
+pub fn visibility_pub_crate() -> ast::Visibility {
+    ast_from_text("pub(crate) struct S")
+}
+
 pub fn fn_def(
+    visibility: Option<ast::Visibility>,
     fn_name: ast::Name,
     type_params: Option<ast::TypeParamList>,
     params: ast::ParamList,
@@ -288,21 +293,11 @@ pub fn fn_def(
 ) -> ast::FnDef {
     let type_params =
         if let Some(type_params) = type_params { format!("<{}>", type_params) } else { "".into() };
-    ast_from_text(&format!("fn {}{}{} {}", fn_name, type_params, params, body))
-}
-
-pub fn add_leading_newlines(amount_of_newlines: usize, t: impl AstNode) -> ast::SourceFile {
-    let newlines = "\n".repeat(amount_of_newlines);
-    ast_from_text(&format!("{}{}", newlines, t.syntax()))
-}
-
-pub fn add_trailing_newlines(amount_of_newlines: usize, t: impl AstNode) -> ast::SourceFile {
-    let newlines = "\n".repeat(amount_of_newlines);
-    ast_from_text(&format!("{}{}", t.syntax(), newlines))
-}
-
-pub fn add_pub_crate_modifier(fn_def: ast::FnDef) -> ast::FnDef {
-    ast_from_text(&format!("pub(crate) {}", fn_def))
+    let visibility = match visibility {
+        None => String::new(),
+        Some(it) => format!("{} ", it),
+    };
+    ast_from_text(&format!("{}fn {}{}{} {}", visibility, fn_name, type_params, params, body))
 }
 
 fn ast_from_text<N: AstNode>(text: &str) -> N {