crates/ide_assists/src/handlers/merge_match_arms.rs

   1 use hir::TypeInfo;
   2 use std::{collections::HashMap, iter::successors};
   3 use syntax::{
   4     algo::neighbor,
   5     ast::{self, AstNode, HasName},
   6     Direction,
   7 };
   8
   9 use crate::{AssistContext, AssistId, AssistKind, Assists, TextRange};
  10
  11 // Assist: merge_match_arms
  12 //
  13 // Merges the current match arm with the following if their bodies are identical.
  14 //
  15 // ```
  16 // enum Action { Move { distance: u32 }, Stop }
  17 //
  18 // fn handle(action: Action) {
  19 //     match action {
  20 //         $0Action::Move(..) => foo(),
  21 //         Action::Stop => foo(),
  22 //     }
  23 // }
  24 // ```
  25 // ->
  26 // ```
  27 // enum Action { Move { distance: u32 }, Stop }
  28 //
  29 // fn handle(action: Action) {
  30 //     match action {
  31 //         Action::Move(..) | Action::Stop => foo(),
  32 //     }
  33 // }
  34 // ```
  35 pub(crate) fn merge_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  36     let current_arm = ctx.find_node_at_offset::<ast::MatchArm>()?;
  37     // Don't try to handle arms with guards for now - can add support for this later
  38     if current_arm.guard().is_some() {
  39         return None;
  40     }
  41     let current_expr = current_arm.expr()?;
  42     let current_text_range = current_arm.syntax().text_range();
  43     let current_arm_types = get_arm_types(&ctx, &current_arm);
  44
  45     // We check if the following match arms match this one. We could, but don't,
  46     // compare to the previous match arm as well.
  47     let arms_to_merge = successors(Some(current_arm), |it| neighbor(it, Direction::Next))
  48         .take_while(|arm| match arm.expr() {
  49             Some(expr) if arm.guard().is_none() => {
  50                 let same_text = expr.syntax().text() == current_expr.syntax().text();
  51                 if !same_text {
  52                     return false;
  53                 }
  54
  55                 are_same_types(&current_arm_types, arm, ctx)
  56             }
  57             _ => false,
  58         })
  59         .collect::<Vec<_>>();
  60
  61     if arms_to_merge.len() <= 1 {
  62         return None;
  63     }
  64
  65     acc.add(
  66         AssistId("merge_match_arms", AssistKind::RefactorRewrite),
  67         "Merge match arms",
  68         current_text_range,
  69         |edit| {
  70             let pats = if arms_to_merge.iter().any(contains_placeholder) {
  71                 "_".into()
  72             } else {
  73                 arms_to_merge
  74                     .iter()
  75                     .filter_map(ast::MatchArm::pat)
  76                     .map(|x| x.syntax().to_string())
  77                     .collect::<Vec<String>>()
  78                     .join(" | ")
  79             };
  80
  81             let arm = format!("{} => {},", pats, current_expr.syntax().text());
  82
  83             if let [first, .., last] = &*arms_to_merge {
  84                 let start = first.syntax().text_range().start();
  85                 let end = last.syntax().text_range().end();
  86
  87                 edit.replace(TextRange::new(start, end), arm);
  88             }
  89         },
  90     )
  91 }
  92
  93 fn contains_placeholder(a: &ast::MatchArm) -> bool {
  94     matches!(a.pat(), Some(ast::Pat::WildcardPat(..)))
  95 }
  96
  97 fn are_same_types(
  98     current_arm_types: &HashMap<String, Option<TypeInfo>>,
  99     arm: &ast::MatchArm,
 100     ctx: &AssistContext,
 101 ) -> bool {
 102     let arm_types = get_arm_types(&ctx, &arm);
 103     for (other_arm_type_name, other_arm_type) in arm_types {
 104         match (current_arm_types.get(&other_arm_type_name), other_arm_type) {
 105             (Some(Some(current_arm_type)), Some(other_arm_type))
 106                 if other_arm_type.original == current_arm_type.original =>
 107             {
 108                 ()
 109             }
 110             _ => return false,
 111         }
 112     }
 113
 114     true
 115 }
 116
 117 fn get_arm_types(
 118     context: &AssistContext,
 119     arm: &ast::MatchArm,
 120 ) -> HashMap<String, Option<TypeInfo>> {
 121     let mut mapping: HashMap<String, Option<TypeInfo>> = HashMap::new();
 122
 123     fn recurse(
 124         map: &mut HashMap<String, Option<TypeInfo>>,
 125         ctx: &AssistContext,
 126         pat: &Option<ast::Pat>,
 127     ) {
 128         if let Some(local_pat) = pat {
 129             match pat {
 130                 Some(ast::Pat::TupleStructPat(tuple)) => {
 131                     for field in tuple.fields() {
 132                         recurse(map, ctx, &Some(field));
 133                     }
 134                 }
 135                 Some(ast::Pat::TuplePat(tuple)) => {
 136                     for field in tuple.fields() {
 137                         recurse(map, ctx, &Some(field));
 138                     }
 139                 }
 140                 Some(ast::Pat::RecordPat(record)) => {
 141                     if let Some(field_list) = record.record_pat_field_list() {
 142                         for field in field_list.fields() {
 143                             recurse(map, ctx, &field.pat());
 144                         }
 145                     }
 146                 }
 147                 Some(ast::Pat::ParenPat(parentheses)) => {
 148                     recurse(map, ctx, &parentheses.pat());
 149                 }
 150                 Some(ast::Pat::SlicePat(slice)) => {
 151                     for slice_pat in slice.pats() {
 152                         recurse(map, ctx, &Some(slice_pat));
 153                     }
 154                 }
 155                 Some(ast::Pat::IdentPat(ident_pat)) => {
 156                     if let Some(name) = ident_pat.name() {
 157                         let pat_type = ctx.sema.type_of_pat(local_pat);
 158                         map.insert(name.text().to_string(), pat_type);
 159                     }
 160                 }
 161                 _ => (),
 162             }
 163         }
 164     }
 165
 166     recurse(&mut mapping, &context, &arm.pat());
 167     mapping
 168 }
 169
 170 #[cfg(test)]
 171 mod tests {
 172     use crate::tests::{check_assist, check_assist_not_applicable};
 173
 174     use super::*;
 175
 176     #[test]
 177     fn merge_match_arms_single_patterns() {
 178         check_assist(
 179             merge_match_arms,
 180             r#"
 181 #[derive(Debug)]
 182 enum X { A, B, C }
 183
 184 fn main() {
 185     let x = X::A;
 186     let y = match x {
 187         X::A => { 1i32$0 }
 188         X::B => { 1i32 }
 189         X::C => { 2i32 }
 190     }
 191 }
 192 "#,
 193             r#"
 194 #[derive(Debug)]
 195 enum X { A, B, C }
 196
 197 fn main() {
 198     let x = X::A;
 199     let y = match x {
 200         X::A | X::B => { 1i32 },
 201         X::C => { 2i32 }
 202     }
 203 }
 204 "#,
 205         );
 206     }
 207
 208     #[test]
 209     fn merge_match_arms_multiple_patterns() {
 210         check_assist(
 211             merge_match_arms,
 212             r#"
 213 #[derive(Debug)]
 214 enum X { A, B, C, D, E }
 215
 216 fn main() {
 217     let x = X::A;
 218     let y = match x {
 219         X::A | X::B => {$0 1i32 },
 220         X::C | X::D => { 1i32 },
 221         X::E => { 2i32 },
 222     }
 223 }
 224 "#,
 225             r#"
 226 #[derive(Debug)]
 227 enum X { A, B, C, D, E }
 228
 229 fn main() {
 230     let x = X::A;
 231     let y = match x {
 232         X::A | X::B | X::C | X::D => { 1i32 },
 233         X::E => { 2i32 },
 234     }
 235 }
 236 "#,
 237         );
 238     }
 239
 240     #[test]
 241     fn merge_match_arms_placeholder_pattern() {
 242         check_assist(
 243             merge_match_arms,
 244             r#"
 245 #[derive(Debug)]
 246 enum X { A, B, C, D, E }
 247
 248 fn main() {
 249     let x = X::A;
 250     let y = match x {
 251         X::A => { 1i32 },
 252         X::B => { 2i$032 },
 253         _ => { 2i32 }
 254     }
 255 }
 256 "#,
 257             r#"
 258 #[derive(Debug)]
 259 enum X { A, B, C, D, E }
 260
 261 fn main() {
 262     let x = X::A;
 263     let y = match x {
 264         X::A => { 1i32 },
 265         _ => { 2i32 },
 266     }
 267 }
 268 "#,
 269         );
 270     }
 271
 272     #[test]
 273     fn merges_all_subsequent_arms() {
 274         check_assist(
 275             merge_match_arms,
 276             r#"
 277 enum X { A, B, C, D, E }
 278
 279 fn main() {
 280     match X::A {
 281         X::A$0 => 92,
 282         X::B => 92,
 283         X::C => 92,
 284         X::D => 62,
 285         _ => panic!(),
 286     }
 287 }
 288 "#,
 289             r#"
 290 enum X { A, B, C, D, E }
 291
 292 fn main() {
 293     match X::A {
 294         X::A | X::B | X::C => 92,
 295         X::D => 62,
 296         _ => panic!(),
 297     }
 298 }
 299 "#,
 300         )
 301     }
 302
 303     #[test]
 304     fn merge_match_arms_rejects_guards() {
 305         check_assist_not_applicable(
 306             merge_match_arms,
 307             r#"
 308 #[derive(Debug)]
 309 enum X {
 310     A(i32),
 311     B,
 312     C
 313 }
 314
 315 fn main() {
 316     let x = X::A;
 317     let y = match x {
 318         X::A(a) if a > 5 => { $01i32 },
 319         X::B => { 1i32 },
 320         X::C => { 2i32 }
 321     }
 322 }
 323 "#,
 324         );
 325     }
 326
 327     #[test]
 328     fn merge_match_arms_different_type() {
 329         check_assist_not_applicable(
 330             merge_match_arms,
 331             r#"
 332 //- minicore: result
 333 fn func() {
 334     match Result::<f64, f32>::Ok(0f64) {
 335         Ok(x) => $0x.classify(),
 336         Err(x) => x.classify()
 337     };
 338 }
 339 "#,
 340         );
 341     }
 342
 343     #[test]
 344     fn merge_match_arms_different_type_multiple_fields() {
 345         check_assist_not_applicable(
 346             merge_match_arms,
 347             r#"
 348 //- minicore: result
 349 fn func() {
 350     match Result::<(f64, f64), (f32, f32)>::Ok((0f64, 0f64)) {
 351         Ok(x) => $0x.1.classify(),
 352         Err(x) => x.1.classify()
 353     };
 354 }
 355 "#,
 356         );
 357     }
 358
 359     #[test]
 360     fn merge_match_arms_same_type_multiple_fields() {
 361         check_assist(
 362             merge_match_arms,
 363             r#"
 364 //- minicore: result
 365 fn func() {
 366     match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
 367         Ok(x) => $0x.1.classify(),
 368         Err(x) => x.1.classify()
 369     };
 370 }
 371 "#,
 372             r#"
 373 fn func() {
 374     match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
 375         Ok(x) | Err(x) => x.1.classify(),
 376     };
 377 }
 378 "#,
 379         );
 380     }
 381
 382     #[test]
 383     fn merge_match_arms_same_type_subsequent_arm_with_different_type_in_other() {
 384         check_assist(
 385             merge_match_arms,
 386             r#"
 387 enum MyEnum {
 388     OptionA(f32),
 389     OptionB(f32),
 390     OptionC(f64)
 391 }
 392
 393 fn func(e: MyEnum) {
 394     match e {
 395         MyEnum::OptionA(x) => $0x.classify(),
 396         MyEnum::OptionB(x) => x.classify(),
 397         MyEnum::OptionC(x) => x.classify(),
 398     };
 399 }
 400 "#,
 401             r#"
 402 enum MyEnum {
 403     OptionA(f32),
 404     OptionB(f32),
 405     OptionC(f64)
 406 }
 407
 408 fn func(e: MyEnum) {
 409     match e {
 410         MyEnum::OptionA(x) | MyEnum::OptionB(x) => x.classify(),
 411         MyEnum::OptionC(x) => x.classify(),
 412     };
 413 }
 414 "#,
 415         );
 416     }
 417
 418     #[test]
 419     fn merge_match_arms_same_type_skip_arm_with_different_type_in_between() {
 420         check_assist_not_applicable(
 421             merge_match_arms,
 422             r#"
 423 enum MyEnum {
 424     OptionA(f32),
 425     OptionB(f64),
 426     OptionC(f32)
 427 }
 428
 429 fn func(e: MyEnum) {
 430     match e {
 431         MyEnum::OptionA(x) => $0x.classify(),
 432         MyEnum::OptionB(x) => x.classify(),
 433         MyEnum::OptionC(x) => x.classify(),
 434     };
 435 }
 436 "#,
 437         );
 438     }
 439
 440     #[test]
 441     fn merge_match_arms_same_type_different_number_of_fields() {
 442         check_assist_not_applicable(
 443             merge_match_arms,
 444             r#"
 445 //- minicore: result
 446 fn func() {
 447     match Result::<(f64, f64, f64), (f64, f64)>::Ok((0f64, 0f64, 0f64)) {
 448         Ok(x) => $0x.1.classify(),
 449         Err(x) => x.1.classify()
 450     };
 451 }
 452 "#,
 453         );
 454     }
 455
 456     #[test]
 457     fn merge_match_same_destructuring_different_types() {
 458         check_assist_not_applicable(
 459             merge_match_arms,
 460             r#"
 461 struct Point {
 462     x: i32,
 463     y: i32,
 464 }
 465
 466 fn func() {
 467     let p = Point { x: 0, y: 7 };
 468
 469     match p {
 470         Point { x, y: 0 } => $0"",
 471         Point { x: 0, y } => "",
 472         Point { x, y } => "",
 473     };
 474 }
 475 "#,
 476         );
 477     }
 478
 479     #[test]
 480     fn merge_match_arms_range() {
 481         check_assist(
 482             merge_match_arms,
 483             r#"
 484 fn func() {
 485     let x = 'c';
 486
 487     match x {
 488         'a'..='j' => $0"",
 489         'c'..='z' => "",
 490         _ => "other",
 491     };
 492 }
 493 "#,
 494             r#"
 495 fn func() {
 496     let x = 'c';
 497
 498     match x {
 499         'a'..='j' | 'c'..='z' => "",
 500         _ => "other",
 501     };
 502 }
 503 "#,
 504         );
 505     }
 506
 507     #[test]
 508     fn merge_match_arms_enum_without_field() {
 509         check_assist_not_applicable(
 510             merge_match_arms,
 511             r#"
 512 enum MyEnum {
 513     NoField,
 514     AField(u8)
 515 }
 516
 517 fn func(x: MyEnum) {
 518     match x {
 519         MyEnum::NoField => $0"",
 520         MyEnum::AField(x) => ""
 521     };
 522 }
 523         "#,
 524         )
 525     }
 526
 527     #[test]
 528     fn merge_match_arms_enum_destructuring_different_types() {
 529         check_assist_not_applicable(
 530             merge_match_arms,
 531             r#"
 532 enum MyEnum {
 533     Move { x: i32, y: i32 },
 534     Write(String),
 535 }
 536
 537 fn func(x: MyEnum) {
 538     match x {
 539         MyEnum::Move { x, y } => $0"",
 540         MyEnum::Write(text) => "",
 541     };
 542 }
 543         "#,
 544         )
 545     }
 546
 547     #[test]
 548     fn merge_match_arms_enum_destructuring_same_types() {
 549         check_assist(
 550             merge_match_arms,
 551             r#"
 552 enum MyEnum {
 553     Move { x: i32, y: i32 },
 554     Crawl { x: i32, y: i32 }
 555 }
 556
 557 fn func(x: MyEnum) {
 558     match x {
 559         MyEnum::Move { x, y } => $0"",
 560         MyEnum::Crawl { x, y } => "",
 561     };
 562 }
 563         "#,
 564             r#"
 565 enum MyEnum {
 566     Move { x: i32, y: i32 },
 567     Crawl { x: i32, y: i32 }
 568 }
 569
 570 fn func(x: MyEnum) {
 571     match x {
 572         MyEnum::Move { x, y } | MyEnum::Crawl { x, y } => "",
 573     };
 574 }
 575         "#,
 576         )
 577     }
 578
 579     #[test]
 580     fn merge_match_arms_enum_destructuring_same_types_different_name() {
 581         check_assist_not_applicable(
 582             merge_match_arms,
 583             r#"
 584 enum MyEnum {
 585     Move { x: i32, y: i32 },
 586     Crawl { a: i32, b: i32 }
 587 }
 588
 589 fn func(x: MyEnum) {
 590     match x {
 591         MyEnum::Move { x, y } => $0"",
 592         MyEnum::Crawl { a, b } => "",
 593     };
 594 }
 595         "#,
 596         )
 597     }
 598
 599     #[test]
 600     fn merge_match_arms_enum_nested_pattern_different_names() {
 601         check_assist_not_applicable(
 602             merge_match_arms,
 603             r#"
 604 enum Color {
 605     Rgb(i32, i32, i32),
 606     Hsv(i32, i32, i32),
 607 }
 608
 609 enum Message {
 610     Quit,
 611     Move { x: i32, y: i32 },
 612     Write(String),
 613     ChangeColor(Color),
 614 }
 615
 616 fn main(msg: Message) {
 617     match msg {
 618         Message::ChangeColor(Color::Rgb(r, g, b)) => $0"",
 619         Message::ChangeColor(Color::Hsv(h, s, v)) => "",
 620         _ => "other"
 621     };
 622 }
 623         "#,
 624         )
 625     }
 626
 627     #[test]
 628     fn merge_match_arms_enum_nested_pattern_same_names() {
 629         check_assist(
 630             merge_match_arms,
 631             r#"
 632 enum Color {
 633     Rgb(i32, i32, i32),
 634     Hsv(i32, i32, i32),
 635 }
 636
 637 enum Message {
 638     Quit,
 639     Move { x: i32, y: i32 },
 640     Write(String),
 641     ChangeColor(Color),
 642 }
 643
 644 fn main(msg: Message) {
 645     match msg {
 646         Message::ChangeColor(Color::Rgb(a, b, c)) => $0"",
 647         Message::ChangeColor(Color::Hsv(a, b, c)) => "",
 648         _ => "other"
 649     };
 650 }
 651         "#,
 652             r#"
 653 enum Color {
 654     Rgb(i32, i32, i32),
 655     Hsv(i32, i32, i32),
 656 }
 657
 658 enum Message {
 659     Quit,
 660     Move { x: i32, y: i32 },
 661     Write(String),
 662     ChangeColor(Color),
 663 }
 664
 665 fn main(msg: Message) {
 666     match msg {
 667         Message::ChangeColor(Color::Rgb(a, b, c)) | Message::ChangeColor(Color::Hsv(a, b, c)) => "",
 668         _ => "other"
 669     };
 670 }
 671         "#,
 672         )
 673     }
 674
 675     #[test]
 676     fn merge_match_arms_enum_destructuring_with_ignore() {
 677         check_assist(
 678             merge_match_arms,
 679             r#"
 680 enum MyEnum {
 681     Move { x: i32, a: i32 },
 682     Crawl { x: i32, b: i32 }
 683 }
 684
 685 fn func(x: MyEnum) {
 686     match x {
 687         MyEnum::Move { x, .. } => $0"",
 688         MyEnum::Crawl { x, .. } => "",
 689     };
 690 }
 691         "#,
 692             r#"
 693 enum MyEnum {
 694     Move { x: i32, a: i32 },
 695     Crawl { x: i32, b: i32 }
 696 }
 697
 698 fn func(x: MyEnum) {
 699     match x {
 700         MyEnum::Move { x, .. } | MyEnum::Crawl { x, .. } => "",
 701     };
 702 }
 703         "#,
 704         )
 705     }
 706
 707     #[test]
 708     fn merge_match_arms_nested_with_conflicting_identifier() {
 709         check_assist_not_applicable(
 710             merge_match_arms,
 711             r#"
 712 enum Color {
 713     Rgb(i32, i32, i32),
 714     Hsv(i32, i32, i32),
 715 }
 716
 717 enum Message {
 718     Move { x: i32, y: i32 },
 719     ChangeColor(u8, Color),
 720 }
 721
 722 fn main(msg: Message) {
 723     match msg {
 724         Message::ChangeColor(x, Color::Rgb(y, b, c)) => $0"",
 725         Message::ChangeColor(y, Color::Hsv(x, b, c)) => "",
 726         _ => "other"
 727     };
 728 }
 729         "#,
 730         )
 731     }
 732
 733     #[test]
 734     fn merge_match_arms_tuple() {
 735         check_assist_not_applicable(
 736             merge_match_arms,
 737             r#"
 738 fn func() {
 739     match (0, "boo") {
 740         (x, y) => $0"",
 741         (y, x) => "",
 742     };
 743 }
 744         "#,
 745         )
 746     }
 747
 748     #[test]
 749     fn merge_match_arms_parentheses() {
 750         check_assist_not_applicable(
 751             merge_match_arms,
 752             r#"
 753 fn func(x: i32) {
 754     let variable = 2;
 755     match x {
 756         1 => $0"",
 757         ((((variable)))) => "",
 758         _ => "other"
 759     };
 760 }
 761         "#,
 762         )
 763     }
 764
 765     #[test]
 766     fn merge_match_arms_refpat() {
 767         check_assist_not_applicable(
 768             merge_match_arms,
 769             r#"
 770 fn func() {
 771     let name = Some(String::from(""));
 772     let n = String::from("");
 773     match name {
 774         Some(ref n) => $0"",
 775         Some(n) => "",
 776         _ => "other",
 777     };
 778 }
 779         "#,
 780         )
 781     }
 782
 783     #[test]
 784     fn merge_match_arms_slice() {
 785         check_assist_not_applicable(
 786             merge_match_arms,
 787             r#"
 788 fn func(binary: &[u8]) {
 789     let space = b' ';
 790     match binary {
 791         [0x7f, b'E', b'L', b'F', ..] => $0"",
 792         [space] => "",
 793         _ => "other",
 794     };
 795 }
 796         "#,
 797         )
 798     }
 799
 800     #[test]
 801     fn merge_match_arms_slice_identical() {
 802         check_assist(
 803             merge_match_arms,
 804             r#"
 805 fn func(binary: &[u8]) {
 806     let space = b' ';
 807     match binary {
 808         [space, 5u8] => $0"",
 809         [space] => "",
 810         _ => "other",
 811     };
 812 }
 813         "#,
 814             r#"
 815 fn func(binary: &[u8]) {
 816     let space = b' ';
 817     match binary {
 818         [space, 5u8] | [space] => "",
 819         _ => "other",
 820     };
 821 }
 822         "#,
 823         )
 824     }
 825 }