crates/ide_assists/src/handlers/merge_match_arms.rs

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