1 files changed, 122 insertions, 33 deletions

diff --git a/crates/ra_assists/src/fill_struct_fields.rs b/crates/ra_assists/src/fill_struct_fields.rs
index 90ce9c577..6e61a58fc 100644
--- a/crates/ra_assists/src/fill_struct_fields.rs
+++ b/crates/ra_assists/src/fill_struct_fields.rs
@@ -1,7 +1,6 @@
 use std::fmt::Write;
 
-use hir::{AdtDef, Ty, source_binder};
-use hir::db::HirDatabase;
+use hir::{AdtDef, Ty, db::HirDatabase, source_binder::function_from_child_node};
 
 use ra_syntax::ast::{self, AstNode};
 
@@ -9,41 +8,89 @@ use crate::{AssistCtx, Assist, AssistId};
 
 pub(crate) fn fill_struct_fields(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
     let struct_lit = ctx.node_at_offset::<ast::StructLit>()?;
-    let named_field_list = struct_lit.named_field_list()?;
-
-    // If we already have existing struct fields, don't provide the assist.
-    if named_field_list.fields().count() > 0 {
+    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()
+}
 
-    let function =
-        source_binder::function_from_child_node(ctx.db, ctx.frange.file_id, struct_lit.syntax())?;
+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 infer_result = function.infer(ctx.db);
-    let source_map = function.body_source_map(ctx.db);
-    let node_expr = source_map.node_expr(struct_lit.into())?;
-    let struct_lit_ty = infer_result[node_expr].clone();
-    let struct_def = match struct_lit_ty {
-        Ty::Adt { def_id: AdtDef::Struct(s), .. } => s,
-        _ => return None,
-    };
+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(())
+    }
 
-    let db = ctx.db;
-    ctx.add_action(AssistId("fill_struct_fields"), "fill struct fields", |edit| {
+    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 function = function_from_child_node(
+            self.ctx.db,
+            self.ctx.frange.file_id,
+            self.struct_lit.syntax(),
+        )?;
+        let infer_result = function.infer(self.ctx.db);
+        let source_map = function.body_source_map(self.ctx.db);
+        let node_expr = source_map.node_expr(self.struct_lit.into())?;
+        let struct_lit_ty = infer_result[node_expr].clone();
+        let struct_def = match struct_lit_ty {
+            Ty::Adt { def_id: 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(())
+    }
+
+    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(())
+    }
+
+    fn struct_fields_string(&mut self) -> Option<String> {
         let mut buf = String::from("{\n");
-        let struct_fields = struct_def.fields(db);
-        for field in struct_fields {
-            let field_name = field.name(db).to_string();
-            write!(&mut buf, "    {}: (),\n", field_name).unwrap();
+        for (name, expr) in &self.struct_fields {
+            write!(&mut buf, "    {}: {},\n", name, expr).unwrap();
         }
         buf.push_str("}");
-
-        edit.target(struct_lit.syntax().range());
-        edit.set_cursor(struct_lit.syntax().range().start());
-        edit.replace_node_and_indent(named_field_list.syntax(), buf);
-    });
-
-    ctx.build()
+        Some(buf)
+    }
 }
 
 #[cfg(test)]
@@ -62,7 +109,7 @@ mod tests {
                 b: String,
                 c: (i32, i32),
                 d: D,
-                r: &'a str,
+                e: &'a str,
             }
 
             fn main() {
@@ -75,7 +122,7 @@ mod tests {
                 b: String,
                 c: (i32, i32),
                 d: D,
-                r: &'a str,
+                e: &'a str,
             }
 
             fn main() {
@@ -84,7 +131,7 @@ mod tests {
                     b: (),
                     c: (),
                     d: (),
-                    r: (),
+                    e: (),
                 }
             }
             "#,
@@ -101,7 +148,7 @@ mod tests {
                 b: String,
                 c: (i32, i32),
                 d: D,
-                r: &'a str,
+                e: &'a str,
             }
 
             fn main() {
@@ -148,4 +195,46 @@ mod tests {
             "#,
         );
     }
+
+    #[test]
+    fn fill_struct_fields_partial() {
+        check_assist(
+            fill_struct_fields,
+            r#"
+            struct S<'a, D> {
+                a: u32,
+                b: String,
+                c: (i32, i32),
+                d: D,
+                e: &'a str,
+            }
+
+            fn main() {
+                let s = S {
+                    c: (1, 2),
+                    e: "foo",<|>
+                }
+            }
+            "#,
+            r#"
+            struct S<'a, D> {
+                a: u32,
+                b: String,
+                c: (i32, i32),
+                d: D,
+                e: &'a str,
+            }
+
+            fn main() {
+                let s = <|>S {
+                    a: (),
+                    b: (),
+                    c: (1, 2),
+                    d: (),
+                    e: "foo",
+                }
+            }
+            "#,
+        );
+    }
 }