diff --git a/crates/ra_assists/src/fill_struct_fields.rs b/crates/ra_assists/src/fill_struct_fields.rs index 90ce9c5775f..6e61a58fced 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) -> Option { let struct_lit = ctx.node_at_offset::()?; - 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 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 { 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", + } + } + "#, + ); + } }