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