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     #[test]
 508     fn fill_match_arms_tuple_of_enum_not_applicable() {
 509         check_assist_not_applicable(
 510             fill_match_arms,
 511             r#"
 512             enum A { One, Two }
 513             enum B { One, Two }
 514
 515             fn main() {
 516                 let a = A::One;
 517                 let b = B::One;
 518                 match (a$0, b) {
 519                     (A::Two, B::One) => {}
 520                     (A::One, B::One) => {}
 521                     (A::One, B::Two) => {}
 522                     (A::Two, B::Two) => {}
 523                 }
 524             }
 525             "#,
 526         );
 527     }
 528
 529     #[test]
 530     fn fill_match_arms_single_element_tuple_of_enum() {
 531         check_assist(
 532             fill_match_arms,
 533             r#"
 534             enum A { One, Two }
 535
 536             fn main() {
 537                 let a = A::One;
 538                 match (a$0, ) {
 539                 }
 540             }
 541             "#,
 542             r#"
 543             enum A { One, Two }
 544
 545             fn main() {
 546                 let a = A::One;
 547                 match (a, ) {
 548                     $0(A::One,) => {}
 549                     (A::Two,) => {}
 550                 }
 551             }
 552             "#,
 553         );
 554     }
 555
 556     #[test]
 557     fn test_fill_match_arm_refs() {
 558         check_assist(
 559             fill_match_arms,
 560             r#"
 561             enum A { As }
 562
 563             fn foo(a: &A) {
 564                 match a$0 {
 565                 }
 566             }
 567             "#,
 568             r#"
 569             enum A { As }
 570
 571             fn foo(a: &A) {
 572                 match a {
 573                     $0A::As => {}
 574                 }
 575             }
 576             "#,
 577         );
 578
 579         check_assist(
 580             fill_match_arms,
 581             r#"
 582             enum A {
 583                 Es { x: usize, y: usize }
 584             }
 585
 586             fn foo(a: &mut A) {
 587                 match a$0 {
 588                 }
 589             }
 590             "#,
 591             r#"
 592             enum A {
 593                 Es { x: usize, y: usize }
 594             }
 595
 596             fn foo(a: &mut A) {
 597                 match a {
 598                     $0A::Es { x, y } => {}
 599                 }
 600             }
 601             "#,
 602         );
 603     }
 604
 605     #[test]
 606     fn fill_match_arms_target() {
 607         check_assist_target(
 608             fill_match_arms,
 609             r#"
 610             enum E { X, Y }
 611
 612             fn main() {
 613                 match E::X$0 {}
 614             }
 615             "#,
 616             "match E::X {}",
 617         );
 618     }
 619
 620     #[test]
 621     fn fill_match_arms_trivial_arm() {
 622         check_assist(
 623             fill_match_arms,
 624             r#"
 625             enum E { X, Y }
 626
 627             fn main() {
 628                 match E::X {
 629                     $0_ => {}
 630                 }
 631             }
 632             "#,
 633             r#"
 634             enum E { X, Y }
 635
 636             fn main() {
 637                 match E::X {
 638                     $0E::X => {}
 639                     E::Y => {}
 640                 }
 641             }
 642             "#,
 643         );
 644     }
 645
 646     #[test]
 647     fn fill_match_arms_qualifies_path() {
 648         check_assist(
 649             fill_match_arms,
 650             r#"
 651             mod foo { pub enum E { X, Y } }
 652             use foo::E::X;
 653
 654             fn main() {
 655                 match X {
 656                     $0
 657                 }
 658             }
 659             "#,
 660             r#"
 661             mod foo { pub enum E { X, Y } }
 662             use foo::E::X;
 663
 664             fn main() {
 665                 match X {
 666                     $0X => {}
 667                     foo::E::Y => {}
 668                 }
 669             }
 670             "#,
 671         );
 672     }
 673
 674     #[test]
 675     fn fill_match_arms_preserves_comments() {
 676         check_assist(
 677             fill_match_arms,
 678             r#"
 679             enum A { One, Two }
 680             fn foo(a: A) {
 681                 match a {
 682                     // foo bar baz$0
 683                     A::One => {}
 684                     // This is where the rest should be
 685                 }
 686             }
 687             "#,
 688             r#"
 689             enum A { One, Two }
 690             fn foo(a: A) {
 691                 match a {
 692                     // foo bar baz
 693                     A::One => {}
 694                     // This is where the rest should be
 695                     $0A::Two => {}
 696                 }
 697             }
 698             "#,
 699         );
 700     }
 701
 702     #[test]
 703     fn fill_match_arms_preserves_comments_empty() {
 704         check_assist(
 705             fill_match_arms,
 706             r#"
 707             enum A { One, Two }
 708             fn foo(a: A) {
 709                 match a {
 710                     // foo bar baz$0
 711                 }
 712             }
 713             "#,
 714             r#"
 715             enum A { One, Two }
 716             fn foo(a: A) {
 717                 match a {
 718                     // foo bar baz
 719                     $0A::One => {}
 720                     A::Two => {}
 721                 }
 722             }
 723             "#,
 724         );
 725     }
 726
 727     #[test]
 728     fn fill_match_arms_placeholder() {
 729         check_assist(
 730             fill_match_arms,
 731             r#"
 732             enum A { One, Two, }
 733             fn foo(a: A) {
 734                 match a$0 {
 735                     _ => (),
 736                 }
 737             }
 738             "#,
 739             r#"
 740             enum A { One, Two, }
 741             fn foo(a: A) {
 742                 match a {
 743                     $0A::One => {}
 744                     A::Two => {}
 745                 }
 746             }
 747             "#,
 748         );
 749     }
 750
 751     #[test]
 752     fn option_order() {
 753         cov_mark::check!(option_order);
 754         let before = r#"
 755 fn foo(opt: Option<i32>) {
 756     match opt$0 {
 757     }
 758 }
 759 "#;
 760         let before = &format!("//- /main.rs crate:main deps:core{}{}", before, FamousDefs::FIXTURE);
 761
 762         check_assist(
 763             fill_match_arms,
 764             before,
 765             r#"
 766 fn foo(opt: Option<i32>) {
 767     match opt {
 768         Some(${0:_}) => {}
 769         None => {}
 770     }
 771 }
 772 "#,
 773         );
 774     }
 775
 776     #[test]
 777     fn works_inside_macro_call() {
 778         check_assist(
 779             fill_match_arms,
 780             r#"
 781 macro_rules! m { ($expr:expr) => {$expr}}
 782 enum Test {
 783     A,
 784     B,
 785     C,
 786 }
 787
 788 fn foo(t: Test) {
 789     m!(match t$0 {});
 790 }"#,
 791             r#"macro_rules! m { ($expr:expr) => {$expr}}
 792 enum Test {
 793     A,
 794     B,
 795     C,
 796 }
 797
 798 fn foo(t: Test) {
 799     m!(match t {
 800     $0Test::A => {}
 801     Test::B => {}
 802     Test::C => {}
 803 });
 804 }"#,
 805         );
 806     }
 807 }