crates/ra_assists/src/handlers/fill_match_arms.rs

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