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 use test_utils::mark;
   9
  10 use crate::{
  11     utils::{does_pat_match_variant, 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 //         $0
  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_with_descend::<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 = ctx.sema.original_range(match_expr.syntax()).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 = ctx.sema.original_range(match_arm_list.syntax()).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 resolve_enum_def(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<hir::Enum> {
 151     sema.type_of_expr(&expr)?.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
 152         Some(Adt::Enum(e)) => Some(e),
 153         _ => None,
 154     })
 155 }
 156
 157 fn resolve_tuple_of_enum_def(
 158     sema: &Semantics<RootDatabase>,
 159     expr: &ast::Expr,
 160 ) -> Option<Vec<hir::Enum>> {
 161     sema.type_of_expr(&expr)?
 162         .tuple_fields(sema.db)
 163         .iter()
 164         .map(|ty| {
 165             ty.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
 166                 Some(Adt::Enum(e)) => Some(e),
 167                 // For now we only handle expansion for a tuple of enums. Here
 168                 // we map non-enum items to None and rely on `collect` to
 169                 // convert Vec<Option<hir::Enum>> into Option<Vec<hir::Enum>>.
 170                 _ => None,
 171             })
 172         })
 173         .collect()
 174 }
 175
 176 fn build_pat(db: &RootDatabase, module: hir::Module, var: hir::Variant) -> Option<ast::Pat> {
 177     let path = mod_path_to_ast(&module.find_use_path(db, ModuleDef::from(var))?);
 178
 179     // FIXME: use HIR for this; it doesn't currently expose struct vs. tuple vs. unit variants though
 180     let pat: ast::Pat = match var.source(db)?.value.kind() {
 181         ast::StructKind::Tuple(field_list) => {
 182             let pats = iter::repeat(make::wildcard_pat().into()).take(field_list.fields().count());
 183             make::tuple_struct_pat(path, pats).into()
 184         }
 185         ast::StructKind::Record(field_list) => {
 186             let pats = field_list.fields().map(|f| make::ident_pat(f.name().unwrap()).into());
 187             make::record_pat(path, pats).into()
 188         }
 189         ast::StructKind::Unit => make::path_pat(path),
 190     };
 191
 192     Some(pat)
 193 }
 194
 195 #[cfg(test)]
 196 mod tests {
 197     use ide_db::helpers::FamousDefs;
 198     use test_utils::mark;
 199
 200     use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};
 201
 202     use super::fill_match_arms;
 203
 204     #[test]
 205     fn all_match_arms_provided() {
 206         check_assist_not_applicable(
 207             fill_match_arms,
 208             r#"
 209             enum A {
 210                 As,
 211                 Bs{x:i32, y:Option<i32>},
 212                 Cs(i32, Option<i32>),
 213             }
 214             fn main() {
 215                 match A::As$0 {
 216                     A::As,
 217                     A::Bs{x,y:Some(_)} => {}
 218                     A::Cs(_, Some(_)) => {}
 219                 }
 220             }
 221             "#,
 222         );
 223     }
 224
 225     #[test]
 226     fn tuple_of_non_enum() {
 227         // for now this case is not handled, although it potentially could be
 228         // in the future
 229         check_assist_not_applicable(
 230             fill_match_arms,
 231             r#"
 232             fn main() {
 233                 match (0, false)$0 {
 234                 }
 235             }
 236             "#,
 237         );
 238     }
 239
 240     #[test]
 241     fn partial_fill_record_tuple() {
 242         check_assist(
 243             fill_match_arms,
 244             r#"
 245             enum A {
 246                 As,
 247                 Bs { x: i32, y: Option<i32> },
 248                 Cs(i32, Option<i32>),
 249             }
 250             fn main() {
 251                 match A::As$0 {
 252                     A::Bs { x, y: Some(_) } => {}
 253                     A::Cs(_, Some(_)) => {}
 254                 }
 255             }
 256             "#,
 257             r#"
 258             enum A {
 259                 As,
 260                 Bs { x: i32, y: Option<i32> },
 261                 Cs(i32, Option<i32>),
 262             }
 263             fn main() {
 264                 match A::As {
 265                     A::Bs { x, y: Some(_) } => {}
 266                     A::Cs(_, Some(_)) => {}
 267                     $0A::As => {}
 268                 }
 269             }
 270             "#,
 271         );
 272     }
 273
 274     #[test]
 275     fn partial_fill_option() {
 276         check_assist(
 277             fill_match_arms,
 278             r#"
 279 enum Option<T> { Some(T), None }
 280 use Option::*;
 281
 282 fn main() {
 283     match None$0 {
 284         None => {}
 285     }
 286 }
 287             "#,
 288             r#"
 289 enum Option<T> { Some(T), None }
 290 use Option::*;
 291
 292 fn main() {
 293     match None {
 294         None => {}
 295         Some(${0:_}) => {}
 296     }
 297 }
 298             "#,
 299         );
 300     }
 301
 302     #[test]
 303     fn partial_fill_or_pat() {
 304         check_assist(
 305             fill_match_arms,
 306             r#"
 307 enum A { As, Bs, Cs(Option<i32>) }
 308 fn main() {
 309     match A::As$0 {
 310         A::Cs(_) | A::Bs => {}
 311     }
 312 }
 313 "#,
 314             r#"
 315 enum A { As, Bs, Cs(Option<i32>) }
 316 fn main() {
 317     match A::As {
 318         A::Cs(_) | A::Bs => {}
 319         $0A::As => {}
 320     }
 321 }
 322 "#,
 323         );
 324     }
 325
 326     #[test]
 327     fn partial_fill() {
 328         check_assist(
 329             fill_match_arms,
 330             r#"
 331 enum A { As, Bs, Cs, Ds(String), Es(B) }
 332 enum B { Xs, Ys }
 333 fn main() {
 334     match A::As$0 {
 335         A::Bs if 0 < 1 => {}
 336         A::Ds(_value) => { let x = 1; }
 337         A::Es(B::Xs) => (),
 338     }
 339 }
 340 "#,
 341             r#"
 342 enum A { As, Bs, Cs, Ds(String), Es(B) }
 343 enum B { Xs, Ys }
 344 fn main() {
 345     match A::As {
 346         A::Bs if 0 < 1 => {}
 347         A::Ds(_value) => { let x = 1; }
 348         A::Es(B::Xs) => (),
 349         $0A::As => {}
 350         A::Cs => {}
 351     }
 352 }
 353 "#,
 354         );
 355     }
 356
 357     #[test]
 358     fn partial_fill_bind_pat() {
 359         check_assist(
 360             fill_match_arms,
 361             r#"
 362 enum A { As, Bs, Cs(Option<i32>) }
 363 fn main() {
 364     match A::As$0 {
 365         A::As(_) => {}
 366         a @ A::Bs(_) => {}
 367     }
 368 }
 369 "#,
 370             r#"
 371 enum A { As, Bs, Cs(Option<i32>) }
 372 fn main() {
 373     match A::As {
 374         A::As(_) => {}
 375         a @ A::Bs(_) => {}
 376         A::Cs(${0:_}) => {}
 377     }
 378 }
 379 "#,
 380         );
 381     }
 382
 383     #[test]
 384     fn fill_match_arms_empty_body() {
 385         check_assist(
 386             fill_match_arms,
 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$0 {}
 393 }
 394 "#,
 395             r#"
 396 enum A { As, Bs, Cs(String), Ds(String, String), Es { x: usize, y: usize } }
 397
 398 fn main() {
 399     let a = A::As;
 400     match a {
 401         $0A::As => {}
 402         A::Bs => {}
 403         A::Cs(_) => {}
 404         A::Ds(_, _) => {}
 405         A::Es { x, y } => {}
 406     }
 407 }
 408 "#,
 409         );
 410     }
 411
 412     #[test]
 413     fn fill_match_arms_tuple_of_enum() {
 414         check_assist(
 415             fill_match_arms,
 416             r#"
 417             enum A { One, Two }
 418             enum B { One, Two }
 419
 420             fn main() {
 421                 let a = A::One;
 422                 let b = B::One;
 423                 match (a$0, b) {}
 424             }
 425             "#,
 426             r#"
 427             enum A { One, Two }
 428             enum B { One, Two }
 429
 430             fn main() {
 431                 let a = A::One;
 432                 let b = B::One;
 433                 match (a, b) {
 434                     $0(A::One, B::One) => {}
 435                     (A::One, B::Two) => {}
 436                     (A::Two, B::One) => {}
 437                     (A::Two, B::Two) => {}
 438                 }
 439             }
 440             "#,
 441         );
 442     }
 443
 444     #[test]
 445     fn fill_match_arms_tuple_of_enum_ref() {
 446         check_assist(
 447             fill_match_arms,
 448             r#"
 449             enum A { One, Two }
 450             enum B { One, Two }
 451
 452             fn main() {
 453                 let a = A::One;
 454                 let b = B::One;
 455                 match (&a$0, &b) {}
 456             }
 457             "#,
 458             r#"
 459             enum A { One, Two }
 460             enum B { One, Two }
 461
 462             fn main() {
 463                 let a = A::One;
 464                 let b = B::One;
 465                 match (&a, &b) {
 466                     $0(A::One, B::One) => {}
 467                     (A::One, B::Two) => {}
 468                     (A::Two, B::One) => {}
 469                     (A::Two, B::Two) => {}
 470                 }
 471             }
 472             "#,
 473         );
 474     }
 475
 476     #[test]
 477     fn fill_match_arms_tuple_of_enum_partial() {
 478         check_assist_not_applicable(
 479             fill_match_arms,
 480             r#"
 481             enum A { One, Two }
 482             enum B { One, Two }
 483
 484             fn main() {
 485                 let a = A::One;
 486                 let b = B::One;
 487                 match (a$0, b) {
 488                     (A::Two, B::One) => {}
 489                 }
 490             }
 491             "#,
 492         );
 493     }
 494
 495     #[test]
 496     fn fill_match_arms_tuple_of_enum_not_applicable() {
 497         check_assist_not_applicable(
 498             fill_match_arms,
 499             r#"
 500             enum A { One, Two }
 501             enum B { One, Two }
 502
 503             fn main() {
 504                 let a = A::One;
 505                 let b = B::One;
 506                 match (a$0, b) {
 507                     (A::Two, B::One) => {}
 508                     (A::One, B::One) => {}
 509                     (A::One, B::Two) => {}
 510                     (A::Two, B::Two) => {}
 511                 }
 512             }
 513             "#,
 514         );
 515     }
 516
 517     #[test]
 518     fn fill_match_arms_single_element_tuple_of_enum() {
 519         // For now we don't hande the case of a single element tuple, but
 520         // we could handle this in the future if `make::tuple_pat` allowed
 521         // creating a tuple with a single pattern.
 522         check_assist_not_applicable(
 523             fill_match_arms,
 524             r#"
 525             enum A { One, Two }
 526
 527             fn main() {
 528                 let a = A::One;
 529                 match (a$0, ) {
 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         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 }