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