crates/ra_assists/src/handlers/fill_match_arms.rs

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