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