src/tools/rust-analyzer/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, target) = match ctx.find_node_at_offset::<ast::RecordField>() {
  51         Some(field) => {
  52             let field_name = field.name()?;
  53             let field_ty = field.ty()?;
  54             (field_name.to_string(), field_ty, field.syntax().text_range())
  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             (field_list_index.to_string(), field_ty, field.syntax().text_range())
  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     for method in methods {
  78         let adt = ast::Adt::Struct(strukt.clone());
  79         let name = method.name(ctx.db()).to_string();
  80         let impl_def = find_struct_impl(ctx, &adt, &[name]).flatten();
  81         acc.add_group(
  82             &GroupLabel("Generate delegate methods…".to_owned()),
  83             AssistId("generate_delegate_methods", AssistKind::Generate),
  84             format!("Generate delegate for `{field_name}.{}()`", method.name(ctx.db())),
  85             target,
  86             |builder| {
  87                 // Create the function
  88                 let method_source = match method.source(ctx.db()) {
  89                     Some(source) => source.value,
  90                     None => return,
  91                 };
  92                 let method_name = method.name(ctx.db());
  93                 let vis = method_source.visibility();
  94                 let name = make::name(&method.name(ctx.db()).to_string());
  95                 let params =
  96                     method_source.param_list().unwrap_or_else(|| make::param_list(None, []));
  97                 let type_params = method_source.generic_param_list();
  98                 let arg_list = match method_source.param_list() {
  99                     Some(list) => convert_param_list_to_arg_list(list),
 100                     None => make::arg_list([]),
 101                 };
 102                 let tail_expr = make::expr_method_call(
 103                     make::ext::field_from_idents(["self", &field_name]).unwrap(), // This unwrap is ok because we have at least 1 arg in the list
 104                     make::name_ref(&method_name.to_string()),
 105                     arg_list,
 106                 );
 107                 let ret_type = method_source.ret_type();
 108                 let is_async = method_source.async_token().is_some();
 109                 let tail_expr_finished =
 110                     if is_async { make::expr_await(tail_expr) } else { tail_expr };
 111                 let body = make::block_expr([], Some(tail_expr_finished));
 112                 let f = make::fn_(vis, name, type_params, params, body, ret_type, is_async)
 113                     .indent(ast::edit::IndentLevel(1))
 114                     .clone_for_update();
 115
 116                 let cursor = Cursor::Before(f.syntax());
 117
 118                 // Create or update an impl block, attach the function to it,
 119                 // then insert into our code.
 120                 match impl_def {
 121                     Some(impl_def) => {
 122                         // Remember where in our source our `impl` block lives.
 123                         let impl_def = impl_def.clone_for_update();
 124                         let old_range = impl_def.syntax().text_range();
 125
 126                         // Attach the function to the impl block
 127                         let assoc_items = impl_def.get_or_create_assoc_item_list();
 128                         assoc_items.add_item(f.clone().into());
 129
 130                         // Update the impl block.
 131                         match ctx.config.snippet_cap {
 132                             Some(cap) => {
 133                                 let snippet = render_snippet(cap, impl_def.syntax(), cursor);
 134                                 builder.replace_snippet(cap, old_range, snippet);
 135                             }
 136                             None => {
 137                                 builder.replace(old_range, impl_def.syntax().to_string());
 138                             }
 139                         }
 140                     }
 141                     None => {
 142                         // Attach the function to the impl block
 143                         let name = &strukt_name.to_string();
 144                         let params = strukt.generic_param_list();
 145                         let ty_params = params.clone();
 146                         let impl_def = make::impl_(make::ext::ident_path(name), params, ty_params)
 147                             .clone_for_update();
 148                         let assoc_items = impl_def.get_or_create_assoc_item_list();
 149                         assoc_items.add_item(f.clone().into());
 150
 151                         // Insert the impl block.
 152                         match ctx.config.snippet_cap {
 153                             Some(cap) => {
 154                                 let offset = strukt.syntax().text_range().end();
 155                                 let snippet = render_snippet(cap, impl_def.syntax(), cursor);
 156                                 let snippet = format!("\n\n{snippet}");
 157                                 builder.insert_snippet(cap, offset, snippet);
 158                             }
 159                             None => {
 160                                 let offset = strukt.syntax().text_range().end();
 161                                 let snippet = format!("\n\n{}", impl_def.syntax());
 162                                 builder.insert(offset, snippet);
 163                             }
 164                         }
 165                     }
 166                 }
 167             },
 168         )?;
 169     }
 170     Some(())
 171 }
 172
 173 #[cfg(test)]
 174 mod tests {
 175     use crate::tests::{check_assist, check_assist_not_applicable};
 176
 177     use super::*;
 178
 179     #[test]
 180     fn test_generate_delegate_create_impl_block() {
 181         check_assist(
 182             generate_delegate_methods,
 183             r#"
 184 struct Age(u8);
 185 impl Age {
 186     fn age(&self) -> u8 {
 187         self.0
 188     }
 189 }
 190
 191 struct Person {
 192     ag$0e: Age,
 193 }"#,
 194             r#"
 195 struct Age(u8);
 196 impl Age {
 197     fn age(&self) -> u8 {
 198         self.0
 199     }
 200 }
 201
 202 struct Person {
 203     age: Age,
 204 }
 205
 206 impl Person {
 207     $0fn age(&self) -> u8 {
 208         self.age.age()
 209     }
 210 }"#,
 211         );
 212     }
 213
 214     #[test]
 215     fn test_generate_delegate_update_impl_block() {
 216         check_assist(
 217             generate_delegate_methods,
 218             r#"
 219 struct Age(u8);
 220 impl Age {
 221     fn age(&self) -> u8 {
 222         self.0
 223     }
 224 }
 225
 226 struct Person {
 227     ag$0e: Age,
 228 }
 229
 230 impl Person {}"#,
 231             r#"
 232 struct Age(u8);
 233 impl Age {
 234     fn age(&self) -> u8 {
 235         self.0
 236     }
 237 }
 238
 239 struct Person {
 240     age: Age,
 241 }
 242
 243 impl Person {
 244     $0fn age(&self) -> u8 {
 245         self.age.age()
 246     }
 247 }"#,
 248         );
 249     }
 250
 251     #[test]
 252     fn test_generate_delegate_tuple_struct() {
 253         check_assist(
 254             generate_delegate_methods,
 255             r#"
 256 struct Age(u8);
 257 impl Age {
 258     fn age(&self) -> u8 {
 259         self.0
 260     }
 261 }
 262
 263 struct Person(A$0ge);"#,
 264             r#"
 265 struct Age(u8);
 266 impl Age {
 267     fn age(&self) -> u8 {
 268         self.0
 269     }
 270 }
 271
 272 struct Person(Age);
 273
 274 impl Person {
 275     $0fn age(&self) -> u8 {
 276         self.0.age()
 277     }
 278 }"#,
 279         );
 280     }
 281
 282     #[test]
 283     fn test_generate_delegate_enable_all_attributes() {
 284         check_assist(
 285             generate_delegate_methods,
 286             r#"
 287 struct Age<T>(T);
 288 impl<T> Age<T> {
 289     pub(crate) async fn age<J, 'a>(&'a mut self, ty: T, arg: J) -> T {
 290         self.0
 291     }
 292 }
 293
 294 struct Person<T> {
 295     ag$0e: Age<T>,
 296 }"#,
 297             r#"
 298 struct Age<T>(T);
 299 impl<T> Age<T> {
 300     pub(crate) async fn age<J, 'a>(&'a mut self, ty: T, arg: J) -> T {
 301         self.0
 302     }
 303 }
 304
 305 struct Person<T> {
 306     age: Age<T>,
 307 }
 308
 309 impl<T> Person<T> {
 310     $0pub(crate) async fn age<J, 'a>(&'a mut self, ty: T, arg: J) -> T {
 311         self.age.age(ty, arg).await
 312     }
 313 }"#,
 314         );
 315     }
 316
 317     #[test]
 318     fn test_generate_delegate_visibility() {
 319         check_assist_not_applicable(
 320             generate_delegate_methods,
 321             r#"
 322 mod m {
 323     pub struct Age(u8);
 324     impl Age {
 325         fn age(&self) -> u8 {
 326             self.0
 327         }
 328     }
 329 }
 330
 331 struct Person {
 332     ag$0e: m::Age,
 333 }"#,
 334         )
 335     }
 336 }