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                     if ctx.config.snippet_cap.is_some() {
 184                         getter_buf = getter_buf.replacen("fn ", "fn $0", 1);
 185                     }
 186                 }
 187
 188                 // For first element we do not merge with '\n', as
 189                 // that can be inserted by impl_def check defined
 190                 // above, for other cases which are:
 191                 //
 192                 // - impl exists but it empty, here we would ideally
 193                 // not want to keep newline between impl <struct> {
 194                 // and fn <fn-name>() { line
 195                 //
 196                 // - next if impl itself does not exist, in this
 197                 // case we ourselves generate a new impl and that
 198                 // again ends up with the same reasoning as above
 199                 // for not keeping newline
 200                 if i == 0 {
 201                     buf = buf + &getter_buf;
 202                 } else {
 203                     buf = buf + "\n" + &getter_buf;
 204                 }
 205
 206                 // We don't insert a new line at the end of
 207                 // last getter as it will end up in the end
 208                 // of an impl where we would not like to keep
 209                 // getter and end of impl ( i.e. `}` ) with an
 210                 // extra line for no reason
 211                 if i < record_fields_count - 1 {
 212                     buf = buf + "\n";
 213                 }
 214             }
 215
 216             let start_offset = getter_info
 217                 .impl_def
 218                 .as_ref()
 219                 .and_then(|impl_def| find_impl_block_end(impl_def.to_owned(), &mut buf))
 220                 .unwrap_or_else(|| {
 221                     buf = generate_impl_text(&ast::Adt::Struct(getter_info.strukt.clone()), &buf);
 222                     getter_info.strukt.syntax().text_range().end()
 223                 });
 224
 225             match ctx.config.snippet_cap {
 226                 Some(cap) => builder.insert_snippet(cap, start_offset, buf),
 227                 None => builder.insert(start_offset, buf),
 228             }
 229         },
 230     )
 231 }
 232
 233 fn generate_getter_from_info(
 234     ctx: &AssistContext<'_>,
 235     info: &GetterInfo,
 236     record_field_info: &RecordFieldInfo,
 237 ) -> String {
 238     let mut buf = String::with_capacity(512);
 239
 240     let vis = info.strukt.visibility().map_or(String::new(), |v| format!("{v} "));
 241     let (ty, body) = if info.mutable {
 242         (
 243             format!("&mut {}", record_field_info.field_ty),
 244             format!("&mut self.{}", record_field_info.field_name),
 245         )
 246     } else {
 247         (|| {
 248             let krate = ctx.sema.scope(record_field_info.field_ty.syntax())?.krate();
 249             let famous_defs = &FamousDefs(&ctx.sema, krate);
 250             ctx.sema
 251                 .resolve_type(&record_field_info.field_ty)
 252                 .and_then(|ty| convert_reference_type(ty, ctx.db(), famous_defs))
 253                 .map(|conversion| {
 254                     cov_mark::hit!(convert_reference_type);
 255                     (
 256                         conversion.convert_type(ctx.db()),
 257                         conversion.getter(record_field_info.field_name.to_string()),
 258                     )
 259                 })
 260         })()
 261         .unwrap_or_else(|| {
 262             (
 263                 format!("&{}", record_field_info.field_ty),
 264                 format!("&self.{}", record_field_info.field_name),
 265             )
 266         })
 267     };
 268
 269     format_to!(
 270         buf,
 271         "    {}fn {}(&{}self) -> {} {{
 272         {}
 273     }}",
 274         vis,
 275         record_field_info.fn_name,
 276         info.mutable.then_some("mut ").unwrap_or_default(),
 277         ty,
 278         body,
 279     );
 280
 281     buf
 282 }
 283
 284 fn extract_and_parse_record_fields(
 285     node: &ast::Struct,
 286     selection_range: TextRange,
 287     mutable: bool,
 288 ) -> Option<(Vec<RecordFieldInfo>, Vec<String>)> {
 289     let mut field_names: Vec<String> = vec![];
 290     let field_list = node.field_list()?;
 291
 292     match field_list {
 293         ast::FieldList::RecordFieldList(ele) => {
 294             let info_of_record_fields_in_selection = ele
 295                 .fields()
 296                 .filter_map(|record_field| {
 297                     if selection_range.contains_range(record_field.syntax().text_range()) {
 298                         let record_field_info = parse_record_field(record_field, mutable)?;
 299                         field_names.push(record_field_info.fn_name.clone());
 300                         return Some(record_field_info);
 301                     }
 302
 303                     None
 304                 })
 305                 .collect::<Vec<RecordFieldInfo>>();
 306
 307             if info_of_record_fields_in_selection.len() == 0 {
 308                 return None;
 309             }
 310
 311             Some((info_of_record_fields_in_selection, field_names))
 312         }
 313         ast::FieldList::TupleFieldList(_) => {
 314             return None;
 315         }
 316     }
 317 }
 318
 319 fn parse_record_field(record_field: ast::RecordField, mutable: bool) -> Option<RecordFieldInfo> {
 320     let field_name = record_field.name()?;
 321     let field_ty = record_field.ty()?;
 322
 323     let mut fn_name = to_lower_snake_case(&field_name.to_string());
 324     if mutable {
 325         format_to!(fn_name, "_mut");
 326     }
 327
 328     let target = record_field.syntax().text_range();
 329
 330     Some(RecordFieldInfo { field_name, field_ty, fn_name, target })
 331 }
 332
 333 #[cfg(test)]
 334 mod tests {
 335     use crate::tests::{check_assist, check_assist_no_snippet_cap, check_assist_not_applicable};
 336
 337     use super::*;
 338
 339     #[test]
 340     fn test_generate_getter_from_field() {
 341         check_assist(
 342             generate_getter,
 343             r#"
 344 struct Context {
 345     dat$0a: Data,
 346 }
 347 "#,
 348             r#"
 349 struct Context {
 350     data: Data,
 351 }
 352
 353 impl Context {
 354     fn $0data(&self) -> &Data {
 355         &self.data
 356     }
 357 }
 358 "#,
 359         );
 360
 361         check_assist(
 362             generate_getter_mut,
 363             r#"
 364 struct Context {
 365     dat$0a: Data,
 366 }
 367 "#,
 368             r#"
 369 struct Context {
 370     data: Data,
 371 }
 372
 373 impl Context {
 374     fn $0data_mut(&mut self) -> &mut Data {
 375         &mut self.data
 376     }
 377 }
 378 "#,
 379         );
 380     }
 381
 382     #[test]
 383     fn test_generate_getter_from_field_no_snippet_cap() {
 384         check_assist_no_snippet_cap(
 385             generate_getter,
 386             r#"
 387 struct Context {
 388     dat$0a: Data,
 389 }
 390 "#,
 391             r#"
 392 struct Context {
 393     data: Data,
 394 }
 395
 396 impl Context {
 397     fn data(&self) -> &Data {
 398         &self.data
 399     }
 400 }
 401 "#,
 402         );
 403
 404         check_assist_no_snippet_cap(
 405             generate_getter_mut,
 406             r#"
 407 struct Context {
 408     dat$0a: Data,
 409 }
 410 "#,
 411             r#"
 412 struct Context {
 413     data: Data,
 414 }
 415
 416 impl Context {
 417     fn data_mut(&mut self) -> &mut Data {
 418         &mut self.data
 419     }
 420 }
 421 "#,
 422         );
 423     }
 424
 425     #[test]
 426     fn test_generate_getter_already_implemented() {
 427         check_assist_not_applicable(
 428             generate_getter,
 429             r#"
 430 struct Context {
 431     dat$0a: Data,
 432 }
 433
 434 impl Context {
 435     fn data(&self) -> &Data {
 436         &self.data
 437     }
 438 }
 439 "#,
 440         );
 441
 442         check_assist_not_applicable(
 443             generate_getter_mut,
 444             r#"
 445 struct Context {
 446     dat$0a: Data,
 447 }
 448
 449 impl Context {
 450     fn data_mut(&mut self) -> &mut Data {
 451         &mut self.data
 452     }
 453 }
 454 "#,
 455         );
 456     }
 457
 458     #[test]
 459     fn test_generate_getter_from_field_with_visibility_marker() {
 460         check_assist(
 461             generate_getter,
 462             r#"
 463 pub(crate) struct Context {
 464     dat$0a: Data,
 465 }
 466 "#,
 467             r#"
 468 pub(crate) struct Context {
 469     data: Data,
 470 }
 471
 472 impl Context {
 473     pub(crate) fn $0data(&self) -> &Data {
 474         &self.data
 475     }
 476 }
 477 "#,
 478         );
 479     }
 480
 481     #[test]
 482     fn test_generate_getter_from_field_with_visibility_marker_no_snippet_cap() {
 483         check_assist_no_snippet_cap(
 484             generate_getter,
 485             r#"
 486 pub(crate) struct Context {
 487     dat$0a: Data,
 488 }
 489 "#,
 490             r#"
 491 pub(crate) struct Context {
 492     data: Data,
 493 }
 494
 495 impl Context {
 496     pub(crate) fn data(&self) -> &Data {
 497         &self.data
 498     }
 499 }
 500 "#,
 501         );
 502     }
 503
 504     #[test]
 505     fn test_multiple_generate_getter() {
 506         check_assist(
 507             generate_getter,
 508             r#"
 509 struct Context {
 510     data: Data,
 511     cou$0nt: usize,
 512 }
 513
 514 impl Context {
 515     fn data(&self) -> &Data {
 516         &self.data
 517     }
 518 }
 519 "#,
 520             r#"
 521 struct Context {
 522     data: Data,
 523     count: usize,
 524 }
 525
 526 impl Context {
 527     fn data(&self) -> &Data {
 528         &self.data
 529     }
 530
 531     fn $0count(&self) -> &usize {
 532         &self.count
 533     }
 534 }
 535 "#,
 536         );
 537     }
 538
 539     #[test]
 540     fn test_multiple_generate_getter_no_snippet_cap() {
 541         check_assist_no_snippet_cap(
 542             generate_getter,
 543             r#"
 544 struct Context {
 545     data: Data,
 546     cou$0nt: usize,
 547 }
 548
 549 impl Context {
 550     fn data(&self) -> &Data {
 551         &self.data
 552     }
 553 }
 554 "#,
 555             r#"
 556 struct Context {
 557     data: Data,
 558     count: usize,
 559 }
 560
 561 impl Context {
 562     fn data(&self) -> &Data {
 563         &self.data
 564     }
 565
 566     fn count(&self) -> &usize {
 567         &self.count
 568     }
 569 }
 570 "#,
 571         );
 572     }
 573
 574     #[test]
 575     fn test_not_a_special_case() {
 576         cov_mark::check_count!(convert_reference_type, 0);
 577         // Fake string which doesn't implement AsRef<str>
 578         check_assist(
 579             generate_getter,
 580             r#"
 581 pub struct String;
 582
 583 struct S { foo: $0String }
 584 "#,
 585             r#"
 586 pub struct String;
 587
 588 struct S { foo: String }
 589
 590 impl S {
 591     fn $0foo(&self) -> &String {
 592         &self.foo
 593     }
 594 }
 595 "#,
 596         );
 597     }
 598
 599     #[test]
 600     fn test_convert_reference_type() {
 601         cov_mark::check_count!(convert_reference_type, 6);
 602
 603         // Copy
 604         check_assist(
 605             generate_getter,
 606             r#"
 607 //- minicore: copy
 608 struct S { foo: $0bool }
 609 "#,
 610             r#"
 611 struct S { foo: bool }
 612
 613 impl S {
 614     fn $0foo(&self) -> bool {
 615         self.foo
 616     }
 617 }
 618 "#,
 619         );
 620
 621         // AsRef<str>
 622         check_assist(
 623             generate_getter,
 624             r#"
 625 //- minicore: as_ref
 626 pub struct String;
 627 impl AsRef<str> for String {
 628     fn as_ref(&self) -> &str {
 629         ""
 630     }
 631 }
 632
 633 struct S { foo: $0String }
 634 "#,
 635             r#"
 636 pub struct String;
 637 impl AsRef<str> for String {
 638     fn as_ref(&self) -> &str {
 639         ""
 640     }
 641 }
 642
 643 struct S { foo: String }
 644
 645 impl S {
 646     fn $0foo(&self) -> &str {
 647         self.foo.as_ref()
 648     }
 649 }
 650 "#,
 651         );
 652
 653         // AsRef<T>
 654         check_assist(
 655             generate_getter,
 656             r#"
 657 //- minicore: as_ref
 658 struct Sweets;
 659
 660 pub struct Box<T>(T);
 661 impl<T> AsRef<T> for Box<T> {
 662     fn as_ref(&self) -> &T {
 663         &self.0
 664     }
 665 }
 666
 667 struct S { foo: $0Box<Sweets> }
 668 "#,
 669             r#"
 670 struct Sweets;
 671
 672 pub struct Box<T>(T);
 673 impl<T> AsRef<T> for Box<T> {
 674     fn as_ref(&self) -> &T {
 675         &self.0
 676     }
 677 }
 678
 679 struct S { foo: Box<Sweets> }
 680
 681 impl S {
 682     fn $0foo(&self) -> &Sweets {
 683         self.foo.as_ref()
 684     }
 685 }
 686 "#,
 687         );
 688
 689         // AsRef<[T]>
 690         check_assist(
 691             generate_getter,
 692             r#"
 693 //- minicore: as_ref
 694 pub struct Vec<T>;
 695 impl<T> AsRef<[T]> for Vec<T> {
 696     fn as_ref(&self) -> &[T] {
 697         &[]
 698     }
 699 }
 700
 701 struct S { foo: $0Vec<()> }
 702 "#,
 703             r#"
 704 pub struct Vec<T>;
 705 impl<T> AsRef<[T]> for Vec<T> {
 706     fn as_ref(&self) -> &[T] {
 707         &[]
 708     }
 709 }
 710
 711 struct S { foo: Vec<()> }
 712
 713 impl S {
 714     fn $0foo(&self) -> &[()] {
 715         self.foo.as_ref()
 716     }
 717 }
 718 "#,
 719         );
 720
 721         // Option
 722         check_assist(
 723             generate_getter,
 724             r#"
 725 //- minicore: option
 726 struct Failure;
 727
 728 struct S { foo: $0Option<Failure> }
 729 "#,
 730             r#"
 731 struct Failure;
 732
 733 struct S { foo: Option<Failure> }
 734
 735 impl S {
 736     fn $0foo(&self) -> Option<&Failure> {
 737         self.foo.as_ref()
 738     }
 739 }
 740 "#,
 741         );
 742
 743         // Result
 744         check_assist(
 745             generate_getter,
 746             r#"
 747 //- minicore: result
 748 struct Context {
 749     dat$0a: Result<bool, i32>,
 750 }
 751 "#,
 752             r#"
 753 struct Context {
 754     data: Result<bool, i32>,
 755 }
 756
 757 impl Context {
 758     fn $0data(&self) -> Result<&bool, &i32> {
 759         self.data.as_ref()
 760     }
 761 }
 762 "#,
 763         );
 764     }
 765
 766     #[test]
 767     fn test_generate_multiple_getters_from_selection() {
 768         check_assist(
 769             generate_getter,
 770             r#"
 771 struct Context {
 772     $0data: Data,
 773     count: usize,$0
 774 }
 775     "#,
 776             r#"
 777 struct Context {
 778     data: Data,
 779     count: usize,
 780 }
 781
 782 impl Context {
 783     fn data(&self) -> &Data {
 784         &self.data
 785     }
 786
 787     fn $0count(&self) -> &usize {
 788         &self.count
 789     }
 790 }
 791     "#,
 792         );
 793     }
 794
 795     #[test]
 796     fn test_generate_multiple_getters_from_selection_one_already_exists() {
 797         // As impl for one of the fields already exist, skip it
 798         check_assist_not_applicable(
 799             generate_getter,
 800             r#"
 801 struct Context {
 802     $0data: Data,
 803     count: usize,$0
 804 }
 805
 806 impl Context {
 807     fn data(&self) -> &Data {
 808         &self.data
 809     }
 810 }
 811     "#,
 812         );
 813     }
 814 }