Only replace NamedFieldList and add test for preserving Self

1 files changed, 45 insertions, 13 deletions

diff --git a/crates/ra_assists/src/fill_struct_fields.rs b/crates/ra_assists/src/fill_struct_fields.rs
index 938322d1d..90ce9c577 100644
--- a/crates/ra_assists/src/fill_struct_fields.rs
+++ b/crates/ra_assists/src/fill_struct_fields.rs
@@ -3,39 +3,34 @@ use std::fmt::Write;
 use hir::{AdtDef, Ty, source_binder};
 use hir::db::HirDatabase;
 
-use ra_syntax::ast::{self, AstNode, Expr};
+use ra_syntax::ast::{self, AstNode};
 
 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.
-    match struct_lit.named_field_list() {
-        Some(named_field_list) if named_field_list.fields().count() > 0 => {
-            return None;
-        }
-        _ => {}
+    if named_field_list.fields().count() > 0 {
+        return None;
     }
 
-    let expr: &Expr = struct_lit.into();
     let function =
         source_binder::function_from_child_node(ctx.db, ctx.frange.file_id, struct_lit.syntax())?;
 
     let infer_result = function.infer(ctx.db);
     let source_map = function.body_source_map(ctx.db);
-    let node_expr = source_map.node_expr(expr)?;
+    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,
     };
 
-    let struct_name = struct_def.name(ctx.db)?;
     let db = ctx.db;
-
     ctx.add_action(AssistId("fill_struct_fields"), "fill struct fields", |edit| {
-        let mut buf = format!("{} {{\n", struct_name);
+        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();
@@ -44,8 +39,8 @@ pub(crate) fn fill_struct_fields(mut ctx: AssistCtx<impl HirDatabase>) -> Option
         buf.push_str("}");
 
         edit.target(struct_lit.syntax().range());
-        edit.set_cursor(expr.syntax().range().start());
-        edit.replace_node_and_indent(struct_lit.syntax(), buf);
+        edit.set_cursor(struct_lit.syntax().range().start());
+        edit.replace_node_and_indent(named_field_list.syntax(), buf);
     });
 
     ctx.build()
@@ -116,4 +111,41 @@ mod tests {
             "S {}",
         );
     }
+
+    #[test]
+    fn fill_struct_fields_preserve_self() {
+        check_assist(
+            fill_struct_fields,
+            r#"
+            struct Foo {
+                foo: u8,
+                bar: String,
+                baz: i128,
+            }
+
+            impl Foo {
+                pub fn new() -> Self {
+                    Self <|>{}
+                }
+            }
+            "#,
+            r#"
+            struct Foo {
+                foo: u8,
+                bar: String,
+                baz: i128,
+            }
+
+            impl Foo {
+                pub fn new() -> Self {
+                    <|>Self {
+                        foo: (),
+                        bar: (),
+                        baz: (),
+                    }
+                }
+            }
+            "#,
+        );
+    }
 }