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                 return 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::IdentPat(ident_pat)) => {
 147                     if let Some(name) = ident_pat.name() {
 148                         let pat_type = ctx.sema.type_of_pat(local_pat);
 149                         map.insert(name.text().to_string(), pat_type);
 150                     }
 151                 }
 152                 _ => (),
 153             }
 154         }
 155     }
 156
 157     recurse(&arm.pat(), &mut mapping, &context);
 158     return mapping;
 159 }
 160
 161 #[cfg(test)]
 162 mod tests {
 163     use crate::tests::{check_assist, check_assist_not_applicable};
 164
 165     use super::*;
 166
 167     #[test]
 168     fn merge_match_arms_single_patterns() {
 169         check_assist(
 170             merge_match_arms,
 171             r#"
 172 #[derive(Debug)]
 173 enum X { A, B, C }
 174
 175 fn main() {
 176     let x = X::A;
 177     let y = match x {
 178         X::A => { 1i32$0 }
 179         X::B => { 1i32 }
 180         X::C => { 2i32 }
 181     }
 182 }
 183 "#,
 184             r#"
 185 #[derive(Debug)]
 186 enum X { A, B, C }
 187
 188 fn main() {
 189     let x = X::A;
 190     let y = match x {
 191         X::A | X::B => { 1i32 },
 192         X::C => { 2i32 }
 193     }
 194 }
 195 "#,
 196         );
 197     }
 198
 199     #[test]
 200     fn merge_match_arms_multiple_patterns() {
 201         check_assist(
 202             merge_match_arms,
 203             r#"
 204 #[derive(Debug)]
 205 enum X { A, B, C, D, E }
 206
 207 fn main() {
 208     let x = X::A;
 209     let y = match x {
 210         X::A | X::B => {$0 1i32 },
 211         X::C | X::D => { 1i32 },
 212         X::E => { 2i32 },
 213     }
 214 }
 215 "#,
 216             r#"
 217 #[derive(Debug)]
 218 enum X { A, B, C, D, E }
 219
 220 fn main() {
 221     let x = X::A;
 222     let y = match x {
 223         X::A | X::B | X::C | X::D => { 1i32 },
 224         X::E => { 2i32 },
 225     }
 226 }
 227 "#,
 228         );
 229     }
 230
 231     #[test]
 232     fn merge_match_arms_placeholder_pattern() {
 233         check_assist(
 234             merge_match_arms,
 235             r#"
 236 #[derive(Debug)]
 237 enum X { A, B, C, D, E }
 238
 239 fn main() {
 240     let x = X::A;
 241     let y = match x {
 242         X::A => { 1i32 },
 243         X::B => { 2i$032 },
 244         _ => { 2i32 }
 245     }
 246 }
 247 "#,
 248             r#"
 249 #[derive(Debug)]
 250 enum X { A, B, C, D, E }
 251
 252 fn main() {
 253     let x = X::A;
 254     let y = match x {
 255         X::A => { 1i32 },
 256         _ => { 2i32 },
 257     }
 258 }
 259 "#,
 260         );
 261     }
 262
 263     #[test]
 264     fn merges_all_subsequent_arms() {
 265         check_assist(
 266             merge_match_arms,
 267             r#"
 268 enum X { A, B, C, D, E }
 269
 270 fn main() {
 271     match X::A {
 272         X::A$0 => 92,
 273         X::B => 92,
 274         X::C => 92,
 275         X::D => 62,
 276         _ => panic!(),
 277     }
 278 }
 279 "#,
 280             r#"
 281 enum X { A, B, C, D, E }
 282
 283 fn main() {
 284     match X::A {
 285         X::A | X::B | X::C => 92,
 286         X::D => 62,
 287         _ => panic!(),
 288     }
 289 }
 290 "#,
 291         )
 292     }
 293
 294     #[test]
 295     fn merge_match_arms_rejects_guards() {
 296         check_assist_not_applicable(
 297             merge_match_arms,
 298             r#"
 299 #[derive(Debug)]
 300 enum X {
 301     A(i32),
 302     B,
 303     C
 304 }
 305
 306 fn main() {
 307     let x = X::A;
 308     let y = match x {
 309         X::A(a) if a > 5 => { $01i32 },
 310         X::B => { 1i32 },
 311         X::C => { 2i32 }
 312     }
 313 }
 314 "#,
 315         );
 316     }
 317
 318     #[test]
 319     fn merge_match_arms_different_type() {
 320         check_assist_not_applicable(
 321             merge_match_arms,
 322             r#"//- minicore: result
 323 fn func() {
 324     match Result::<f64, f32>::Ok(0f64) {
 325         Ok(x) => $0x.classify(),
 326         Err(x) => x.classify()
 327     };
 328 }
 329 "#,
 330         );
 331     }
 332
 333     #[test]
 334     fn merge_match_arms_different_type_multiple_fields() {
 335         check_assist_not_applicable(
 336             merge_match_arms,
 337             r#"//- minicore: result
 338 fn func() {
 339     match Result::<(f64, f64), (f32, f32)>::Ok((0f64, 0f64)) {
 340         Ok(x) => $0x.1.classify(),
 341         Err(x) => x.1.classify()
 342     };
 343 }
 344 "#,
 345         );
 346     }
 347
 348     #[test]
 349     fn merge_match_arms_same_type_multiple_fields() {
 350         check_assist(
 351             merge_match_arms,
 352             r#"//- minicore: result
 353 fn func() {
 354     match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
 355         Ok(x) => $0x.1.classify(),
 356         Err(x) => x.1.classify()
 357     };
 358 }
 359 "#,
 360             r#"
 361 fn func() {
 362     match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
 363         Ok(x) | Err(x) => x.1.classify(),
 364     };
 365 }
 366 "#,
 367         );
 368     }
 369
 370     #[test]
 371     fn merge_match_arms_same_type_subsequent_arm_with_different_type_in_other() {
 372         check_assist(
 373             merge_match_arms,
 374             r#"
 375 enum MyEnum {
 376     OptionA(f32),
 377     OptionB(f32),
 378     OptionC(f64)
 379 }
 380
 381 fn func(e: MyEnum) {
 382     match e {
 383         MyEnum::OptionA(x) => $0x.classify(),
 384         MyEnum::OptionB(x) => x.classify(),
 385         MyEnum::OptionC(x) => x.classify(),
 386     };
 387 }
 388 "#,
 389             r#"
 390 enum MyEnum {
 391     OptionA(f32),
 392     OptionB(f32),
 393     OptionC(f64)
 394 }
 395
 396 fn func(e: MyEnum) {
 397     match e {
 398         MyEnum::OptionA(x) | MyEnum::OptionB(x) => x.classify(),
 399         MyEnum::OptionC(x) => x.classify(),
 400     };
 401 }
 402 "#,
 403         );
 404     }
 405
 406     #[test]
 407     fn merge_match_arms_same_type_skip_arm_with_different_type_in_between() {
 408         check_assist_not_applicable(
 409             merge_match_arms,
 410             r#"
 411 enum MyEnum {
 412     OptionA(f32),
 413     OptionB(f64),
 414     OptionC(f32)
 415 }
 416
 417 fn func(e: MyEnum) {
 418     match e {
 419         MyEnum::OptionA(x) => $0x.classify(),
 420         MyEnum::OptionB(x) => x.classify(),
 421         MyEnum::OptionC(x) => x.classify(),
 422     };
 423 }
 424 "#,
 425         );
 426     }
 427
 428     #[test]
 429     fn merge_match_arms_same_type_different_number_of_fields() {
 430         check_assist_not_applicable(
 431             merge_match_arms,
 432             r#"//- minicore: result
 433 fn func() {
 434     match Result::<(f64, f64, f64), (f64, f64)>::Ok((0f64, 0f64, 0f64)) {
 435         Ok(x) => $0x.1.classify(),
 436         Err(x) => x.1.classify()
 437     };
 438 }
 439 "#,
 440         );
 441     }
 442
 443     #[test]
 444     fn merge_match_same_destructuring_different_types() {
 445         check_assist_not_applicable(
 446             merge_match_arms,
 447             r#"
 448 struct Point {
 449     x: i32,
 450     y: i32,
 451 }
 452
 453 fn func() {
 454     let p = Point { x: 0, y: 7 };
 455
 456     match p {
 457         Point { x, y: 0 } => $0"",
 458         Point { x: 0, y } => "",
 459         Point { x, y } => "",
 460     };
 461 }
 462 "#,
 463         );
 464     }
 465
 466     #[test]
 467     fn merge_match_arms_range() {
 468         check_assist(
 469             merge_match_arms,
 470             r#"
 471 fn func() {
 472     let x = 'c';
 473
 474     match x {
 475         'a'..='j' => $0"",
 476         'c'..='z' => "",
 477         _ => "other",
 478     };
 479 }
 480 "#,
 481             r#"
 482 fn func() {
 483     let x = 'c';
 484
 485     match x {
 486         'a'..='j' | 'c'..='z' => "",
 487         _ => "other",
 488     };
 489 }
 490 "#,
 491         );
 492     }
 493
 494     #[test]
 495     fn merge_match_arms_enum_without_field() {
 496         check_assist_not_applicable(
 497             merge_match_arms,
 498             r#"
 499 enum MyEnum {
 500     NoField,
 501     AField(u8)
 502 }
 503
 504 fn func(x: MyEnum) {
 505     match x {
 506         MyEnum::NoField => $0"",
 507         MyEnum::AField(x) => ""
 508     };
 509 }
 510         "#,
 511         )
 512     }
 513
 514     #[test]
 515     fn merge_match_arms_enum_destructuring_different_types() {
 516         check_assist_not_applicable(
 517             merge_match_arms,
 518             r#"
 519 enum MyEnum {
 520     Move { x: i32, y: i32 },
 521     Write(String),
 522 }
 523
 524 fn func(x: MyEnum) {
 525     match x {
 526         MyEnum::Move { x, y } => $0"",
 527         MyEnum::Write(text) => "",
 528     };
 529 }
 530         "#,
 531         )
 532     }
 533
 534     #[test]
 535     fn merge_match_arms_enum_destructuring_same_types() {
 536         check_assist(
 537             merge_match_arms,
 538             r#"
 539 enum MyEnum {
 540     Move { x: i32, y: i32 },
 541     Crawl { x: i32, y: i32 }
 542 }
 543
 544 fn func(x: MyEnum) {
 545     match x {
 546         MyEnum::Move { x, y } => $0"",
 547         MyEnum::Crawl { x, y } => "",
 548     };
 549 }
 550         "#,
 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 } | MyEnum::Crawl { x, y } => "",
 560     };
 561 }
 562         "#,
 563         )
 564     }
 565
 566     #[test]
 567     fn merge_match_arms_enum_destructuring_same_types_different_name() {
 568         check_assist_not_applicable(
 569             merge_match_arms,
 570             r#"
 571 enum MyEnum {
 572     Move { x: i32, y: i32 },
 573     Crawl { a: i32, b: i32 }
 574 }
 575
 576 fn func(x: MyEnum) {
 577     match x {
 578         MyEnum::Move { x, y } => $0"",
 579         MyEnum::Crawl { a, b } => "",
 580     };
 581 }
 582         "#,
 583         )
 584     }
 585
 586     #[test]
 587     fn merge_match_arms_enum_nested_pattern_different_names() {
 588         check_assist_not_applicable(
 589             merge_match_arms,
 590             r#"
 591 enum Color {
 592     Rgb(i32, i32, i32),
 593     Hsv(i32, i32, i32),
 594 }
 595
 596 enum Message {
 597     Quit,
 598     Move { x: i32, y: i32 },
 599     Write(String),
 600     ChangeColor(Color),
 601 }
 602
 603 fn main(msg: Message) {
 604     match msg {
 605         Message::ChangeColor(Color::Rgb(r, g, b)) => $0"",
 606         Message::ChangeColor(Color::Hsv(h, s, v)) => "",
 607         _ => "other"
 608     };
 609 }
 610         "#,
 611         )
 612     }
 613
 614     #[test]
 615     fn merge_match_arms_enum_nested_pattern_same_names() {
 616         check_assist(
 617             merge_match_arms,
 618             r#"
 619 enum Color {
 620     Rgb(i32, i32, i32),
 621     Hsv(i32, i32, i32),
 622 }
 623
 624 enum Message {
 625     Quit,
 626     Move { x: i32, y: i32 },
 627     Write(String),
 628     ChangeColor(Color),
 629 }
 630
 631 fn main(msg: Message) {
 632     match msg {
 633         Message::ChangeColor(Color::Rgb(a, b, c)) => $0"",
 634         Message::ChangeColor(Color::Hsv(a, b, c)) => "",
 635         _ => "other"
 636     };
 637 }
 638         "#,
 639             r#"
 640 enum Color {
 641     Rgb(i32, i32, i32),
 642     Hsv(i32, i32, i32),
 643 }
 644
 645 enum Message {
 646     Quit,
 647     Move { x: i32, y: i32 },
 648     Write(String),
 649     ChangeColor(Color),
 650 }
 651
 652 fn main(msg: Message) {
 653     match msg {
 654         Message::ChangeColor(Color::Rgb(a, b, c)) | Message::ChangeColor(Color::Hsv(a, b, c)) => "",
 655         _ => "other"
 656     };
 657 }
 658         "#,
 659         )
 660     }
 661
 662     #[test]
 663     fn merge_match_arms_enum_destructuring_with_ignore() {
 664         check_assist(
 665             merge_match_arms,
 666             r#"
 667 enum MyEnum {
 668     Move { x: i32, a: i32 },
 669     Crawl { x: i32, b: i32 }
 670 }
 671
 672 fn func(x: MyEnum) {
 673     match x {
 674         MyEnum::Move { x, .. } => $0"",
 675         MyEnum::Crawl { x, .. } => "",
 676     };
 677 }
 678         "#,
 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, .. } | MyEnum::Crawl { x, .. } => "",
 688     };
 689 }
 690         "#,
 691         )
 692     }
 693
 694     #[test]
 695     fn merge_match_arms_nested_with_conflicting_identifier() {
 696         check_assist_not_applicable(
 697             merge_match_arms,
 698             r#"
 699 enum Color {
 700     Rgb(i32, i32, i32),
 701     Hsv(i32, i32, i32),
 702 }
 703
 704 enum Message {
 705     Move { x: i32, y: i32 },
 706     ChangeColor(u8, Color),
 707 }
 708
 709 fn main(msg: Message) {
 710     match msg {
 711         Message::ChangeColor(x, Color::Rgb(y, b, c)) => $0"",
 712         Message::ChangeColor(y, Color::Hsv(x, b, c)) => "",
 713         _ => "other"
 714     };
 715 }
 716         "#,
 717         )
 718     }
 719
 720     #[test]
 721     fn merge_match_arms_tuple() {
 722         check_assist_not_applicable(
 723             merge_match_arms,
 724             r#"
 725 fn func() {
 726     match (0, "boo") {
 727         (x, y) => $0"",
 728         (y, x) => "",
 729     };
 730 }
 731         "#,
 732         )
 733     }
 734 }