crates/ra_assists/src/handlers/fill_match_arms.rs

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