crates/ide_assists/src/handlers/generate_delegate_methods.rs

   1 use hir::{self, HasCrate, HasSource, HasVisibility};
   2 use syntax::ast::{self, make, AstNode, HasGenericParams, HasName, HasVisibility as _};
   3
   4 use crate::{
   5     utils::{convert_param_list_to_arg_list, find_struct_impl, render_snippet, Cursor},
   6     AssistContext, AssistId, AssistKind, Assists, GroupLabel,
   7 };
   8 use syntax::ast::edit::AstNodeEdit;
   9
  10 // Assist: generate_delegate_methods
  11 //
  12 // Generate delegate methods.
  13 //
  14 // ```
  15 // struct Age(u8);
  16 // impl Age {
  17 //     fn age(&self) -> u8 {
  18 //         self.0
  19 //     }
  20 // }
  21 //
  22 // struct Person {
  23 //     ag$0e: Age,
  24 // }
  25 // ```
  26 // ->
  27 // ```
  28 // struct Age(u8);
  29 // impl Age {
  30 //     fn age(&self) -> u8 {
  31 //         self.0
  32 //     }
  33 // }
  34 //
  35 // struct Person {
  36 //     age: Age,
  37 // }
  38 //
  39 // impl Person {
  40 //     $0fn age(&self) -> u8 {
  41 //         self.age.age()
  42 //     }
  43 // }
  44 // ```
  45 pub(crate) fn generate_delegate_methods(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  46     let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
  47     let strukt_name = strukt.name()?;
  48     let current_module = ctx.sema.scope(strukt.syntax()).module()?;
  49
  50     let (field_name, field_ty) = match ctx.find_node_at_offset::<ast::RecordField>() {
  51         Some(field) => {
  52             let field_name = field.name()?;
  53             let field_ty = field.ty()?;
  54             (format!("{}", field_name), field_ty)
  55         }
  56         None => {
  57             let field = ctx.find_node_at_offset::<ast::TupleField>()?;
  58             let field_list = ctx.find_node_at_offset::<ast::TupleFieldList>()?;
  59             let field_list_index = field_list.fields().position(|it| it == field)?;
  60             let field_ty = field.ty()?;
  61             (format!("{}", field_list_index), field_ty)
  62         }
  63     };
  64
  65     let sema_field_ty = ctx.sema.resolve_type(&field_ty)?;
  66     let krate = sema_field_ty.krate(ctx.db());
  67     let mut methods = vec![];
  68     sema_field_ty.iterate_assoc_items(ctx.db(), krate, |item| {
  69         if let hir::AssocItem::Function(f) = item {
  70             if f.self_param(ctx.db()).is_some() && f.is_visible_from(ctx.db(), current_module) {
  71                 methods.push(f)
  72             }
  73         }
  74         Option::<()>::None
  75     });
  76
  77     let target = field_ty.syntax().text_range();
  78     for method in methods {
  79         let adt = ast::Adt::Struct(strukt.clone());
  80         let name = method.name(ctx.db()).to_string();
  81         let impl_def = find_struct_impl(ctx, &adt, &name).flatten();
  82         acc.add_group(
  83             &GroupLabel("Generate delegate methods…".to_owned()),
  84             AssistId("generate_delegate_methods", AssistKind::Generate),
  85             format!("Generate delegate for `{}.{}()`", field_name, method.name(ctx.db())),
  86             target,
  87             |builder| {
  88                 // Create the function
  89                 let method_source = match method.source(ctx.db()) {
  90                     Some(source) => source.value,
  91                     None => return,
  92                 };
  93                 let method_name = method.name(ctx.db());
  94                 let vis = method_source.visibility();
  95                 let name = make::name(&method.name(ctx.db()).to_string());
  96                 let params =
  97                     method_source.param_list().unwrap_or_else(|| make::param_list(None, []));
  98                 let type_params = method_source.generic_param_list();
  99                 let arg_list = match method_source.param_list() {
 100                     Some(list) => convert_param_list_to_arg_list(list),
 101                     None => make::arg_list([]),
 102                 };
 103                 let tail_expr = make::expr_method_call(
 104                     make::ext::field_from_idents(["self", &field_name]).unwrap(), // This unwrap is ok because we have at least 1 arg in the list
 105                     make::name_ref(&method_name.to_string()),
 106                     arg_list,
 107                 );
 108                 let body = make::block_expr([], Some(tail_expr));
 109                 let ret_type = method_source.ret_type();
 110                 let is_async = method_source.async_token().is_some();
 111                 let f = make::fn_(vis, name, type_params, params, body, ret_type, is_async)
 112                     .indent(ast::edit::IndentLevel(1))
 113                     .clone_for_update();
 114
 115                 let cursor = Cursor::Before(f.syntax());
 116
 117                 // Create or update an impl block, attach the function to it,
 118                 // then insert into our code.
 119                 match impl_def {
 120                     Some(impl_def) => {
 121                         // Remember where in our source our `impl` block lives.
 122                         let impl_def = impl_def.clone_for_update();
 123                         let old_range = impl_def.syntax().text_range();
 124
 125                         // Attach the function to the impl block
 126                         let assoc_items = impl_def.get_or_create_assoc_item_list();
 127                         assoc_items.add_item(f.clone().into());
 128
 129                         // Update the impl block.
 130                         match ctx.config.snippet_cap {
 131                             Some(cap) => {
 132                                 let snippet = render_snippet(cap, impl_def.syntax(), cursor);
 133                                 builder.replace_snippet(cap, old_range, snippet);
 134                             }
 135                             None => {
 136                                 builder.replace(old_range, impl_def.syntax().to_string());
 137                             }
 138                         }
 139                     }
 140                     None => {
 141                         // Attach the function to the impl block
 142                         let name = &strukt_name.to_string();
 143                         let params = strukt.generic_param_list();
 144                         let ty_params = params.clone();
 145                         let impl_def = make::impl_(make::ext::ident_path(name), params, ty_params)
 146                             .clone_for_update();
 147                         let assoc_items = impl_def.get_or_create_assoc_item_list();
 148                         assoc_items.add_item(f.clone().into());
 149
 150                         // Insert the impl block.
 151                         match ctx.config.snippet_cap {
 152                             Some(cap) => {
 153                                 let offset = strukt.syntax().text_range().end();
 154                                 let snippet = render_snippet(cap, impl_def.syntax(), cursor);
 155                                 let snippet = format!("\n\n{}", snippet);
 156                                 builder.insert_snippet(cap, offset, snippet);
 157                             }
 158                             None => {
 159                                 let offset = strukt.syntax().text_range().end();
 160                                 let snippet = format!("\n\n{}", impl_def.syntax().to_string());
 161                                 builder.insert(offset, snippet);
 162                             }
 163                         }
 164                     }
 165                 }
 166             },
 167         )?;
 168     }
 169     Some(())
 170 }
 171
 172 #[cfg(test)]
 173 mod tests {
 174     use crate::tests::{check_assist, check_assist_not_applicable};
 175
 176     use super::*;
 177
 178     #[test]
 179     fn test_generate_delegate_create_impl_block() {
 180         check_assist(
 181             generate_delegate_methods,
 182             r#"
 183 struct Age(u8);
 184 impl Age {
 185     fn age(&self) -> u8 {
 186         self.0
 187     }
 188 }
 189
 190 struct Person {
 191     ag$0e: Age,
 192 }"#,
 193             r#"
 194 struct Age(u8);
 195 impl Age {
 196     fn age(&self) -> u8 {
 197         self.0
 198     }
 199 }
 200
 201 struct Person {
 202     age: Age,
 203 }
 204
 205 impl Person {
 206     $0fn age(&self) -> u8 {
 207         self.age.age()
 208     }
 209 }"#,
 210         );
 211     }
 212
 213     #[test]
 214     fn test_generate_delegate_update_impl_block() {
 215         check_assist(
 216             generate_delegate_methods,
 217             r#"
 218 struct Age(u8);
 219 impl Age {
 220     fn age(&self) -> u8 {
 221         self.0
 222     }
 223 }
 224
 225 struct Person {
 226     ag$0e: Age,
 227 }
 228
 229 impl Person {}"#,
 230             r#"
 231 struct Age(u8);
 232 impl Age {
 233     fn age(&self) -> u8 {
 234         self.0
 235     }
 236 }
 237
 238 struct Person {
 239     age: Age,
 240 }
 241
 242 impl Person {
 243     $0fn age(&self) -> u8 {
 244         self.age.age()
 245     }
 246 }"#,
 247         );
 248     }
 249
 250     #[test]
 251     fn test_generate_delegate_tuple_struct() {
 252         check_assist(
 253             generate_delegate_methods,
 254             r#"
 255 struct Age(u8);
 256 impl Age {
 257     fn age(&self) -> u8 {
 258         self.0
 259     }
 260 }
 261
 262 struct Person(A$0ge);"#,
 263             r#"
 264 struct Age(u8);
 265 impl Age {
 266     fn age(&self) -> u8 {
 267         self.0
 268     }
 269 }
 270
 271 struct Person(Age);
 272
 273 impl Person {
 274     $0fn age(&self) -> u8 {
 275         self.0.age()
 276     }
 277 }"#,
 278         );
 279     }
 280
 281     #[test]
 282     fn test_generate_delegate_enable_all_attributes() {
 283         check_assist(
 284             generate_delegate_methods,
 285             r#"
 286 struct Age<T>(T);
 287 impl<T> Age<T> {
 288     pub(crate) async fn age<J, 'a>(&'a mut self, ty: T, arg: J) -> T {
 289         self.0
 290     }
 291 }
 292
 293 struct Person<T> {
 294     ag$0e: Age<T>,
 295 }"#,
 296             r#"
 297 struct Age<T>(T);
 298 impl<T> Age<T> {
 299     pub(crate) async fn age<J, 'a>(&'a mut self, ty: T, arg: J) -> T {
 300         self.0
 301     }
 302 }
 303
 304 struct Person<T> {
 305     age: Age<T>,
 306 }
 307
 308 impl<T> Person<T> {
 309     $0pub(crate) async fn age<J, 'a>(&'a mut self, ty: T, arg: J) -> T {
 310         self.age.age(ty, arg)
 311     }
 312 }"#,
 313         );
 314     }
 315
 316     #[test]
 317     fn test_generate_delegate_visibility() {
 318         check_assist_not_applicable(
 319             generate_delegate_methods,
 320             r#"
 321 mod m {
 322     pub struct Age(u8);
 323     impl Age {
 324         fn age(&self) -> u8 {
 325             self.0
 326         }
 327     }
 328 }
 329
 330 struct Person {
 331     ag$0e: m::Age,
 332 }"#,
 333         )
 334     }
 335 }