1 files changed, 98 insertions, 42 deletions

diff --git a/crates/ra_assists/src/ast_editor.rs b/crates/ra_assists/src/ast_editor.rs
index 7347c1738..5bf1fc0ed 100644
--- a/crates/ra_assists/src/ast_editor.rs
+++ b/crates/ra_assists/src/ast_editor.rs
@@ -1,6 +1,9 @@
+use std::iter;
+
 use arrayvec::ArrayVec;
-use ra_text_edit::{TextEdit, TextEditBuilder};
+use ra_text_edit::TextEditBuilder;
 use ra_syntax::{AstNode, TreeArc, ast, SyntaxKind::*, SyntaxElement, SourceFile, InsertPosition, Direction};
+use ra_fmt::leading_indent;
 
 pub struct AstEditor<N: AstNode> {
     original_ast: TreeArc<N>,
@@ -12,18 +15,25 @@ impl<N: AstNode> AstEditor<N> {
         AstEditor { original_ast: node.to_owned(), ast: node.to_owned() }
     }
 
-    pub fn into_text_edit(self) -> TextEdit {
+    pub fn into_text_edit(self, builder: &mut TextEditBuilder) {
         // FIXME: compute a more fine-grained diff here.
         // If *you* know a nice algorithm to compute diff between two syntax
         // tree, tell me about it!
-        let mut builder = TextEditBuilder::default();
         builder.replace(self.original_ast.syntax().range(), self.ast().syntax().text().to_string());
-        builder.finish()
     }
 
     pub fn ast(&self) -> &N {
         &*self.ast
     }
+
+    fn insert_children<'a>(
+        &self,
+        position: InsertPosition<SyntaxElement<'_>>,
+        to_insert: impl Iterator<Item = SyntaxElement<'a>>,
+    ) -> TreeArc<N> {
+        let new_syntax = self.ast().syntax().insert_children(position, to_insert);
+        N::cast(&new_syntax).unwrap().to_owned()
+    }
 }
 
 impl AstEditor<ast::NamedFieldList> {
@@ -31,41 +41,61 @@ impl AstEditor<ast::NamedFieldList> {
         self.insert_field(InsertPosition::Last, field)
     }
 
+    pub fn make_multiline(&mut self) {
+        let l_curly = match self.l_curly() {
+            Some(it) => it,
+            None => return,
+        };
+        let sibling = match l_curly.next_sibling_or_token() {
+            Some(it) => it,
+            None => return,
+        };
+        if sibling.as_token().map(|it| it.text().contains('\n')) == Some(true) {
+            return;
+        }
+
+        let ws = tokens::WsBuilder::new(&format!(
+            "\n{}",
+            leading_indent(self.ast().syntax()).unwrap_or("")
+        ));
+        self.ast = self.insert_children(InsertPosition::After(l_curly), iter::once(ws.ws().into()));
+    }
+
     pub fn insert_field(
         &mut self,
         position: InsertPosition<&'_ ast::NamedField>,
         field: &ast::NamedField,
     ) {
-        let mut to_insert: ArrayVec<[SyntaxElement; 2]> =
-            [field.syntax().into(), tokens::comma().into()].into();
-        let position = match position {
-            InsertPosition::First => {
-                let anchor = match self
-                    .ast()
-                    .syntax()
-                    .children_with_tokens()
-                    .find(|it| it.kind() == L_CURLY)
-                {
+        let is_multiline = self.ast().syntax().text().contains('\n');
+        let ws;
+        let space = if is_multiline {
+            ws = tokens::WsBuilder::new(&format!(
+                "\n{}    ",
+                leading_indent(self.ast().syntax()).unwrap_or("")
+            ));
+            ws.ws()
+        } else {
+            tokens::single_space()
+        };
+
+        let mut to_insert: ArrayVec<[SyntaxElement; 4]> = ArrayVec::new();
+        to_insert.push(space.into());
+        to_insert.push(field.syntax().into());
+        to_insert.push(tokens::comma().into());
+
+        macro_rules! after_l_curly {
+            () => {{
+                let anchor = match self.l_curly() {
                     Some(it) => it,
                     None => return,
                 };
                 InsertPosition::After(anchor)
-            }
-            InsertPosition::Last => {
-                let anchor = match self
-                    .ast()
-                    .syntax()
-                    .children_with_tokens()
-                    .find(|it| it.kind() == R_CURLY)
-                {
-                    Some(it) => it,
-                    None => return,
-                };
-                InsertPosition::Before(anchor)
-            }
-            InsertPosition::Before(anchor) => InsertPosition::Before(anchor.syntax().into()),
-            InsertPosition::After(anchor) => {
-                if let Some(comma) = anchor
+            }};
+        }
+
+        macro_rules! after_field {
+            ($anchor:expr) => {
+                if let Some(comma) = $anchor
                     .syntax()
                     .siblings_with_tokens(Direction::Next)
                     .find(|it| it.kind() == COMMA)
@@ -73,21 +103,26 @@ impl AstEditor<ast::NamedFieldList> {
                     InsertPosition::After(comma)
                 } else {
                     to_insert.insert(0, tokens::comma().into());
-                    InsertPosition::After(anchor.syntax().into())
+                    InsertPosition::After($anchor.syntax().into())
                 }
-            }
+            };
+        };
+
+        let position = match position {
+            InsertPosition::First => after_l_curly!(),
+            InsertPosition::Last => match self.ast().fields().last() {
+                Some(it) => after_field!(it),
+                None => after_l_curly!(),
+            },
+            InsertPosition::Before(anchor) => InsertPosition::Before(anchor.syntax().into()),
+            InsertPosition::After(anchor) => after_field!(anchor),
         };
-        self.ast = insert_children_into_ast(self.ast(), position, to_insert.iter().cloned());
+        self.ast = self.insert_children(position, to_insert.iter().cloned());
     }
-}
 
-fn insert_children_into_ast<'a, N: AstNode>(
-    node: &N,
-    position: InsertPosition<SyntaxElement<'_>>,
-    to_insert: impl Iterator<Item = SyntaxElement<'a>>,
-) -> TreeArc<N> {
-    let new_syntax = node.syntax().insert_children(position, to_insert);
-    N::cast(&new_syntax).unwrap().to_owned()
+    fn l_curly(&self) -> Option<SyntaxElement> {
+        self.ast().syntax().children_with_tokens().find(|it| it.kind() == L_CURLY)
+    }
 }
 
 pub struct AstBuilder<N: AstNode> {
@@ -111,7 +146,7 @@ mod tokens {
     use ra_syntax::{AstNode, SourceFile, TreeArc, SyntaxToken, SyntaxKind::*};
 
     lazy_static! {
-        static ref SOURCE_FILE: TreeArc<SourceFile> = SourceFile::parse(",");
+        static ref SOURCE_FILE: TreeArc<SourceFile> = SourceFile::parse(",\n; ;");
     }
 
     pub(crate) fn comma() -> SyntaxToken<'static> {
@@ -122,6 +157,27 @@ mod tokens {
             .find(|it| it.kind() == COMMA)
             .unwrap()
     }
+
+    pub(crate) fn single_space() -> SyntaxToken<'static> {
+        SOURCE_FILE
+            .syntax()
+            .descendants_with_tokens()
+            .filter_map(|it| it.as_token())
+            .find(|it| it.kind() == WHITESPACE && it.text().as_str() == " ")
+            .unwrap()
+    }
+
+    pub(crate) struct WsBuilder(TreeArc<SourceFile>);
+
+    impl WsBuilder {
+        pub(crate) fn new(text: &str) -> WsBuilder {
+            WsBuilder(SourceFile::parse(text))
+        }
+        pub(crate) fn ws(&self) -> SyntaxToken<'_> {
+            self.0.syntax().first_child_or_token().unwrap().as_token().unwrap()
+        }
+    }
+
 }
 
 #[cfg(test)]