use format_buf::format; use hir::{db::HirDatabase, FromSource}; use join_to_string::join; use ra_syntax::{ ast::{ self, AstNode, NameOwner, StructKind, TypeAscriptionOwner, TypeParamsOwner, VisibilityOwner, }, TextUnit, T, }; use std::fmt::Write; use crate::{Assist, AssistCtx, AssistId}; // Assist: add_new // // Adds a new inherent impl for a type. // // ``` // struct Ctx { // data: T,<|> // } // ``` // -> // ``` // struct Ctx { // data: T, // } // // impl Ctx { // fn new(data: T) -> Self { Self { data } } // } // // ``` pub(crate) fn add_new(ctx: AssistCtx) -> Option { let strukt = ctx.find_node_at_offset::()?; // We want to only apply this to non-union structs with named fields let field_list = match (strukt.kind(), strukt.is_union()) { (StructKind::Named(named), false) => named, _ => return None, }; // Return early if we've found an existing new fn let impl_block = find_struct_impl(&ctx, &strukt)?; ctx.add_assist(AssistId("add_new"), "add new fn", |edit| { edit.target(strukt.syntax().text_range()); let mut buf = String::with_capacity(512); if impl_block.is_some() { buf.push('\n'); } let vis = strukt.visibility().map(|v| format!("{} ", v.syntax())); let vis = vis.as_ref().map(String::as_str).unwrap_or(""); write!(&mut buf, " {}fn new(", vis).unwrap(); join(field_list.fields().map(|f| { format!( "{}: {}", f.name().unwrap().syntax().text(), f.ascribed_type().unwrap().syntax().text() ) })) .separator(", ") .to_buf(&mut buf); buf.push_str(") -> Self { Self {"); join(field_list.fields().map(|f| f.name().unwrap().syntax().text())) .separator(", ") .surround_with(" ", " ") .to_buf(&mut buf); buf.push_str("} }"); let (start_offset, end_offset) = if let Some(impl_block) = impl_block { buf.push('\n'); let start = impl_block .syntax() .descendants_with_tokens() .find(|t| t.kind() == T!['{']) .unwrap() .text_range() .end(); (start, TextUnit::from_usize(1)) } else { buf = generate_impl_text(&strukt, &buf); let start = strukt.syntax().text_range().end(); (start, TextUnit::from_usize(3)) }; edit.set_cursor(start_offset + TextUnit::of_str(&buf) - end_offset); edit.insert(start_offset, buf); }) } // Generates the surrounding `impl Type { }` including type and lifetime // parameters fn generate_impl_text(strukt: &ast::StructDef, code: &str) -> String { let type_params = strukt.type_param_list(); let mut buf = String::with_capacity(code.len()); buf.push_str("\n\nimpl"); if let Some(type_params) = &type_params { format!(buf, "{}", type_params.syntax()); } buf.push_str(" "); buf.push_str(strukt.name().unwrap().text().as_str()); if let Some(type_params) = type_params { let lifetime_params = type_params .lifetime_params() .filter_map(|it| it.lifetime_token()) .map(|it| it.text().clone()); let type_params = type_params.type_params().filter_map(|it| it.name()).map(|it| it.text().clone()); join(lifetime_params.chain(type_params)).surround_with("<", ">").to_buf(&mut buf); } format!(&mut buf, " {{\n{}\n}}\n", code); buf } // Uses a syntax-driven approach to find any impl blocks for the struct that // exist within the module/file // // Returns `None` if we've found an existing `new` fn // // FIXME: change the new fn checking to a more semantic approach when that's more // viable (e.g. we process proc macros, etc) fn find_struct_impl( ctx: &AssistCtx, strukt: &ast::StructDef, ) -> Option> { let db = ctx.db; let module = strukt.syntax().ancestors().find(|node| { ast::Module::can_cast(node.kind()) || ast::SourceFile::can_cast(node.kind()) })?; let struct_ty = { let src = hir::Source { file_id: ctx.frange.file_id.into(), ast: strukt.clone() }; hir::Struct::from_source(db, src).unwrap().ty(db) }; let mut found_new_fn = false; let block = module.descendants().filter_map(ast::ImplBlock::cast).find(|impl_blk| { if found_new_fn { return false; } let src = hir::Source { file_id: ctx.frange.file_id.into(), ast: impl_blk.clone() }; let blk = hir::ImplBlock::from_source(db, src).unwrap(); let same_ty = blk.target_ty(db) == struct_ty; let not_trait_impl = blk.target_trait(db).is_none(); found_new_fn = has_new_fn(impl_blk); same_ty && not_trait_impl }); if found_new_fn { None } else { Some(block) } } fn has_new_fn(imp: &ast::ImplBlock) -> bool { if let Some(il) = imp.item_list() { for item in il.impl_items() { if let ast::ImplItem::FnDef(f) = item { if f.name().unwrap().text().eq_ignore_ascii_case("new") { return true; } } } } false } #[cfg(test)] mod tests { use super::*; use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target}; #[test] #[rustfmt::skip] fn test_add_new() { // Check output of generation check_assist( add_new, "struct Foo {<|>}", "struct Foo {} impl Foo { fn new() -> Self { Self { } }<|> } ", ); check_assist( add_new, "struct Foo {<|>}", "struct Foo {} impl Foo { fn new() -> Self { Self { } }<|> } ", ); check_assist( add_new, "struct Foo<'a, T: Foo<'a>> {<|>}", "struct Foo<'a, T: Foo<'a>> {} impl<'a, T: Foo<'a>> Foo<'a, T> { fn new() -> Self { Self { } }<|> } ", ); check_assist( add_new, "struct Foo { baz: String <|>}", "struct Foo { baz: String } impl Foo { fn new(baz: String) -> Self { Self { baz } }<|> } ", ); check_assist( add_new, "struct Foo { baz: String, qux: Vec <|>}", "struct Foo { baz: String, qux: Vec } impl Foo { fn new(baz: String, qux: Vec) -> Self { Self { baz, qux } }<|> } ", ); // Check that visibility modifiers don't get brought in for fields check_assist( add_new, "struct Foo { pub baz: String, pub qux: Vec <|>}", "struct Foo { pub baz: String, pub qux: Vec } impl Foo { fn new(baz: String, qux: Vec) -> Self { Self { baz, qux } }<|> } ", ); // Check that it reuses existing impls check_assist( add_new, "struct Foo {<|>} impl Foo {} ", "struct Foo {} impl Foo { fn new() -> Self { Self { } }<|> } ", ); check_assist( add_new, "struct Foo {<|>} impl Foo { fn qux(&self) {} } ", "struct Foo {} impl Foo { fn new() -> Self { Self { } }<|> fn qux(&self) {} } ", ); check_assist( add_new, "struct Foo {<|>} impl Foo { fn qux(&self) {} fn baz() -> i32 { 5 } } ", "struct Foo {} impl Foo { fn new() -> Self { Self { } }<|> fn qux(&self) {} fn baz() -> i32 { 5 } } ", ); // Check visibility of new fn based on struct check_assist( add_new, "pub struct Foo {<|>}", "pub struct Foo {} impl Foo { pub fn new() -> Self { Self { } }<|> } ", ); check_assist( add_new, "pub(crate) struct Foo {<|>}", "pub(crate) struct Foo {} impl Foo { pub(crate) fn new() -> Self { Self { } }<|> } ", ); } #[test] fn add_new_not_applicable_if_fn_exists() { check_assist_not_applicable( add_new, " struct Foo {<|>} impl Foo { fn new() -> Self { Self } }", ); check_assist_not_applicable( add_new, " struct Foo {<|>} impl Foo { fn New() -> Self { Self } }", ); } #[test] fn add_new_target() { check_assist_target( add_new, " struct SomeThingIrrelevant; /// Has a lifetime parameter struct Foo<'a, T: Foo<'a>> {<|>} struct EvenMoreIrrelevant; ", "/// Has a lifetime parameter struct Foo<'a, T: Foo<'a>> {}", ); } }