crates/ide_assists/src/handlers/fill_match_arms.rs

   1 use std::iter;
   2
   3 use hir::{Adt, HasSource, ModuleDef, Semantics};
   4 use ide_db::helpers::{mod_path_to_ast, FamousDefs};
   5 use ide_db::RootDatabase;
   6 use itertools::Itertools;
   7 use syntax::ast::{self, make, AstNode, MatchArm, NameOwner, Pat};
   8
   9 use crate::{
  10     utils::{does_pat_match_variant, render_snippet, Cursor},
  11     AssistContext, AssistId, AssistKind, 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 //         $0
  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_with_descend::<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::WildcardPat(..)) = 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, Some(module.krate())).core_option_Option() {
  63             // Match `Some` variant first.
  64             cov_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 = ctx.sema.original_range(match_expr.syntax()).range;
 106     acc.add(
 107         AssistId("fill_match_arms", AssistKind::QuickFix),
 108         "Fill match arms",
 109         target,
 110         |builder| {
 111             let new_arm_list = match_arm_list.remove_placeholder();
 112             let n_old_arms = new_arm_list.arms().count();
 113             let new_arm_list = new_arm_list.append_arms(missing_arms);
 114             let first_new_arm = new_arm_list.arms().nth(n_old_arms);
 115             let old_range = ctx.sema.original_range(match_arm_list.syntax()).range;
 116             match (first_new_arm, ctx.config.snippet_cap) {
 117                 (Some(first_new_arm), Some(cap)) => {
 118                     let extend_lifetime;
 119                     let cursor =
 120                         match first_new_arm.syntax().descendants().find_map(ast::WildcardPat::cast)
 121                         {
 122                             Some(it) => {
 123                                 extend_lifetime = it.syntax().clone();
 124                                 Cursor::Replace(&extend_lifetime)
 125                             }
 126                             None => Cursor::Before(first_new_arm.syntax()),
 127                         };
 128                     let snippet = render_snippet(cap, new_arm_list.syntax(), cursor);
 129                     builder.replace_snippet(cap, old_range, snippet);
 130                 }
 131                 _ => builder.replace(old_range, new_arm_list.to_string()),
 132             }
 133         },
 134     )
 135 }
 136
 137 fn is_variant_missing(existing_arms: &mut Vec<MatchArm>, var: &Pat) -> bool {
 138     existing_arms.iter().filter_map(|arm| arm.pat()).all(|pat| {
 139         // Special casee OrPat as separate top-level pats
 140         let top_level_pats: Vec<Pat> = match pat {
 141             Pat::OrPat(pats) => pats.pats().collect::<Vec<_>>(),
 142             _ => vec![pat],
 143         };
 144
 145         !top_level_pats.iter().any(|pat| does_pat_match_variant(pat, var))
 146     })
 147 }
 148
 149 fn resolve_enum_def(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<hir::Enum> {
 150     sema.type_of_expr(&expr)?.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
 151         Some(Adt::Enum(e)) => Some(e),
 152         _ => None,
 153     })
 154 }
 155
 156 fn resolve_tuple_of_enum_def(
 157     sema: &Semantics<RootDatabase>,
 158     expr: &ast::Expr,
 159 ) -> Option<Vec<hir::Enum>> {
 160     sema.type_of_expr(&expr)?
 161         .tuple_fields(sema.db)
 162         .iter()
 163         .map(|ty| {
 164             ty.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
 165                 Some(Adt::Enum(e)) => Some(e),
 166                 // For now we only handle expansion for a tuple of enums. Here
 167                 // we map non-enum items to None and rely on `collect` to
 168                 // convert Vec<Option<hir::Enum>> into Option<Vec<hir::Enum>>.
 169                 _ => None,
 170             })
 171         })
 172         .collect()
 173 }
 174
 175 fn build_pat(db: &RootDatabase, module: hir::Module, var: hir::Variant) -> Option<ast::Pat> {
 176     let path = mod_path_to_ast(&module.find_use_path(db, ModuleDef::from(var))?);
 177
 178     // FIXME: use HIR for this; it doesn't currently expose struct vs. tuple vs. unit variants though
 179     let pat: ast::Pat = match var.source(db)?.value.kind() {
 180         ast::StructKind::Tuple(field_list) => {
 181             let pats = iter::repeat(make::wildcard_pat().into()).take(field_list.fields().count());
 182             make::tuple_struct_pat(path, pats).into()
 183         }
 184         ast::StructKind::Record(field_list) => {
 185             let pats = field_list.fields().map(|f| make::ident_pat(f.name().unwrap()).into());
 186             make::record_pat(path, pats).into()
 187         }
 188         ast::StructKind::Unit => make::path_pat(path),
 189     };
 190
 191     Some(pat)
 192 }
 193
 194 #[cfg(test)]
 195 mod tests {
 196     use ide_db::helpers::FamousDefs;
 197
 198     use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};
 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$0 {
 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)$0 {
 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$0 {
 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_option() {
 274         check_assist(
 275             fill_match_arms,
 276             r#"
 277 enum Option<T> { Some(T), None }
 278 use Option::*;
 279
 280 fn main() {
 281     match None$0 {
 282         None => {}
 283     }
 284 }
 285             "#,
 286             r#"
 287 enum Option<T> { Some(T), None }
 288 use Option::*;
 289
 290 fn main() {
 291     match None {
 292         None => {}
 293         Some(${0:_}) => {}
 294     }
 295 }
 296             "#,
 297         );
 298     }
 299
 300     #[test]
 301     fn partial_fill_or_pat() {
 302         check_assist(
 303             fill_match_arms,
 304             r#"
 305 enum A { As, Bs, Cs(Option<i32>) }
 306 fn main() {
 307     match A::As$0 {
 308         A::Cs(_) | A::Bs => {}
 309     }
 310 }
 311 "#,
 312             r#"
 313 enum A { As, Bs, Cs(Option<i32>) }
 314 fn main() {
 315     match A::As {
 316         A::Cs(_) | A::Bs => {}
 317         $0A::As => {}
 318     }
 319 }
 320 "#,
 321         );
 322     }
 323
 324     #[test]
 325     fn partial_fill() {
 326         check_assist(
 327             fill_match_arms,
 328             r#"
 329 enum A { As, Bs, Cs, Ds(String), Es(B) }
 330 enum B { Xs, Ys }
 331 fn main() {
 332     match A::As$0 {
 333         A::Bs if 0 < 1 => {}
 334         A::Ds(_value) => { let x = 1; }
 335         A::Es(B::Xs) => (),
 336     }
 337 }
 338 "#,
 339             r#"
 340 enum A { As, Bs, Cs, Ds(String), Es(B) }
 341 enum B { Xs, Ys }
 342 fn main() {
 343     match A::As {
 344         A::Bs if 0 < 1 => {}
 345         A::Ds(_value) => { let x = 1; }
 346         A::Es(B::Xs) => (),
 347         $0A::As => {}
 348         A::Cs => {}
 349     }
 350 }
 351 "#,
 352         );
 353     }
 354
 355     #[test]
 356     fn partial_fill_bind_pat() {
 357         check_assist(
 358             fill_match_arms,
 359             r#"
 360 enum A { As, Bs, Cs(Option<i32>) }
 361 fn main() {
 362     match A::As$0 {
 363         A::As(_) => {}
 364         a @ A::Bs(_) => {}
 365     }
 366 }
 367 "#,
 368             r#"
 369 enum A { As, Bs, Cs(Option<i32>) }
 370 fn main() {
 371     match A::As {
 372         A::As(_) => {}
 373         a @ A::Bs(_) => {}
 374         A::Cs(${0:_}) => {}
 375     }
 376 }
 377 "#,
 378         );
 379     }
 380
 381     #[test]
 382     fn fill_match_arms_empty_body() {
 383         check_assist(
 384             fill_match_arms,
 385             r#"
 386 enum A { As, Bs, Cs(String), Ds(String, String), Es { x: usize, y: usize } }
 387
 388 fn main() {
 389     let a = A::As;
 390     match a$0 {}
 391 }
 392 "#,
 393             r#"
 394 enum A { As, Bs, Cs(String), Ds(String, String), Es { x: usize, y: usize } }
 395
 396 fn main() {
 397     let a = A::As;
 398     match a {
 399         $0A::As => {}
 400         A::Bs => {}
 401         A::Cs(_) => {}
 402         A::Ds(_, _) => {}
 403         A::Es { x, y } => {}
 404     }
 405 }
 406 "#,
 407         );
 408     }
 409
 410     #[test]
 411     fn fill_match_arms_tuple_of_enum() {
 412         check_assist(
 413             fill_match_arms,
 414             r#"
 415             enum A { One, Two }
 416             enum B { One, Two }
 417
 418             fn main() {
 419                 let a = A::One;
 420                 let b = B::One;
 421                 match (a$0, b) {}
 422             }
 423             "#,
 424             r#"
 425             enum A { One, Two }
 426             enum B { One, Two }
 427
 428             fn main() {
 429                 let a = A::One;
 430                 let b = B::One;
 431                 match (a, b) {
 432                     $0(A::One, B::One) => {}
 433                     (A::One, B::Two) => {}
 434                     (A::Two, B::One) => {}
 435                     (A::Two, B::Two) => {}
 436                 }
 437             }
 438             "#,
 439         );
 440     }
 441
 442     #[test]
 443     fn fill_match_arms_tuple_of_enum_ref() {
 444         check_assist(
 445             fill_match_arms,
 446             r#"
 447             enum A { One, Two }
 448             enum B { One, Two }
 449
 450             fn main() {
 451                 let a = A::One;
 452                 let b = B::One;
 453                 match (&a$0, &b) {}
 454             }
 455             "#,
 456             r#"
 457             enum A { One, Two }
 458             enum B { One, Two }
 459
 460             fn main() {
 461                 let a = A::One;
 462                 let b = B::One;
 463                 match (&a, &b) {
 464                     $0(A::One, B::One) => {}
 465                     (A::One, B::Two) => {}
 466                     (A::Two, B::One) => {}
 467                     (A::Two, B::Two) => {}
 468                 }
 469             }
 470             "#,
 471         );
 472     }
 473
 474     #[test]
 475     fn fill_match_arms_tuple_of_enum_partial() {
 476         check_assist_not_applicable(
 477             fill_match_arms,
 478             r#"
 479             enum A { One, Two }
 480             enum B { One, Two }
 481
 482             fn main() {
 483                 let a = A::One;
 484                 let b = B::One;
 485                 match (a$0, b) {
 486                     (A::Two, B::One) => {}
 487                 }
 488             }
 489             "#,
 490         );
 491     }
 492
 493     #[test]
 494     fn fill_match_arms_tuple_of_enum_not_applicable() {
 495         check_assist_not_applicable(
 496             fill_match_arms,
 497             r#"
 498             enum A { One, Two }
 499             enum B { One, Two }
 500
 501             fn main() {
 502                 let a = A::One;
 503                 let b = B::One;
 504                 match (a$0, b) {
 505                     (A::Two, B::One) => {}
 506                     (A::One, B::One) => {}
 507                     (A::One, B::Two) => {}
 508                     (A::Two, B::Two) => {}
 509                 }
 510             }
 511             "#,
 512         );
 513     }
 514
 515     #[test]
 516     fn fill_match_arms_single_element_tuple_of_enum() {
 517         // For now we don't hande the case of a single element tuple, but
 518         // we could handle this in the future if `make::tuple_pat` allowed
 519         // creating a tuple with a single pattern.
 520         check_assist_not_applicable(
 521             fill_match_arms,
 522             r#"
 523             enum A { One, Two }
 524
 525             fn main() {
 526                 let a = A::One;
 527                 match (a$0, ) {
 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 { As }
 540
 541             fn foo(a: &A) {
 542                 match a$0 {
 543                 }
 544             }
 545             "#,
 546             r#"
 547             enum A { As }
 548
 549             fn foo(a: &A) {
 550                 match a {
 551                     $0A::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$0 {
 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                     $0A::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$0 {}
 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                     $0_ => {}
 608                 }
 609             }
 610             "#,
 611             r#"
 612             enum E { X, Y }
 613
 614             fn main() {
 615                 match E::X {
 616                     $0E::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                     $0
 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                     $0X => {}
 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 { One, Two }
 658             fn foo(a: A) {
 659                 match a {
 660                     // foo bar baz$0
 661                     A::One => {}
 662                     // This is where the rest should be
 663                 }
 664             }
 665             "#,
 666             r#"
 667             enum A { One, Two }
 668             fn foo(a: A) {
 669                 match a {
 670                     // foo bar baz
 671                     A::One => {}
 672                     // This is where the rest should be
 673                     $0A::Two => {}
 674                 }
 675             }
 676             "#,
 677         );
 678     }
 679
 680     #[test]
 681     fn fill_match_arms_preserves_comments_empty() {
 682         check_assist(
 683             fill_match_arms,
 684             r#"
 685             enum A { One, Two }
 686             fn foo(a: A) {
 687                 match a {
 688                     // foo bar baz$0
 689                 }
 690             }
 691             "#,
 692             r#"
 693             enum A { One, Two }
 694             fn foo(a: A) {
 695                 match a {
 696                     // foo bar baz
 697                     $0A::One => {}
 698                     A::Two => {}
 699                 }
 700             }
 701             "#,
 702         );
 703     }
 704
 705     #[test]
 706     fn fill_match_arms_placeholder() {
 707         check_assist(
 708             fill_match_arms,
 709             r#"
 710             enum A { One, Two, }
 711             fn foo(a: A) {
 712                 match a$0 {
 713                     _ => (),
 714                 }
 715             }
 716             "#,
 717             r#"
 718             enum A { One, Two, }
 719             fn foo(a: A) {
 720                 match a {
 721                     $0A::One => {}
 722                     A::Two => {}
 723                 }
 724             }
 725             "#,
 726         );
 727     }
 728
 729     #[test]
 730     fn option_order() {
 731         cov_mark::check!(option_order);
 732         let before = r#"
 733 fn foo(opt: Option<i32>) {
 734     match opt$0 {
 735     }
 736 }
 737 "#;
 738         let before = &format!("//- /main.rs crate:main deps:core{}{}", before, FamousDefs::FIXTURE);
 739
 740         check_assist(
 741             fill_match_arms,
 742             before,
 743             r#"
 744 fn foo(opt: Option<i32>) {
 745     match opt {
 746         Some(${0:_}) => {}
 747         None => {}
 748     }
 749 }
 750 "#,
 751         );
 752     }
 753
 754     #[test]
 755     fn works_inside_macro_call() {
 756         check_assist(
 757             fill_match_arms,
 758             r#"
 759 macro_rules! m { ($expr:expr) => {$expr}}
 760 enum Test {
 761     A,
 762     B,
 763     C,
 764 }
 765
 766 fn foo(t: Test) {
 767     m!(match t$0 {});
 768 }"#,
 769             r#"macro_rules! m { ($expr:expr) => {$expr}}
 770 enum Test {
 771     A,
 772     B,
 773     C,
 774 }
 775
 776 fn foo(t: Test) {
 777     m!(match t {
 778     $0Test::A => {}
 779     Test::B => {}
 780     Test::C => {}
 781 });
 782 }"#,
 783         );
 784     }
 785 }