use hir::{self, HasCrate, HasSource, HirDisplay}; use syntax::ast::{self, make, AstNode, HasGenericParams, HasName, HasVisibility}; use crate::{ utils::{find_struct_impl, render_snippet, Cursor}, AssistContext, AssistId, AssistKind, Assists, GroupLabel, }; use syntax::ast::edit::AstNodeEdit; // Assist: generate_delegate_methods // // Generate delegate methods. // // ``` // struct Age(u8); // impl Age { // fn age(&self) -> u8 { // self.0 // } // } // // struct Person { // ag$0e: Age, // } // ``` // -> // ``` // struct Age(u8); // impl Age { // fn age(&self) -> u8 { // self.0 // } // } // // struct Person { // age: Age, // } // // impl Person { // $0fn age(&self) -> u8 { // self.age.age() // } // } // ``` pub(crate) fn generate_delegate_methods(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { let cap = ctx.config.snippet_cap?; let strukt = ctx.find_node_at_offset::()?; let strukt_name = strukt.name()?; let (field_name, field_ty) = match ctx.find_node_at_offset::() { Some(field) => { let field_name = field.name()?; let field_ty = field.ty()?; (format!("{}", field_name), field_ty) } None => { let field = ctx.find_node_at_offset::()?; let field_list = ctx.find_node_at_offset::()?; let field_list_index = field_list.syntax().children().into_iter().position(|s| &s == field.syntax())?; let field_ty = field.ty()?; (format!("{}", field_list_index), field_ty) } }; let sema_field_ty = ctx.sema.resolve_type(&field_ty)?; let krate = sema_field_ty.krate(ctx.db()); let mut methods = vec![]; sema_field_ty.iterate_assoc_items(ctx.db(), krate, |item| { if let hir::AssocItem::Function(f) = item { if f.self_param(ctx.db()).is_some() { methods.push(f) } } Some(()) }); let target = field_ty.syntax().text_range(); for method in methods { let impl_def = find_struct_impl( ctx, &ast::Adt::Struct(strukt.clone()), &method.name(ctx.db()).to_string(), )?; acc.add_group( &GroupLabel("Generate delegate methods…".to_owned()), AssistId("generate_delegate_methods", AssistKind::Generate), format!("Generate delegate for `{}.{}()`", field_name, method.name(ctx.db())), target, |builder| { // Create the function let method_source = match method.source(ctx.db()) { Some(source) => source.value, None => return, }; let method_name = method.name(ctx.db()); let vis = method_source.visibility(); let name = make::name(&method.name(ctx.db()).to_string()); let params = method_source.param_list().unwrap_or_else(|| make::param_list(None, [])); let tail_expr = make::expr_method_call( make::ext::field_from_idents(["self", &field_name]).unwrap(), // This unwrap is ok because we have at least 1 arg in the list make::name_ref(&method_name.to_string()), make::arg_list([]), ); let type_params = method_source.generic_param_list(); let body = make::block_expr([], Some(tail_expr)); let ret_type = method.ret_type(ctx.db()); let ret_type = if ret_type.is_unknown() { Some(make::ret_type(make::ty_placeholder())) } else { let ret_type = &ret_type.display(ctx.db()).to_string(); Some(make::ret_type(make::ty(ret_type))) }; let is_async = method_source.async_token().is_some(); let f = make::fn_(vis, name, type_params, params, body, ret_type, is_async) .indent(ast::edit::IndentLevel(1)) .clone_for_update(); let cursor = Cursor::Before(f.syntax()); // Create or update an impl block, attach the function to it, // then insert into our code. match impl_def { Some(impl_def) => { // Remember where in our source our `impl` block lives. let impl_def = impl_def.clone_for_update(); let old_range = impl_def.syntax().text_range(); // Attach the function to the impl block let assoc_items = impl_def.get_or_create_assoc_item_list(); assoc_items.add_item(f.clone().into()); // Update the impl block. let snippet = render_snippet(cap, impl_def.syntax(), cursor); builder.replace_snippet(cap, old_range, snippet); } None => { // Attach the function to the impl block let name = &strukt_name.to_string(); let params = strukt.generic_param_list(); let ty_params = params.clone(); let impl_def = make::impl_(make::ext::ident_path(name), params, ty_params) .clone_for_update(); let assoc_items = impl_def.get_or_create_assoc_item_list(); assoc_items.add_item(f.clone().into()); // Insert the impl block. let offset = strukt.syntax().text_range().end(); let snippet = render_snippet(cap, impl_def.syntax(), cursor); let snippet = format!("\n\n{}", snippet); builder.insert_snippet(cap, offset, snippet); } } }, )?; } Some(()) } #[cfg(test)] mod tests { use crate::tests::check_assist; use super::*; #[test] fn test_generate_delegate_create_impl_block() { check_assist( generate_delegate_methods, r#" struct Age(u8); impl Age { fn age(&self) -> u8 { self.0 } } struct Person { ag$0e: Age, }"#, r#" struct Age(u8); impl Age { fn age(&self) -> u8 { self.0 } } struct Person { age: Age, } impl Person { $0fn age(&self) -> u8 { self.age.age() } }"#, ); } #[test] fn test_generate_delegate_update_impl_block() { check_assist( generate_delegate_methods, r#" struct Age(u8); impl Age { fn age(&self) -> u8 { self.0 } } struct Person { ag$0e: Age, } impl Person {}"#, r#" struct Age(u8); impl Age { fn age(&self) -> u8 { self.0 } } struct Person { age: Age, } impl Person { $0fn age(&self) -> u8 { self.age.age() } }"#, ); } #[test] fn test_generate_delegate_tuple_struct() { check_assist( generate_delegate_methods, r#" struct Age(u8); impl Age { fn age(&self) -> u8 { self.0 } } struct Person(A$0ge);"#, r#" struct Age(u8); impl Age { fn age(&self) -> u8 { self.0 } } struct Person(Age); impl Person { $0fn age(&self) -> u8 { self.0.age() } }"#, ); } #[test] fn test_generate_delegate_enable_all_attributes() { check_assist( generate_delegate_methods, r#" struct Age(T); impl Age { pub(crate) async fn age(&'a mut self, ty: T, arg: J) -> T { self.0 } } struct Person { ag$0e: Age, }"#, r#" struct Age(T); impl Age { pub(crate) async fn age(&'a mut self, ty: T, arg: J) -> T { self.0 } } struct Person { age: Age, } impl Person { $0pub(crate) async fn age(&'a mut self, ty: T, arg: J) -> _ { self.age.age() } }"#, ); } }