src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_getter.rs

   1 use ide_db::famous_defs::FamousDefs;
   2 use stdx::{format_to, to_lower_snake_case};
   3 use syntax::{
   4     ast::{self, AstNode, HasName, HasVisibility},
   5     TextRange,
   6 };
   7
   8 use crate::{
   9     utils::{convert_reference_type, find_impl_block_end, find_struct_impl, generate_impl_text},
  10     AssistContext, AssistId, AssistKind, Assists, GroupLabel,
  11 };
  12
  13 // Assist: generate_getter
  14 //
  15 // Generate a getter method.
  16 //
  17 // ```
  18 // # //- minicore: as_ref
  19 // # pub struct String;
  20 // # impl AsRef<str> for String {
  21 // #     fn as_ref(&self) -> &str {
  22 // #         ""
  23 // #     }
  24 // # }
  25 // #
  26 // struct Person {
  27 //     nam$0e: String,
  28 // }
  29 // ```
  30 // ->
  31 // ```
  32 // # pub struct String;
  33 // # impl AsRef<str> for String {
  34 // #     fn as_ref(&self) -> &str {
  35 // #         ""
  36 // #     }
  37 // # }
  38 // #
  39 // struct Person {
  40 //     name: String,
  41 // }
  42 //
  43 // impl Person {
  44 //     fn $0name(&self) -> &str {
  45 //         self.name.as_ref()
  46 //     }
  47 // }
  48 // ```
  49 pub(crate) fn generate_getter(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  50     generate_getter_impl(acc, ctx, false)
  51 }
  52
  53 // Assist: generate_getter_mut
  54 //
  55 // Generate a mut getter method.
  56 //
  57 // ```
  58 // struct Person {
  59 //     nam$0e: String,
  60 // }
  61 // ```
  62 // ->
  63 // ```
  64 // struct Person {
  65 //     name: String,
  66 // }
  67 //
  68 // impl Person {
  69 //     fn $0name_mut(&mut self) -> &mut String {
  70 //         &mut self.name
  71 //     }
  72 // }
  73 // ```
  74 pub(crate) fn generate_getter_mut(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  75     generate_getter_impl(acc, ctx, true)
  76 }
  77
  78 #[derive(Clone, Debug)]
  79 struct RecordFieldInfo {
  80     field_name: syntax::ast::Name,
  81     field_ty: syntax::ast::Type,
  82     fn_name: String,
  83     target: TextRange,
  84 }
  85
  86 struct GetterInfo {
  87     impl_def: Option<ast::Impl>,
  88     strukt: ast::Struct,
  89     mutable: bool,
  90 }
  91
  92 pub(crate) fn generate_getter_impl(
  93     acc: &mut Assists,
  94     ctx: &AssistContext<'_>,
  95     mutable: bool,
  96 ) -> Option<()> {
  97     // This if condition denotes two modes this assist can work in:
  98     // - First is acting upon selection of record fields
  99     // - Next is acting upon a single record field
 100     //
 101     // This is the only part where implementation diverges a bit,
 102     // subsequent code is generic for both of these modes
 103
 104     let (strukt, info_of_record_fields, fn_names) = if !ctx.has_empty_selection() {
 105         // Selection Mode
 106         let node = ctx.covering_element();
 107
 108         let node = match node {
 109             syntax::NodeOrToken::Node(n) => n,
 110             syntax::NodeOrToken::Token(t) => t.parent()?,
 111         };
 112
 113         let parent_struct = node.ancestors().find_map(ast::Struct::cast)?;
 114
 115         let (info_of_record_fields, field_names) =
 116             extract_and_parse_record_fields(&parent_struct, ctx.selection_trimmed(), mutable)?;
 117
 118         (parent_struct, info_of_record_fields, field_names)
 119     } else {
 120         // Single Record Field mode
 121         let strukt = ctx.find_node_at_offset::<ast::Struct>()?;
 122         let field = ctx.find_node_at_offset::<ast::RecordField>()?;
 123
 124         let record_field_info = parse_record_field(field, mutable)?;
 125
 126         let fn_name = record_field_info.fn_name.clone();
 127
 128         (strukt, vec![record_field_info], vec![fn_name])
 129     };
 130
 131     // No record fields to do work on :(
 132     if info_of_record_fields.len() == 0 {
 133         return None;
 134     }
 135
 136     let impl_def = find_struct_impl(ctx, &ast::Adt::Struct(strukt.clone()), &fn_names)?;
 137
 138     let (id, label) = if mutable {
 139         ("generate_getter_mut", "Generate a mut getter method")
 140     } else {
 141         ("generate_getter", "Generate a getter method")
 142     };
 143
 144     // Computing collective text range of all record fields in selected region
 145     let target: TextRange = info_of_record_fields
 146         .iter()
 147         .map(|record_field_info| record_field_info.target)
 148         .reduce(|acc, target| acc.cover(target))?;
 149
 150     let getter_info = GetterInfo { impl_def, strukt, mutable };
 151
 152     acc.add_group(
 153         &GroupLabel("Generate getter/setter".to_owned()),
 154         AssistId(id, AssistKind::Generate),
 155         label,
 156         target,
 157         |builder| {
 158             let record_fields_count = info_of_record_fields.len();
 159
 160             let mut buf = String::with_capacity(512);
 161
 162             // Check if an impl exists
 163             if let Some(impl_def) = &getter_info.impl_def {
 164                 // Check if impl is empty
 165                 if let Some(assoc_item_list) = impl_def.assoc_item_list() {
 166                     if assoc_item_list.assoc_items().next().is_some() {
 167                         // If not empty then only insert a new line
 168                         buf.push('\n');
 169                     }
 170                 }
 171             }
 172
 173             for (i, record_field_info) in info_of_record_fields.iter().enumerate() {
 174                 // this buf inserts a newline at the end of a getter
 175                 // automatically, if one wants to add one more newline
 176                 // for separating it from other assoc items, that needs
 177                 // to be handled spearately
 178                 let mut getter_buf =
 179                     generate_getter_from_info(ctx, &getter_info, &record_field_info);
 180
 181                 // Insert `$0` only for last getter we generate
 182                 if i == record_fields_count - 1 {
 183                     getter_buf = getter_buf.replacen("fn ", "fn $0", 1);
 184                 }
 185
 186                 // For first element we do not merge with '\n', as
 187                 // that can be inserted by impl_def check defined
 188                 // above, for other cases which are:
 189                 //
 190                 // - impl exists but it empty, here we would ideally
 191                 // not want to keep newline between impl <struct> {
 192                 // and fn <fn-name>() { line
 193                 //
 194                 // - next if impl itself does not exist, in this
 195                 // case we ourselves generate a new impl and that
 196                 // again ends up with the same reasoning as above
 197                 // for not keeping newline
 198                 if i == 0 {
 199                     buf = buf + &getter_buf;
 200                 } else {
 201                     buf = buf + "\n" + &getter_buf;
 202                 }
 203
 204                 // We don't insert a new line at the end of
 205                 // last getter as it will end up in the end
 206                 // of an impl where we would not like to keep
 207                 // getter and end of impl ( i.e. `}` ) with an
 208                 // extra line for no reason
 209                 if i < record_fields_count - 1 {
 210                     buf = buf + "\n";
 211                 }
 212             }
 213
 214             let start_offset = getter_info
 215                 .impl_def
 216                 .as_ref()
 217                 .and_then(|impl_def| find_impl_block_end(impl_def.to_owned(), &mut buf))
 218                 .unwrap_or_else(|| {
 219                     buf = generate_impl_text(&ast::Adt::Struct(getter_info.strukt.clone()), &buf);
 220                     getter_info.strukt.syntax().text_range().end()
 221                 });
 222
 223             match ctx.config.snippet_cap {
 224                 Some(cap) => builder.insert_snippet(cap, start_offset, buf),
 225                 None => builder.insert(start_offset, buf),
 226             }
 227         },
 228     )
 229 }
 230
 231 fn generate_getter_from_info(
 232     ctx: &AssistContext<'_>,
 233     info: &GetterInfo,
 234     record_field_info: &RecordFieldInfo,
 235 ) -> String {
 236     let mut buf = String::with_capacity(512);
 237
 238     let vis = info.strukt.visibility().map_or(String::new(), |v| format!("{} ", v));
 239     let (ty, body) = if info.mutable {
 240         (
 241             format!("&mut {}", record_field_info.field_ty),
 242             format!("&mut self.{}", record_field_info.field_name),
 243         )
 244     } else {
 245         (|| {
 246             let krate = ctx.sema.scope(record_field_info.field_ty.syntax())?.krate();
 247             let famous_defs = &FamousDefs(&ctx.sema, krate);
 248             ctx.sema
 249                 .resolve_type(&record_field_info.field_ty)
 250                 .and_then(|ty| convert_reference_type(ty, ctx.db(), famous_defs))
 251                 .map(|conversion| {
 252                     cov_mark::hit!(convert_reference_type);
 253                     (
 254                         conversion.convert_type(ctx.db()),
 255                         conversion.getter(record_field_info.field_name.to_string()),
 256                     )
 257                 })
 258         })()
 259         .unwrap_or_else(|| {
 260             (
 261                 format!("&{}", record_field_info.field_ty),
 262                 format!("&self.{}", record_field_info.field_name),
 263             )
 264         })
 265     };
 266
 267     format_to!(
 268         buf,
 269         "    {}fn {}(&{}self) -> {} {{
 270         {}
 271     }}",
 272         vis,
 273         record_field_info.fn_name,
 274         info.mutable.then(|| "mut ").unwrap_or_default(),
 275         ty,
 276         body,
 277     );
 278
 279     buf
 280 }
 281
 282 fn extract_and_parse_record_fields(
 283     node: &ast::Struct,
 284     selection_range: TextRange,
 285     mutable: bool,
 286 ) -> Option<(Vec<RecordFieldInfo>, Vec<String>)> {
 287     let mut field_names: Vec<String> = vec![];
 288     let field_list = node.field_list()?;
 289
 290     match field_list {
 291         ast::FieldList::RecordFieldList(ele) => {
 292             let info_of_record_fields_in_selection = ele
 293                 .fields()
 294                 .filter_map(|record_field| {
 295                     if selection_range.contains_range(record_field.syntax().text_range()) {
 296                         let record_field_info = parse_record_field(record_field, mutable)?;
 297                         field_names.push(record_field_info.fn_name.clone());
 298                         return Some(record_field_info);
 299                     }
 300
 301                     None
 302                 })
 303                 .collect::<Vec<RecordFieldInfo>>();
 304
 305             if info_of_record_fields_in_selection.len() == 0 {
 306                 return None;
 307             }
 308
 309             Some((info_of_record_fields_in_selection, field_names))
 310         }
 311         ast::FieldList::TupleFieldList(_) => {
 312             return None;
 313         }
 314     }
 315 }
 316
 317 fn parse_record_field(record_field: ast::RecordField, mutable: bool) -> Option<RecordFieldInfo> {
 318     let field_name = record_field.name()?;
 319     let field_ty = record_field.ty()?;
 320
 321     let mut fn_name = to_lower_snake_case(&field_name.to_string());
 322     if mutable {
 323         format_to!(fn_name, "_mut");
 324     }
 325
 326     let target = record_field.syntax().text_range();
 327
 328     Some(RecordFieldInfo { field_name, field_ty, fn_name, target })
 329 }
 330
 331 #[cfg(test)]
 332 mod tests {
 333     use crate::tests::{check_assist, check_assist_not_applicable};
 334
 335     use super::*;
 336
 337     #[test]
 338     fn test_generate_getter_from_field() {
 339         check_assist(
 340             generate_getter,
 341             r#"
 342 struct Context {
 343     dat$0a: Data,
 344 }
 345 "#,
 346             r#"
 347 struct Context {
 348     data: Data,
 349 }
 350
 351 impl Context {
 352     fn $0data(&self) -> &Data {
 353         &self.data
 354     }
 355 }
 356 "#,
 357         );
 358
 359         check_assist(
 360             generate_getter_mut,
 361             r#"
 362 struct Context {
 363     dat$0a: Data,
 364 }
 365 "#,
 366             r#"
 367 struct Context {
 368     data: Data,
 369 }
 370
 371 impl Context {
 372     fn $0data_mut(&mut self) -> &mut Data {
 373         &mut self.data
 374     }
 375 }
 376 "#,
 377         );
 378     }
 379
 380     #[test]
 381     fn test_generate_getter_already_implemented() {
 382         check_assist_not_applicable(
 383             generate_getter,
 384             r#"
 385 struct Context {
 386     dat$0a: Data,
 387 }
 388
 389 impl Context {
 390     fn data(&self) -> &Data {
 391         &self.data
 392     }
 393 }
 394 "#,
 395         );
 396
 397         check_assist_not_applicable(
 398             generate_getter_mut,
 399             r#"
 400 struct Context {
 401     dat$0a: Data,
 402 }
 403
 404 impl Context {
 405     fn data_mut(&mut self) -> &mut Data {
 406         &mut self.data
 407     }
 408 }
 409 "#,
 410         );
 411     }
 412
 413     #[test]
 414     fn test_generate_getter_from_field_with_visibility_marker() {
 415         check_assist(
 416             generate_getter,
 417             r#"
 418 pub(crate) struct Context {
 419     dat$0a: Data,
 420 }
 421 "#,
 422             r#"
 423 pub(crate) struct Context {
 424     data: Data,
 425 }
 426
 427 impl Context {
 428     pub(crate) fn $0data(&self) -> &Data {
 429         &self.data
 430     }
 431 }
 432 "#,
 433         );
 434     }
 435
 436     #[test]
 437     fn test_multiple_generate_getter() {
 438         check_assist(
 439             generate_getter,
 440             r#"
 441 struct Context {
 442     data: Data,
 443     cou$0nt: usize,
 444 }
 445
 446 impl Context {
 447     fn data(&self) -> &Data {
 448         &self.data
 449     }
 450 }
 451 "#,
 452             r#"
 453 struct Context {
 454     data: Data,
 455     count: usize,
 456 }
 457
 458 impl Context {
 459     fn data(&self) -> &Data {
 460         &self.data
 461     }
 462
 463     fn $0count(&self) -> &usize {
 464         &self.count
 465     }
 466 }
 467 "#,
 468         );
 469     }
 470
 471     #[test]
 472     fn test_not_a_special_case() {
 473         cov_mark::check_count!(convert_reference_type, 0);
 474         // Fake string which doesn't implement AsRef<str>
 475         check_assist(
 476             generate_getter,
 477             r#"
 478 pub struct String;
 479
 480 struct S { foo: $0String }
 481 "#,
 482             r#"
 483 pub struct String;
 484
 485 struct S { foo: String }
 486
 487 impl S {
 488     fn $0foo(&self) -> &String {
 489         &self.foo
 490     }
 491 }
 492 "#,
 493         );
 494     }
 495
 496     #[test]
 497     fn test_convert_reference_type() {
 498         cov_mark::check_count!(convert_reference_type, 6);
 499
 500         // Copy
 501         check_assist(
 502             generate_getter,
 503             r#"
 504 //- minicore: copy
 505 struct S { foo: $0bool }
 506 "#,
 507             r#"
 508 struct S { foo: bool }
 509
 510 impl S {
 511     fn $0foo(&self) -> bool {
 512         self.foo
 513     }
 514 }
 515 "#,
 516         );
 517
 518         // AsRef<str>
 519         check_assist(
 520             generate_getter,
 521             r#"
 522 //- minicore: as_ref
 523 pub struct String;
 524 impl AsRef<str> for String {
 525     fn as_ref(&self) -> &str {
 526         ""
 527     }
 528 }
 529
 530 struct S { foo: $0String }
 531 "#,
 532             r#"
 533 pub struct String;
 534 impl AsRef<str> for String {
 535     fn as_ref(&self) -> &str {
 536         ""
 537     }
 538 }
 539
 540 struct S { foo: String }
 541
 542 impl S {
 543     fn $0foo(&self) -> &str {
 544         self.foo.as_ref()
 545     }
 546 }
 547 "#,
 548         );
 549
 550         // AsRef<T>
 551         check_assist(
 552             generate_getter,
 553             r#"
 554 //- minicore: as_ref
 555 struct Sweets;
 556
 557 pub struct Box<T>(T);
 558 impl<T> AsRef<T> for Box<T> {
 559     fn as_ref(&self) -> &T {
 560         &self.0
 561     }
 562 }
 563
 564 struct S { foo: $0Box<Sweets> }
 565 "#,
 566             r#"
 567 struct Sweets;
 568
 569 pub struct Box<T>(T);
 570 impl<T> AsRef<T> for Box<T> {
 571     fn as_ref(&self) -> &T {
 572         &self.0
 573     }
 574 }
 575
 576 struct S { foo: Box<Sweets> }
 577
 578 impl S {
 579     fn $0foo(&self) -> &Sweets {
 580         self.foo.as_ref()
 581     }
 582 }
 583 "#,
 584         );
 585
 586         // AsRef<[T]>
 587         check_assist(
 588             generate_getter,
 589             r#"
 590 //- minicore: as_ref
 591 pub struct Vec<T>;
 592 impl<T> AsRef<[T]> for Vec<T> {
 593     fn as_ref(&self) -> &[T] {
 594         &[]
 595     }
 596 }
 597
 598 struct S { foo: $0Vec<()> }
 599 "#,
 600             r#"
 601 pub struct Vec<T>;
 602 impl<T> AsRef<[T]> for Vec<T> {
 603     fn as_ref(&self) -> &[T] {
 604         &[]
 605     }
 606 }
 607
 608 struct S { foo: Vec<()> }
 609
 610 impl S {
 611     fn $0foo(&self) -> &[()] {
 612         self.foo.as_ref()
 613     }
 614 }
 615 "#,
 616         );
 617
 618         // Option
 619         check_assist(
 620             generate_getter,
 621             r#"
 622 //- minicore: option
 623 struct Failure;
 624
 625 struct S { foo: $0Option<Failure> }
 626 "#,
 627             r#"
 628 struct Failure;
 629
 630 struct S { foo: Option<Failure> }
 631
 632 impl S {
 633     fn $0foo(&self) -> Option<&Failure> {
 634         self.foo.as_ref()
 635     }
 636 }
 637 "#,
 638         );
 639
 640         // Result
 641         check_assist(
 642             generate_getter,
 643             r#"
 644 //- minicore: result
 645 struct Context {
 646     dat$0a: Result<bool, i32>,
 647 }
 648 "#,
 649             r#"
 650 struct Context {
 651     data: Result<bool, i32>,
 652 }
 653
 654 impl Context {
 655     fn $0data(&self) -> Result<&bool, &i32> {
 656         self.data.as_ref()
 657     }
 658 }
 659 "#,
 660         );
 661     }
 662
 663     #[test]
 664     fn test_generate_multiple_getters_from_selection() {
 665         check_assist(
 666             generate_getter,
 667             r#"
 668 struct Context {
 669     $0data: Data,
 670     count: usize,$0
 671 }
 672     "#,
 673             r#"
 674 struct Context {
 675     data: Data,
 676     count: usize,
 677 }
 678
 679 impl Context {
 680     fn data(&self) -> &Data {
 681         &self.data
 682     }
 683
 684     fn $0count(&self) -> &usize {
 685         &self.count
 686     }
 687 }
 688     "#,
 689         );
 690     }
 691
 692     #[test]
 693     fn test_generate_multiple_getters_from_selection_one_already_exists() {
 694         // As impl for one of the fields already exist, skip it
 695         check_assist_not_applicable(
 696             generate_getter,
 697             r#"
 698 struct Context {
 699     $0data: Data,
 700     count: usize,$0
 701 }
 702
 703 impl Context {
 704     fn data(&self) -> &Data {
 705         &self.data
 706     }
 707 }
 708     "#,
 709         );
 710     }
 711 }