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