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