3 files changed, 138 insertions, 119 deletions

diff --git a/crates/ra_assists/src/assist_ctx.rs b/crates/ra_assists/src/assist_ctx.rs
index e80e35738..f46de61eb 100644
--- a/crates/ra_assists/src/assist_ctx.rs
+++ b/crates/ra_assists/src/assist_ctx.rs
@@ -161,6 +161,10 @@ impl AssistBuilder {
         self.target = Some(target)
     }
 
+    pub(crate) fn text_edit_builder(&mut self) -> &mut TextEditBuilder {
+        &mut self.edit
+    }
+
     fn build(self) -> AssistAction {
         AssistAction {
             edit: self.edit.finish(),
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)]
diff --git a/crates/ra_assists/src/fill_struct_fields.rs b/crates/ra_assists/src/fill_struct_fields.rs
index 663b4f669..ca128168a 100644
--- a/crates/ra_assists/src/fill_struct_fields.rs
+++ b/crates/ra_assists/src/fill_struct_fields.rs
@@ -1,94 +1,53 @@
-use std::fmt::Write;
-
 use hir::{AdtDef, db::HirDatabase};
 
 use ra_syntax::ast::{self, AstNode};
 
-use crate::{AssistCtx, Assist, AssistId};
+use crate::{AssistCtx, Assist, AssistId, ast_editor::{AstEditor, AstBuilder}};
 
 pub(crate) fn fill_struct_fields(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
     let struct_lit = ctx.node_at_offset::<ast::StructLit>()?;
-    let mut fsf = FillStructFields {
-        ctx: &mut ctx,
-        named_field_list: struct_lit.named_field_list()?,
-        struct_fields: vec![],
-        struct_lit,
-    };
-    fsf.evaluate_struct_def_fields()?;
-    if fsf.struct_lit_and_def_have_the_same_number_of_fields() {
-        return None;
-    }
-    fsf.remove_already_included_fields()?;
-    fsf.add_action()?;
-    ctx.build()
-}
-
-struct FillStructFields<'a, 'b: 'a, DB> {
-    ctx: &'a mut AssistCtx<'b, DB>,
-    named_field_list: &'a ast::NamedFieldList,
-    struct_fields: Vec<(String, String)>,
-    struct_lit: &'a ast::StructLit,
-}
+    let named_field_list = struct_lit.named_field_list()?;
 
-impl<DB> FillStructFields<'_, '_, DB>
-where
-    DB: HirDatabase,
-{
-    fn add_action(&mut self) -> Option<()> {
-        let named_field_list = self.named_field_list;
-        let struct_fields_string = self.struct_fields_string()?;
-        let struct_lit = self.struct_lit;
-        self.ctx.add_action(AssistId("fill_struct_fields"), "fill struct fields", |edit| {
-            edit.target(struct_lit.syntax().range());
-            edit.set_cursor(struct_lit.syntax().range().start());
-            edit.replace_node_and_indent(named_field_list.syntax(), struct_fields_string);
-        });
-        Some(())
-    }
-
-    fn struct_lit_and_def_have_the_same_number_of_fields(&self) -> bool {
-        self.named_field_list.fields().count() == self.struct_fields.len()
-    }
-
-    fn evaluate_struct_def_fields(&mut self) -> Option<()> {
-        let analyzer = hir::SourceAnalyzer::new(
-            self.ctx.db,
-            self.ctx.frange.file_id,
-            self.struct_lit.syntax(),
-            None,
-        );
-        let struct_lit_ty = analyzer.type_of(self.ctx.db, self.struct_lit.into())?;
+    // Collect all fields from struct definition
+    let mut fields = {
+        let analyzer =
+            hir::SourceAnalyzer::new(ctx.db, ctx.frange.file_id, struct_lit.syntax(), None);
+        let struct_lit_ty = analyzer.type_of(ctx.db, struct_lit.into())?;
         let struct_def = match struct_lit_ty.as_adt() {
             Some((AdtDef::Struct(s), _)) => s,
             _ => return None,
         };
-        self.struct_fields = struct_def
-            .fields(self.ctx.db)
-            .into_iter()
-            .map(|f| (f.name(self.ctx.db).to_string(), "()".into()))
-            .collect();
-        Some(())
-    }
+        struct_def.fields(ctx.db)
+    };
 
-    fn remove_already_included_fields(&mut self) -> Option<()> {
-        for ast_field in self.named_field_list.fields() {
-            let expr = ast_field.expr()?.syntax().text().to_string();
-            let name_from_ast = ast_field.name_ref()?.text().to_string();
-            if let Some(idx) = self.struct_fields.iter().position(|(n, _)| n == &name_from_ast) {
-                self.struct_fields[idx] = (name_from_ast, expr);
-            }
-        }
-        Some(())
+    // Filter out existing fields
+    for ast_field in named_field_list.fields() {
+        let name_from_ast = ast_field.name_ref()?.text().to_string();
+        fields.retain(|field| field.name(ctx.db).to_string() != name_from_ast);
+    }
+    if fields.is_empty() {
+        return None;
     }
 
-    fn struct_fields_string(&mut self) -> Option<String> {
-        let mut buf = String::from("{\n");
-        for (name, expr) in &self.struct_fields {
-            write!(&mut buf, "    {}: {},\n", name, expr).unwrap();
+    let db = ctx.db;
+    ctx.add_action(AssistId("fill_struct_fields"), "fill struct fields", |edit| {
+        let mut ast_editor = AstEditor::new(named_field_list);
+        if named_field_list.fields().count() == 0 && fields.len() > 2 {
+            ast_editor.make_multiline();
+        };
+
+        for field in fields {
+            let field =
+                AstBuilder::<ast::NamedField>::from_text(&format!("{}: ()", field.name(db)));
+            ast_editor.append_field(&field);
         }
-        buf.push_str("}");
-        Some(buf)
-    }
+
+        edit.target(struct_lit.syntax().range());
+        edit.set_cursor(struct_lit.syntax().range().start());
+
+        ast_editor.into_text_edit(edit.text_edit_builder());
+    });
+    ctx.build()
 }
 
 #[cfg(test)]
@@ -225,11 +184,11 @@ mod tests {
 
             fn main() {
                 let s = <|>S {
+                    c: (1, 2),
+                    e: "foo",
                     a: (),
                     b: (),
-                    c: (1, 2),
                     d: (),
-                    e: "foo",
                 }
             }
             "#,