crates/ide_assists/src/handlers/fill_match_arms.rs

   1 use std::iter;
   2
   3 use either::Either;
   4 use hir::{Adt, HasSource, ModuleDef, Semantics};
   5 use ide_db::helpers::{mod_path_to_ast, FamousDefs};
   6 use ide_db::RootDatabase;
   7 use itertools::Itertools;
   8 use syntax::ast::{self, make, AstNode, MatchArm, NameOwner, Pat};
   9
  10 use crate::{
  11     utils::{does_pat_match_variant, render_snippet, Cursor},
  12     AssistContext, AssistId, AssistKind, Assists,
  13 };
  14
  15 // Assist: fill_match_arms
  16 //
  17 // Adds missing clauses to a `match` expression.
  18 //
  19 // ```
  20 // enum Action { Move { distance: u32 }, Stop }
  21 //
  22 // fn handle(action: Action) {
  23 //     match action {
  24 //         $0
  25 //     }
  26 // }
  27 // ```
  28 // ->
  29 // ```
  30 // enum Action { Move { distance: u32 }, Stop }
  31 //
  32 // fn handle(action: Action) {
  33 //     match action {
  34 //         $0Action::Move { distance } => {}
  35 //         Action::Stop => {}
  36 //     }
  37 // }
  38 // ```
  39 pub(crate) fn fill_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  40     let match_expr = ctx.find_node_at_offset_with_descend::<ast::MatchExpr>()?;
  41     let match_arm_list = match_expr.match_arm_list()?;
  42
  43     let expr = match_expr.expr()?;
  44
  45     let mut arms: Vec<MatchArm> = match_arm_list.arms().collect();
  46     if arms.len() == 1 {
  47         if let Some(Pat::WildcardPat(..)) = arms[0].pat() {
  48             arms.clear();
  49         }
  50     }
  51
  52     let top_lvl_pats: Vec<_> = arms
  53         .iter()
  54         .filter_map(ast::MatchArm::pat)
  55         .flat_map(|pat| match pat {
  56             // Special casee OrPat as separate top-level pats
  57             Pat::OrPat(or_pat) => Either::Left(or_pat.pats()),
  58             _ => Either::Right(iter::once(pat)),
  59         })
  60         .collect();
  61
  62     let module = ctx.sema.scope(expr.syntax()).module()?;
  63
  64     let missing_arms: Vec<MatchArm> = if let Some(enum_def) = resolve_enum_def(&ctx.sema, &expr) {
  65         let variants = enum_def.variants(ctx.db());
  66
  67         let mut variants = variants
  68             .into_iter()
  69             .filter_map(|variant| build_pat(ctx.db(), module, variant))
  70             .filter(|variant_pat| is_variant_missing(&top_lvl_pats, variant_pat))
  71             .map(|pat| make::match_arm(iter::once(pat), make::expr_empty_block()))
  72             .collect::<Vec<_>>();
  73         if Some(enum_def) == FamousDefs(&ctx.sema, Some(module.krate())).core_option_Option() {
  74             // Match `Some` variant first.
  75             cov_mark::hit!(option_order);
  76             variants.reverse()
  77         }
  78         variants
  79     } else if let Some(enum_defs) = resolve_tuple_of_enum_def(&ctx.sema, &expr) {
  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(&top_lvl_pats, 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_pats: &[Pat], var: &Pat) -> bool {
 138     !existing_pats.iter().any(|pat| match (pat, var) {
 139         (Pat::TuplePat(tpat), Pat::TuplePat(tvar)) => {
 140             // `does_pat_match_variant` gives false positives for tuple patterns
 141             // Fixme: this is still somewhat limited
 142             tpat.fields().zip(tvar.fields()).all(|(p, v)| does_pat_match_variant(&p, &v))
 143         }
 144         _ => does_pat_match_variant(pat, var),
 145     })
 146 }
 147
 148 fn resolve_enum_def(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<hir::Enum> {
 149     sema.type_of_expr(&expr)?.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
 150         Some(Adt::Enum(e)) => Some(e),
 151         _ => None,
 152     })
 153 }
 154
 155 fn resolve_tuple_of_enum_def(
 156     sema: &Semantics<RootDatabase>,
 157     expr: &ast::Expr,
 158 ) -> Option<Vec<hir::Enum>> {
 159     sema.type_of_expr(&expr)?
 160         .tuple_fields(sema.db)
 161         .iter()
 162         .map(|ty| {
 163             ty.autoderef(sema.db).find_map(|ty| match ty.as_adt() {
 164                 Some(Adt::Enum(e)) => Some(e),
 165                 // For now we only handle expansion for a tuple of enums. Here
 166                 // we map non-enum items to None and rely on `collect` to
 167                 // convert Vec<Option<hir::Enum>> into Option<Vec<hir::Enum>>.
 168                 _ => None,
 169             })
 170         })
 171         .collect()
 172 }
 173
 174 fn build_pat(db: &RootDatabase, module: hir::Module, var: hir::Variant) -> Option<ast::Pat> {
 175     let path = mod_path_to_ast(&module.find_use_path(db, ModuleDef::from(var))?);
 176
 177     // FIXME: use HIR for this; it doesn't currently expose struct vs. tuple vs. unit variants though
 178     let pat: ast::Pat = match var.source(db)?.value.kind() {
 179         ast::StructKind::Tuple(field_list) => {
 180             let pats = iter::repeat(make::wildcard_pat().into()).take(field_list.fields().count());
 181             make::tuple_struct_pat(path, pats).into()
 182         }
 183         ast::StructKind::Record(field_list) => {
 184             let pats = field_list.fields().map(|f| make::ident_pat(f.name().unwrap()).into());
 185             make::record_pat(path, pats).into()
 186         }
 187         ast::StructKind::Unit => make::path_pat(path),
 188     };
 189
 190     Some(pat)
 191 }
 192
 193 #[cfg(test)]
 194 mod tests {
 195     use ide_db::helpers::FamousDefs;
 196
 197     use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};
 198
 199     use super::fill_match_arms;
 200
 201     #[test]
 202     fn all_match_arms_provided() {
 203         check_assist_not_applicable(
 204             fill_match_arms,
 205             r#"
 206             enum A {
 207                 As,
 208                 Bs{x:i32, y:Option<i32>},
 209                 Cs(i32, Option<i32>),
 210             }
 211             fn main() {
 212                 match A::As$0 {
 213                     A::As,
 214                     A::Bs{x,y:Some(_)} => {}
 215                     A::Cs(_, Some(_)) => {}
 216                 }
 217             }
 218             "#,
 219         );
 220     }
 221
 222     #[test]
 223     fn tuple_of_non_enum() {
 224         // for now this case is not handled, although it potentially could be
 225         // in the future
 226         check_assist_not_applicable(
 227             fill_match_arms,
 228             r#"
 229             fn main() {
 230                 match (0, false)$0 {
 231                 }
 232             }
 233             "#,
 234         );
 235     }
 236
 237     #[test]
 238     fn partial_fill_record_tuple() {
 239         check_assist(
 240             fill_match_arms,
 241             r#"
 242             enum A {
 243                 As,
 244                 Bs { x: i32, y: Option<i32> },
 245                 Cs(i32, Option<i32>),
 246             }
 247             fn main() {
 248                 match A::As$0 {
 249                     A::Bs { x, y: Some(_) } => {}
 250                     A::Cs(_, Some(_)) => {}
 251                 }
 252             }
 253             "#,
 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                     $0A::As => {}
 265                 }
 266             }
 267             "#,
 268         );
 269     }
 270
 271     #[test]
 272     fn partial_fill_option() {
 273         check_assist(
 274             fill_match_arms,
 275             r#"
 276 enum Option<T> { Some(T), None }
 277 use Option::*;
 278
 279 fn main() {
 280     match None$0 {
 281         None => {}
 282     }
 283 }
 284             "#,
 285             r#"
 286 enum Option<T> { Some(T), None }
 287 use Option::*;
 288
 289 fn main() {
 290     match None {
 291         None => {}
 292         Some(${0:_}) => {}
 293     }
 294 }
 295             "#,
 296         );
 297     }
 298
 299     #[test]
 300     fn partial_fill_or_pat() {
 301         check_assist(
 302             fill_match_arms,
 303             r#"
 304 enum A { As, Bs, Cs(Option<i32>) }
 305 fn main() {
 306     match A::As$0 {
 307         A::Cs(_) | A::Bs => {}
 308     }
 309 }
 310 "#,
 311             r#"
 312 enum A { As, Bs, Cs(Option<i32>) }
 313 fn main() {
 314     match A::As {
 315         A::Cs(_) | A::Bs => {}
 316         $0A::As => {}
 317     }
 318 }
 319 "#,
 320         );
 321     }
 322
 323     #[test]
 324     fn partial_fill() {
 325         check_assist(
 326             fill_match_arms,
 327             r#"
 328 enum A { As, Bs, Cs, Ds(String), Es(B) }
 329 enum B { Xs, Ys }
 330 fn main() {
 331     match A::As$0 {
 332         A::Bs if 0 < 1 => {}
 333         A::Ds(_value) => { let x = 1; }
 334         A::Es(B::Xs) => (),
 335     }
 336 }
 337 "#,
 338             r#"
 339 enum A { As, Bs, Cs, Ds(String), Es(B) }
 340 enum B { Xs, Ys }
 341 fn main() {
 342     match A::As {
 343         A::Bs if 0 < 1 => {}
 344         A::Ds(_value) => { let x = 1; }
 345         A::Es(B::Xs) => (),
 346         $0A::As => {}
 347         A::Cs => {}
 348     }
 349 }
 350 "#,
 351         );
 352     }
 353
 354     #[test]
 355     fn partial_fill_bind_pat() {
 356         check_assist(
 357             fill_match_arms,
 358             r#"
 359 enum A { As, Bs, Cs(Option<i32>) }
 360 fn main() {
 361     match A::As$0 {
 362         A::As(_) => {}
 363         a @ A::Bs(_) => {}
 364     }
 365 }
 366 "#,
 367             r#"
 368 enum A { As, Bs, Cs(Option<i32>) }
 369 fn main() {
 370     match A::As {
 371         A::As(_) => {}
 372         a @ A::Bs(_) => {}
 373         A::Cs(${0:_}) => {}
 374     }
 375 }
 376 "#,
 377         );
 378     }
 379
 380     #[test]
 381     fn fill_match_arms_empty_body() {
 382         check_assist(
 383             fill_match_arms,
 384             r#"
 385 enum A { As, Bs, Cs(String), Ds(String, String), Es { x: usize, y: usize } }
 386
 387 fn main() {
 388     let a = A::As;
 389     match a$0 {}
 390 }
 391 "#,
 392             r#"
 393 enum A { As, Bs, Cs(String), Ds(String, String), Es { x: usize, y: usize } }
 394
 395 fn main() {
 396     let a = A::As;
 397     match a {
 398         $0A::As => {}
 399         A::Bs => {}
 400         A::Cs(_) => {}
 401         A::Ds(_, _) => {}
 402         A::Es { x, y } => {}
 403     }
 404 }
 405 "#,
 406         );
 407     }
 408
 409     #[test]
 410     fn fill_match_arms_tuple_of_enum() {
 411         check_assist(
 412             fill_match_arms,
 413             r#"
 414             enum A { One, Two }
 415             enum B { One, Two }
 416
 417             fn main() {
 418                 let a = A::One;
 419                 let b = B::One;
 420                 match (a$0, b) {}
 421             }
 422             "#,
 423             r#"
 424             enum A { One, Two }
 425             enum B { One, Two }
 426
 427             fn main() {
 428                 let a = A::One;
 429                 let b = B::One;
 430                 match (a, b) {
 431                     $0(A::One, B::One) => {}
 432                     (A::One, B::Two) => {}
 433                     (A::Two, B::One) => {}
 434                     (A::Two, B::Two) => {}
 435                 }
 436             }
 437             "#,
 438         );
 439     }
 440
 441     #[test]
 442     fn fill_match_arms_tuple_of_enum_ref() {
 443         check_assist(
 444             fill_match_arms,
 445             r#"
 446             enum A { One, Two }
 447             enum B { One, Two }
 448
 449             fn main() {
 450                 let a = A::One;
 451                 let b = B::One;
 452                 match (&a$0, &b) {}
 453             }
 454             "#,
 455             r#"
 456             enum A { One, Two }
 457             enum B { One, Two }
 458
 459             fn main() {
 460                 let a = A::One;
 461                 let b = B::One;
 462                 match (&a, &b) {
 463                     $0(A::One, B::One) => {}
 464                     (A::One, B::Two) => {}
 465                     (A::Two, B::One) => {}
 466                     (A::Two, B::Two) => {}
 467                 }
 468             }
 469             "#,
 470         );
 471     }
 472
 473     #[test]
 474     fn fill_match_arms_tuple_of_enum_partial() {
 475         check_assist(
 476             fill_match_arms,
 477             r#"
 478             enum A { One, Two }
 479             enum B { One, Two }
 480
 481             fn main() {
 482                 let a = A::One;
 483                 let b = B::One;
 484                 match (a$0, b) {
 485                     (A::Two, B::One) => {}
 486                 }
 487             }
 488             "#,
 489             r#"
 490             enum A { One, Two }
 491             enum B { One, Two }
 492
 493             fn main() {
 494                 let a = A::One;
 495                 let b = B::One;
 496                 match (a, b) {
 497                     (A::Two, B::One) => {}
 498                     $0(A::One, B::One) => {}
 499                     (A::One, B::Two) => {}
 500                     (A::Two, B::Two) => {}
 501                 }
 502             }
 503             "#,
 504         );
 505     }
 506
 507     // Fixme: This fails with extra useless match arms added.
 508     // To fix, it needs full fledged match exhaustiveness checking from
 509     // hir_ty::diagnostics::match_check
 510     // see https://github.com/rust-analyzer/rust-analyzer/issues/8493
 511     #[ignore]
 512     #[test]
 513     fn fill_match_arms_tuple_of_enum_partial_with_wildcards() {
 514         let ra_fixture = r#"
 515 fn main() {
 516     let a = Some(1);
 517     let b = Some(());
 518     match (a$0, b) {
 519         (Some(_), _) => {}
 520         (None, Some(_)) => {}
 521     }
 522 }
 523 "#;
 524         check_assist(
 525             fill_match_arms,
 526             &format!("//- /main.rs crate:main deps:core{}{}", ra_fixture, FamousDefs::FIXTURE),
 527             r#"
 528 fn main() {
 529     let a = Some(1);
 530     let b = Some(());
 531     match (a, b) {
 532         (Some(_), _) => {}
 533         (None, Some(_)) => {}
 534         $0(None, None) => {}
 535     }
 536 }
 537 "#,
 538         );
 539     }
 540
 541     #[test]
 542     fn fill_match_arms_tuple_of_enum_not_applicable() {
 543         check_assist_not_applicable(
 544             fill_match_arms,
 545             r#"
 546             enum A { One, Two }
 547             enum B { One, Two }
 548
 549             fn main() {
 550                 let a = A::One;
 551                 let b = B::One;
 552                 match (a$0, b) {
 553                     (A::Two, B::One) => {}
 554                     (A::One, B::One) => {}
 555                     (A::One, B::Two) => {}
 556                     (A::Two, B::Two) => {}
 557                 }
 558             }
 559             "#,
 560         );
 561     }
 562
 563     #[test]
 564     fn fill_match_arms_single_element_tuple_of_enum() {
 565         check_assist(
 566             fill_match_arms,
 567             r#"
 568             enum A { One, Two }
 569
 570             fn main() {
 571                 let a = A::One;
 572                 match (a$0, ) {
 573                 }
 574             }
 575             "#,
 576             r#"
 577             enum A { One, Two }
 578
 579             fn main() {
 580                 let a = A::One;
 581                 match (a, ) {
 582                     $0(A::One,) => {}
 583                     (A::Two,) => {}
 584                 }
 585             }
 586             "#,
 587         );
 588     }
 589
 590     #[test]
 591     fn test_fill_match_arm_refs() {
 592         check_assist(
 593             fill_match_arms,
 594             r#"
 595             enum A { As }
 596
 597             fn foo(a: &A) {
 598                 match a$0 {
 599                 }
 600             }
 601             "#,
 602             r#"
 603             enum A { As }
 604
 605             fn foo(a: &A) {
 606                 match a {
 607                     $0A::As => {}
 608                 }
 609             }
 610             "#,
 611         );
 612
 613         check_assist(
 614             fill_match_arms,
 615             r#"
 616             enum A {
 617                 Es { x: usize, y: usize }
 618             }
 619
 620             fn foo(a: &mut A) {
 621                 match a$0 {
 622                 }
 623             }
 624             "#,
 625             r#"
 626             enum A {
 627                 Es { x: usize, y: usize }
 628             }
 629
 630             fn foo(a: &mut A) {
 631                 match a {
 632                     $0A::Es { x, y } => {}
 633                 }
 634             }
 635             "#,
 636         );
 637     }
 638
 639     #[test]
 640     fn fill_match_arms_target() {
 641         check_assist_target(
 642             fill_match_arms,
 643             r#"
 644             enum E { X, Y }
 645
 646             fn main() {
 647                 match E::X$0 {}
 648             }
 649             "#,
 650             "match E::X {}",
 651         );
 652     }
 653
 654     #[test]
 655     fn fill_match_arms_trivial_arm() {
 656         check_assist(
 657             fill_match_arms,
 658             r#"
 659             enum E { X, Y }
 660
 661             fn main() {
 662                 match E::X {
 663                     $0_ => {}
 664                 }
 665             }
 666             "#,
 667             r#"
 668             enum E { X, Y }
 669
 670             fn main() {
 671                 match E::X {
 672                     $0E::X => {}
 673                     E::Y => {}
 674                 }
 675             }
 676             "#,
 677         );
 678     }
 679
 680     #[test]
 681     fn fill_match_arms_qualifies_path() {
 682         check_assist(
 683             fill_match_arms,
 684             r#"
 685             mod foo { pub enum E { X, Y } }
 686             use foo::E::X;
 687
 688             fn main() {
 689                 match X {
 690                     $0
 691                 }
 692             }
 693             "#,
 694             r#"
 695             mod foo { pub enum E { X, Y } }
 696             use foo::E::X;
 697
 698             fn main() {
 699                 match X {
 700                     $0X => {}
 701                     foo::E::Y => {}
 702                 }
 703             }
 704             "#,
 705         );
 706     }
 707
 708     #[test]
 709     fn fill_match_arms_preserves_comments() {
 710         check_assist(
 711             fill_match_arms,
 712             r#"
 713             enum A { One, Two }
 714             fn foo(a: A) {
 715                 match a {
 716                     // foo bar baz$0
 717                     A::One => {}
 718                     // This is where the rest should be
 719                 }
 720             }
 721             "#,
 722             r#"
 723             enum A { One, Two }
 724             fn foo(a: A) {
 725                 match a {
 726                     // foo bar baz
 727                     A::One => {}
 728                     // This is where the rest should be
 729                     $0A::Two => {}
 730                 }
 731             }
 732             "#,
 733         );
 734     }
 735
 736     #[test]
 737     fn fill_match_arms_preserves_comments_empty() {
 738         check_assist(
 739             fill_match_arms,
 740             r#"
 741             enum A { One, Two }
 742             fn foo(a: A) {
 743                 match a {
 744                     // foo bar baz$0
 745                 }
 746             }
 747             "#,
 748             r#"
 749             enum A { One, Two }
 750             fn foo(a: A) {
 751                 match a {
 752                     // foo bar baz
 753                     $0A::One => {}
 754                     A::Two => {}
 755                 }
 756             }
 757             "#,
 758         );
 759     }
 760
 761     #[test]
 762     fn fill_match_arms_placeholder() {
 763         check_assist(
 764             fill_match_arms,
 765             r#"
 766             enum A { One, Two, }
 767             fn foo(a: A) {
 768                 match a$0 {
 769                     _ => (),
 770                 }
 771             }
 772             "#,
 773             r#"
 774             enum A { One, Two, }
 775             fn foo(a: A) {
 776                 match a {
 777                     $0A::One => {}
 778                     A::Two => {}
 779                 }
 780             }
 781             "#,
 782         );
 783     }
 784
 785     #[test]
 786     fn option_order() {
 787         cov_mark::check!(option_order);
 788         let before = r#"
 789 fn foo(opt: Option<i32>) {
 790     match opt$0 {
 791     }
 792 }
 793 "#;
 794         let before = &format!("//- /main.rs crate:main deps:core{}{}", before, FamousDefs::FIXTURE);
 795
 796         check_assist(
 797             fill_match_arms,
 798             before,
 799             r#"
 800 fn foo(opt: Option<i32>) {
 801     match opt {
 802         Some(${0:_}) => {}
 803         None => {}
 804     }
 805 }
 806 "#,
 807         );
 808     }
 809
 810     #[test]
 811     fn works_inside_macro_call() {
 812         check_assist(
 813             fill_match_arms,
 814             r#"
 815 macro_rules! m { ($expr:expr) => {$expr}}
 816 enum Test {
 817     A,
 818     B,
 819     C,
 820 }
 821
 822 fn foo(t: Test) {
 823     m!(match t$0 {});
 824 }"#,
 825             r#"macro_rules! m { ($expr:expr) => {$expr}}
 826 enum Test {
 827     A,
 828     B,
 829     C,
 830 }
 831
 832 fn foo(t: Test) {
 833     m!(match t {
 834     $0Test::A => {}
 835     Test::B => {}
 836     Test::C => {}
 837 });
 838 }"#,
 839         );
 840     }
 841 }