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