crates/ide_assists/src/handlers/fill_match_arms.rs

   1 use std::iter;
   2
   3 use hir::{Adt, HasSource, ModuleDef, Semantics};
   4 use ide_db::helpers::{mod_path_to_ast, FamousDefs};
   5 use ide_db::RootDatabase;
   6 use itertools::Itertools;
   7 use syntax::ast::{self, make, AstNode, MatchArm, NameOwner, Pat};
   8
   9 use crate::{
  10     utils::{does_pat_match_variant, render_snippet, Cursor},
  11     AssistContext, AssistId, AssistKind, Assists,
  12 };
  13
  14 // Assist: fill_match_arms
  15 //
  16 // Adds missing clauses to a `match` expression.
  17 //
  18 // ```
  19 // enum Action { Move { distance: u32 }, Stop }
  20 //
  21 // fn handle(action: Action) {
  22 //     match action {
  23 //         $0
  24 //     }
  25 // }
  26 // ```
  27 // ->
  28 // ```
  29 // enum Action { Move { distance: u32 }, Stop }
  30 //
  31 // fn handle(action: Action) {
  32 //     match action {
  33 //         $0Action::Move { distance } => {}
  34 //         Action::Stop => {}
  35 //     }
  36 // }
  37 // ```
  38 pub(crate) fn fill_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  39     let match_expr = ctx.find_node_at_offset_with_descend::<ast::MatchExpr>()?;
  40     let match_arm_list = match_expr.match_arm_list()?;
  41
  42     let expr = match_expr.expr()?;
  43
  44     let mut arms: Vec<MatchArm> = match_arm_list.arms().collect();
  45     if arms.len() == 1 {
  46         if let Some(Pat::WildcardPat(..)) = arms[0].pat() {
  47             arms.clear();
  48         }
  49     }
  50
  51     let module = ctx.sema.scope(expr.syntax()).module()?;
  52
  53     let missing_arms: Vec<MatchArm> = if let Some(enum_def) = resolve_enum_def(&ctx.sema, &expr) {
  54         let variants = enum_def.variants(ctx.db());
  55
  56         let mut variants = variants
  57             .into_iter()
  58             .filter_map(|variant| build_pat(ctx.db(), module, variant))
  59             .filter(|variant_pat| is_variant_missing(&mut arms, variant_pat))
  60             .map(|pat| make::match_arm(iter::once(pat), make::expr_empty_block()))
  61             .collect::<Vec<_>>();
  62         if Some(enum_def) == FamousDefs(&ctx.sema, Some(module.krate())).core_option_Option() {
  63             // Match `Some` variant first.
  64             cov_mark::hit!(option_order);
  65             variants.reverse()
  66         }
  67         variants
  68     } else if let Some(enum_defs) = resolve_tuple_of_enum_def(&ctx.sema, &expr) {
  69         // Partial fill not currently supported for tuple of enums.
  70         if !arms.is_empty() {
  71             return None;
  72         }
  73
  74         // When calculating the match arms for a tuple of enums, we want
  75         // to create a match arm for each possible combination of enum
  76         // values. The `multi_cartesian_product` method transforms
  77         // Vec<Vec<EnumVariant>> into Vec<(EnumVariant, .., EnumVariant)>
  78         // where each tuple represents a proposed match arm.
  79         enum_defs
  80             .into_iter()
  81             .map(|enum_def| enum_def.variants(ctx.db()))
  82             .multi_cartesian_product()
  83             .map(|variants| {
  84                 let patterns =
  85                     variants.into_iter().filter_map(|variant| build_pat(ctx.db(), module, variant));
  86                 ast::Pat::from(make::tuple_pat(patterns))
  87             })
  88             .filter(|variant_pat| is_variant_missing(&mut arms, variant_pat))
  89             .map(|pat| make::match_arm(iter::once(pat), make::expr_empty_block()))
  90             .collect()
  91     } else {
  92         return None;
  93     };
  94
  95     if missing_arms.is_empty() {
  96         return None;
  97     }
  98
  99     let target = ctx.sema.original_range(match_expr.syntax()).range;
 100     acc.add(
 101         AssistId("fill_match_arms", AssistKind::QuickFix),
 102         "Fill match arms",
 103         target,
 104         |builder| {
 105             let new_arm_list = match_arm_list.remove_placeholder();
 106             let n_old_arms = new_arm_list.arms().count();
 107             let new_arm_list = new_arm_list.append_arms(missing_arms);
 108             let first_new_arm = new_arm_list.arms().nth(n_old_arms);
 109             let old_range = ctx.sema.original_range(match_arm_list.syntax()).range;
 110             match (first_new_arm, ctx.config.snippet_cap) {
 111                 (Some(first_new_arm), Some(cap)) => {
 112                     let extend_lifetime;
 113                     let cursor =
 114                         match first_new_arm.syntax().descendants().find_map(ast::WildcardPat::cast)
 115                         {
 116                             Some(it) => {
 117                                 extend_lifetime = it.syntax().clone();
 118                                 Cursor::Replace(&extend_lifetime)
 119                             }
 120                             None => Cursor::Before(first_new_arm.syntax()),
 121                         };
 122                     let snippet = render_snippet(cap, new_arm_list.syntax(), cursor);
 123                     builder.replace_snippet(cap, old_range, snippet);
 124                 }
 125                 _ => builder.replace(old_range, new_arm_list.to_string()),
 126             }
 127         },
 128     )
 129 }
 130
 131 fn is_variant_missing(existing_arms: &mut Vec<MatchArm>, var: &Pat) -> bool {
 132     existing_arms.iter().filter_map(|arm| arm.pat()).all(|pat| {
 133         // Special casee OrPat as separate top-level pats
 134         let top_level_pats: Vec<Pat> = match pat {
 135             Pat::OrPat(pats) => pats.pats().collect::<Vec<_>>(),
 136             _ => vec![pat],
 137         };
 138
 139         !top_level_pats.iter().any(|pat| does_pat_match_variant(pat, var))
 140     })
 141 }
 142
 143 fn resolve_enum_def(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<hir::Enum> {
 144     sema.type_of_expr(&expr)?.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
 145         Some(Adt::Enum(e)) => Some(e),
 146         _ => None,
 147     })
 148 }
 149
 150 fn resolve_tuple_of_enum_def(
 151     sema: &Semantics<RootDatabase>,
 152     expr: &ast::Expr,
 153 ) -> Option<Vec<hir::Enum>> {
 154     sema.type_of_expr(&expr)?
 155         .tuple_fields(sema.db)
 156         .iter()
 157         .map(|ty| {
 158             ty.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
 159                 Some(Adt::Enum(e)) => Some(e),
 160                 // For now we only handle expansion for a tuple of enums. Here
 161                 // we map non-enum items to None and rely on `collect` to
 162                 // convert Vec<Option<hir::Enum>> into Option<Vec<hir::Enum>>.
 163                 _ => None,
 164             })
 165         })
 166         .collect()
 167 }
 168
 169 fn build_pat(db: &RootDatabase, module: hir::Module, var: hir::Variant) -> Option<ast::Pat> {
 170     let path = mod_path_to_ast(&module.find_use_path(db, ModuleDef::from(var))?);
 171
 172     // FIXME: use HIR for this; it doesn't currently expose struct vs. tuple vs. unit variants though
 173     let pat: ast::Pat = match var.source(db)?.value.kind() {
 174         ast::StructKind::Tuple(field_list) => {
 175             let pats = iter::repeat(make::wildcard_pat().into()).take(field_list.fields().count());
 176             make::tuple_struct_pat(path, pats).into()
 177         }
 178         ast::StructKind::Record(field_list) => {
 179             let pats = field_list.fields().map(|f| make::ident_pat(f.name().unwrap()).into());
 180             make::record_pat(path, pats).into()
 181         }
 182         ast::StructKind::Unit => make::path_pat(path),
 183     };
 184
 185     Some(pat)
 186 }
 187
 188 #[cfg(test)]
 189 mod tests {
 190     use ide_db::helpers::FamousDefs;
 191
 192     use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};
 193
 194     use super::fill_match_arms;
 195
 196     #[test]
 197     fn all_match_arms_provided() {
 198         check_assist_not_applicable(
 199             fill_match_arms,
 200             r#"
 201             enum A {
 202                 As,
 203                 Bs{x:i32, y:Option<i32>},
 204                 Cs(i32, Option<i32>),
 205             }
 206             fn main() {
 207                 match A::As$0 {
 208                     A::As,
 209                     A::Bs{x,y:Some(_)} => {}
 210                     A::Cs(_, Some(_)) => {}
 211                 }
 212             }
 213             "#,
 214         );
 215     }
 216
 217     #[test]
 218     fn tuple_of_non_enum() {
 219         // for now this case is not handled, although it potentially could be
 220         // in the future
 221         check_assist_not_applicable(
 222             fill_match_arms,
 223             r#"
 224             fn main() {
 225                 match (0, false)$0 {
 226                 }
 227             }
 228             "#,
 229         );
 230     }
 231
 232     #[test]
 233     fn partial_fill_record_tuple() {
 234         check_assist(
 235             fill_match_arms,
 236             r#"
 237             enum A {
 238                 As,
 239                 Bs { x: i32, y: Option<i32> },
 240                 Cs(i32, Option<i32>),
 241             }
 242             fn main() {
 243                 match A::As$0 {
 244                     A::Bs { x, y: Some(_) } => {}
 245                     A::Cs(_, Some(_)) => {}
 246                 }
 247             }
 248             "#,
 249             r#"
 250             enum A {
 251                 As,
 252                 Bs { x: i32, y: Option<i32> },
 253                 Cs(i32, Option<i32>),
 254             }
 255             fn main() {
 256                 match A::As {
 257                     A::Bs { x, y: Some(_) } => {}
 258                     A::Cs(_, Some(_)) => {}
 259                     $0A::As => {}
 260                 }
 261             }
 262             "#,
 263         );
 264     }
 265
 266     #[test]
 267     fn partial_fill_option() {
 268         check_assist(
 269             fill_match_arms,
 270             r#"
 271 enum Option<T> { Some(T), None }
 272 use Option::*;
 273
 274 fn main() {
 275     match None$0 {
 276         None => {}
 277     }
 278 }
 279             "#,
 280             r#"
 281 enum Option<T> { Some(T), None }
 282 use Option::*;
 283
 284 fn main() {
 285     match None {
 286         None => {}
 287         Some(${0:_}) => {}
 288     }
 289 }
 290             "#,
 291         );
 292     }
 293
 294     #[test]
 295     fn partial_fill_or_pat() {
 296         check_assist(
 297             fill_match_arms,
 298             r#"
 299 enum A { As, Bs, Cs(Option<i32>) }
 300 fn main() {
 301     match A::As$0 {
 302         A::Cs(_) | A::Bs => {}
 303     }
 304 }
 305 "#,
 306             r#"
 307 enum A { As, Bs, Cs(Option<i32>) }
 308 fn main() {
 309     match A::As {
 310         A::Cs(_) | A::Bs => {}
 311         $0A::As => {}
 312     }
 313 }
 314 "#,
 315         );
 316     }
 317
 318     #[test]
 319     fn partial_fill() {
 320         check_assist(
 321             fill_match_arms,
 322             r#"
 323 enum A { As, Bs, Cs, Ds(String), Es(B) }
 324 enum B { Xs, Ys }
 325 fn main() {
 326     match A::As$0 {
 327         A::Bs if 0 < 1 => {}
 328         A::Ds(_value) => { let x = 1; }
 329         A::Es(B::Xs) => (),
 330     }
 331 }
 332 "#,
 333             r#"
 334 enum A { As, Bs, Cs, Ds(String), Es(B) }
 335 enum B { Xs, Ys }
 336 fn main() {
 337     match A::As {
 338         A::Bs if 0 < 1 => {}
 339         A::Ds(_value) => { let x = 1; }
 340         A::Es(B::Xs) => (),
 341         $0A::As => {}
 342         A::Cs => {}
 343     }
 344 }
 345 "#,
 346         );
 347     }
 348
 349     #[test]
 350     fn partial_fill_bind_pat() {
 351         check_assist(
 352             fill_match_arms,
 353             r#"
 354 enum A { As, Bs, Cs(Option<i32>) }
 355 fn main() {
 356     match A::As$0 {
 357         A::As(_) => {}
 358         a @ A::Bs(_) => {}
 359     }
 360 }
 361 "#,
 362             r#"
 363 enum A { As, Bs, Cs(Option<i32>) }
 364 fn main() {
 365     match A::As {
 366         A::As(_) => {}
 367         a @ A::Bs(_) => {}
 368         A::Cs(${0:_}) => {}
 369     }
 370 }
 371 "#,
 372         );
 373     }
 374
 375     #[test]
 376     fn fill_match_arms_empty_body() {
 377         check_assist(
 378             fill_match_arms,
 379             r#"
 380 enum A { As, Bs, Cs(String), Ds(String, String), Es { x: usize, y: usize } }
 381
 382 fn main() {
 383     let a = A::As;
 384     match a$0 {}
 385 }
 386 "#,
 387             r#"
 388 enum A { As, Bs, Cs(String), Ds(String, String), Es { x: usize, y: usize } }
 389
 390 fn main() {
 391     let a = A::As;
 392     match a {
 393         $0A::As => {}
 394         A::Bs => {}
 395         A::Cs(_) => {}
 396         A::Ds(_, _) => {}
 397         A::Es { x, y } => {}
 398     }
 399 }
 400 "#,
 401         );
 402     }
 403
 404     #[test]
 405     fn fill_match_arms_tuple_of_enum() {
 406         check_assist(
 407             fill_match_arms,
 408             r#"
 409             enum A { One, Two }
 410             enum B { One, Two }
 411
 412             fn main() {
 413                 let a = A::One;
 414                 let b = B::One;
 415                 match (a$0, b) {}
 416             }
 417             "#,
 418             r#"
 419             enum A { One, Two }
 420             enum B { One, Two }
 421
 422             fn main() {
 423                 let a = A::One;
 424                 let b = B::One;
 425                 match (a, b) {
 426                     $0(A::One, B::One) => {}
 427                     (A::One, B::Two) => {}
 428                     (A::Two, B::One) => {}
 429                     (A::Two, B::Two) => {}
 430                 }
 431             }
 432             "#,
 433         );
 434     }
 435
 436     #[test]
 437     fn fill_match_arms_tuple_of_enum_ref() {
 438         check_assist(
 439             fill_match_arms,
 440             r#"
 441             enum A { One, Two }
 442             enum B { One, Two }
 443
 444             fn main() {
 445                 let a = A::One;
 446                 let b = B::One;
 447                 match (&a$0, &b) {}
 448             }
 449             "#,
 450             r#"
 451             enum A { One, Two }
 452             enum B { One, Two }
 453
 454             fn main() {
 455                 let a = A::One;
 456                 let b = B::One;
 457                 match (&a, &b) {
 458                     $0(A::One, B::One) => {}
 459                     (A::One, B::Two) => {}
 460                     (A::Two, B::One) => {}
 461                     (A::Two, B::Two) => {}
 462                 }
 463             }
 464             "#,
 465         );
 466     }
 467
 468     #[test]
 469     fn fill_match_arms_tuple_of_enum_partial() {
 470         check_assist_not_applicable(
 471             fill_match_arms,
 472             r#"
 473             enum A { One, Two }
 474             enum B { One, Two }
 475
 476             fn main() {
 477                 let a = A::One;
 478                 let b = B::One;
 479                 match (a$0, b) {
 480                     (A::Two, B::One) => {}
 481                 }
 482             }
 483             "#,
 484         );
 485     }
 486
 487     #[test]
 488     fn fill_match_arms_tuple_of_enum_not_applicable() {
 489         check_assist_not_applicable(
 490             fill_match_arms,
 491             r#"
 492             enum A { One, Two }
 493             enum B { One, Two }
 494
 495             fn main() {
 496                 let a = A::One;
 497                 let b = B::One;
 498                 match (a$0, b) {
 499                     (A::Two, B::One) => {}
 500                     (A::One, B::One) => {}
 501                     (A::One, B::Two) => {}
 502                     (A::Two, B::Two) => {}
 503                 }
 504             }
 505             "#,
 506         );
 507     }
 508
 509     #[test]
 510     fn fill_match_arms_single_element_tuple_of_enum() {
 511         check_assist(
 512             fill_match_arms,
 513             r#"
 514             enum A { One, Two }
 515
 516             fn main() {
 517                 let a = A::One;
 518                 match (a$0, ) {
 519                 }
 520             }
 521             "#,
 522             r#"
 523             enum A { One, Two }
 524
 525             fn main() {
 526                 let a = A::One;
 527                 match (a, ) {
 528                     $0(A::One,) => {}
 529                     (A::Two,) => {}
 530                 }
 531             }
 532             "#,
 533         );
 534     }
 535
 536     #[test]
 537     fn test_fill_match_arm_refs() {
 538         check_assist(
 539             fill_match_arms,
 540             r#"
 541             enum A { As }
 542
 543             fn foo(a: &A) {
 544                 match a$0 {
 545                 }
 546             }
 547             "#,
 548             r#"
 549             enum A { As }
 550
 551             fn foo(a: &A) {
 552                 match a {
 553                     $0A::As => {}
 554                 }
 555             }
 556             "#,
 557         );
 558
 559         check_assist(
 560             fill_match_arms,
 561             r#"
 562             enum A {
 563                 Es { x: usize, y: usize }
 564             }
 565
 566             fn foo(a: &mut A) {
 567                 match a$0 {
 568                 }
 569             }
 570             "#,
 571             r#"
 572             enum A {
 573                 Es { x: usize, y: usize }
 574             }
 575
 576             fn foo(a: &mut A) {
 577                 match a {
 578                     $0A::Es { x, y } => {}
 579                 }
 580             }
 581             "#,
 582         );
 583     }
 584
 585     #[test]
 586     fn fill_match_arms_target() {
 587         check_assist_target(
 588             fill_match_arms,
 589             r#"
 590             enum E { X, Y }
 591
 592             fn main() {
 593                 match E::X$0 {}
 594             }
 595             "#,
 596             "match E::X {}",
 597         );
 598     }
 599
 600     #[test]
 601     fn fill_match_arms_trivial_arm() {
 602         check_assist(
 603             fill_match_arms,
 604             r#"
 605             enum E { X, Y }
 606
 607             fn main() {
 608                 match E::X {
 609                     $0_ => {}
 610                 }
 611             }
 612             "#,
 613             r#"
 614             enum E { X, Y }
 615
 616             fn main() {
 617                 match E::X {
 618                     $0E::X => {}
 619                     E::Y => {}
 620                 }
 621             }
 622             "#,
 623         );
 624     }
 625
 626     #[test]
 627     fn fill_match_arms_qualifies_path() {
 628         check_assist(
 629             fill_match_arms,
 630             r#"
 631             mod foo { pub enum E { X, Y } }
 632             use foo::E::X;
 633
 634             fn main() {
 635                 match X {
 636                     $0
 637                 }
 638             }
 639             "#,
 640             r#"
 641             mod foo { pub enum E { X, Y } }
 642             use foo::E::X;
 643
 644             fn main() {
 645                 match X {
 646                     $0X => {}
 647                     foo::E::Y => {}
 648                 }
 649             }
 650             "#,
 651         );
 652     }
 653
 654     #[test]
 655     fn fill_match_arms_preserves_comments() {
 656         check_assist(
 657             fill_match_arms,
 658             r#"
 659             enum A { One, Two }
 660             fn foo(a: A) {
 661                 match a {
 662                     // foo bar baz$0
 663                     A::One => {}
 664                     // This is where the rest should be
 665                 }
 666             }
 667             "#,
 668             r#"
 669             enum A { One, Two }
 670             fn foo(a: A) {
 671                 match a {
 672                     // foo bar baz
 673                     A::One => {}
 674                     // This is where the rest should be
 675                     $0A::Two => {}
 676                 }
 677             }
 678             "#,
 679         );
 680     }
 681
 682     #[test]
 683     fn fill_match_arms_preserves_comments_empty() {
 684         check_assist(
 685             fill_match_arms,
 686             r#"
 687             enum A { One, Two }
 688             fn foo(a: A) {
 689                 match a {
 690                     // foo bar baz$0
 691                 }
 692             }
 693             "#,
 694             r#"
 695             enum A { One, Two }
 696             fn foo(a: A) {
 697                 match a {
 698                     // foo bar baz
 699                     $0A::One => {}
 700                     A::Two => {}
 701                 }
 702             }
 703             "#,
 704         );
 705     }
 706
 707     #[test]
 708     fn fill_match_arms_placeholder() {
 709         check_assist(
 710             fill_match_arms,
 711             r#"
 712             enum A { One, Two, }
 713             fn foo(a: A) {
 714                 match a$0 {
 715                     _ => (),
 716                 }
 717             }
 718             "#,
 719             r#"
 720             enum A { One, Two, }
 721             fn foo(a: A) {
 722                 match a {
 723                     $0A::One => {}
 724                     A::Two => {}
 725                 }
 726             }
 727             "#,
 728         );
 729     }
 730
 731     #[test]
 732     fn option_order() {
 733         cov_mark::check!(option_order);
 734         let before = r#"
 735 fn foo(opt: Option<i32>) {
 736     match opt$0 {
 737     }
 738 }
 739 "#;
 740         let before = &format!("//- /main.rs crate:main deps:core{}{}", before, FamousDefs::FIXTURE);
 741
 742         check_assist(
 743             fill_match_arms,
 744             before,
 745             r#"
 746 fn foo(opt: Option<i32>) {
 747     match opt {
 748         Some(${0:_}) => {}
 749         None => {}
 750     }
 751 }
 752 "#,
 753         );
 754     }
 755
 756     #[test]
 757     fn works_inside_macro_call() {
 758         check_assist(
 759             fill_match_arms,
 760             r#"
 761 macro_rules! m { ($expr:expr) => {$expr}}
 762 enum Test {
 763     A,
 764     B,
 765     C,
 766 }
 767
 768 fn foo(t: Test) {
 769     m!(match t$0 {});
 770 }"#,
 771             r#"macro_rules! m { ($expr:expr) => {$expr}}
 772 enum Test {
 773     A,
 774     B,
 775     C,
 776 }
 777
 778 fn foo(t: Test) {
 779     m!(match t {
 780     $0Test::A => {}
 781     Test::B => {}
 782     Test::C => {}
 783 });
 784 }"#,
 785         );
 786     }
 787 }