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 = crate::ast_transform::path_to_ast(module.find_use_path(db, ModuleDef::from(var))?);
 158
 159     // FIXME: use HIR for this; it doesn't currently expose struct vs. tuple vs. unit variants though
 160     let pat: ast::Pat = match var.source(db).value.kind() {
 161         ast::StructKind::Tuple(field_list) => {
 162             let pats =
 163                 iter::repeat(make::placeholder_pat().into()).take(field_list.fields().count());
 164             make::tuple_struct_pat(path, pats).into()
 165         }
 166         ast::StructKind::Record(field_list) => {
 167             let pats = field_list.fields().map(|f| make::bind_pat(f.name().unwrap()).into());
 168             make::record_pat(path, pats).into()
 169         }
 170         ast::StructKind::Unit => make::path_pat(path),
 171     };
 172
 173     Some(pat)
 174 }
 175
 176 #[cfg(test)]
 177 mod tests {
 178     use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_target};
 179
 180     use super::fill_match_arms;
 181
 182     #[test]
 183     fn all_match_arms_provided() {
 184         check_assist_not_applicable(
 185             fill_match_arms,
 186             r#"
 187             enum A {
 188                 As,
 189                 Bs{x:i32, y:Option<i32>},
 190                 Cs(i32, Option<i32>),
 191             }
 192             fn main() {
 193                 match A::As<|> {
 194                     A::As,
 195                     A::Bs{x,y:Some(_)} => (),
 196                     A::Cs(_, Some(_)) => (),
 197                 }
 198             }
 199             "#,
 200         );
 201     }
 202
 203     #[test]
 204     fn tuple_of_non_enum() {
 205         // for now this case is not handled, although it potentially could be
 206         // in the future
 207         check_assist_not_applicable(
 208             fill_match_arms,
 209             r#"
 210             fn main() {
 211                 match (0, false)<|> {
 212                 }
 213             }
 214             "#,
 215         );
 216     }
 217
 218     #[test]
 219     fn partial_fill_record_tuple() {
 220         check_assist(
 221             fill_match_arms,
 222             r#"
 223             enum A {
 224                 As,
 225                 Bs{x:i32, y:Option<i32>},
 226                 Cs(i32, Option<i32>),
 227             }
 228             fn main() {
 229                 match A::As<|> {
 230                     A::Bs{x,y:Some(_)} => (),
 231                     A::Cs(_, Some(_)) => (),
 232                 }
 233             }
 234             "#,
 235             r#"
 236             enum A {
 237                 As,
 238                 Bs{x:i32, y:Option<i32>},
 239                 Cs(i32, Option<i32>),
 240             }
 241             fn main() {
 242                 match <|>A::As {
 243                     A::Bs{x,y:Some(_)} => (),
 244                     A::Cs(_, Some(_)) => (),
 245                     A::As => (),
 246                 }
 247             }
 248             "#,
 249         );
 250     }
 251
 252     #[test]
 253     fn partial_fill_or_pat() {
 254         check_assist(
 255             fill_match_arms,
 256             r#"
 257             enum A {
 258                 As,
 259                 Bs,
 260                 Cs(Option<i32>),
 261             }
 262             fn main() {
 263                 match A::As<|> {
 264                     A::Cs(_) | A::Bs => (),
 265                 }
 266             }
 267             "#,
 268             r#"
 269             enum A {
 270                 As,
 271                 Bs,
 272                 Cs(Option<i32>),
 273             }
 274             fn main() {
 275                 match <|>A::As {
 276                     A::Cs(_) | A::Bs => (),
 277                     A::As => (),
 278                 }
 279             }
 280             "#,
 281         );
 282     }
 283
 284     #[test]
 285     fn partial_fill() {
 286         check_assist(
 287             fill_match_arms,
 288             r#"
 289             enum A {
 290                 As,
 291                 Bs,
 292                 Cs,
 293                 Ds(String),
 294                 Es(B),
 295             }
 296             enum B {
 297                 Xs,
 298                 Ys,
 299             }
 300             fn main() {
 301                 match A::As<|> {
 302                     A::Bs if 0 < 1 => (),
 303                     A::Ds(_value) => (),
 304                     A::Es(B::Xs) => (),
 305                 }
 306             }
 307             "#,
 308             r#"
 309             enum A {
 310                 As,
 311                 Bs,
 312                 Cs,
 313                 Ds(String),
 314                 Es(B),
 315             }
 316             enum B {
 317                 Xs,
 318                 Ys,
 319             }
 320             fn main() {
 321                 match <|>A::As {
 322                     A::Bs if 0 < 1 => (),
 323                     A::Ds(_value) => (),
 324                     A::Es(B::Xs) => (),
 325                     A::As => (),
 326                     A::Cs => (),
 327                 }
 328             }
 329             "#,
 330         );
 331     }
 332
 333     #[test]
 334     fn fill_match_arms_empty_body() {
 335         check_assist(
 336             fill_match_arms,
 337             r#"
 338             enum A {
 339                 As,
 340                 Bs,
 341                 Cs(String),
 342                 Ds(String, String),
 343                 Es{ x: usize, y: usize }
 344             }
 345
 346             fn main() {
 347                 let a = A::As;
 348                 match a<|> {}
 349             }
 350             "#,
 351             r#"
 352             enum A {
 353                 As,
 354                 Bs,
 355                 Cs(String),
 356                 Ds(String, String),
 357                 Es{ x: usize, y: usize }
 358             }
 359
 360             fn main() {
 361                 let a = A::As;
 362                 match <|>a {
 363                     A::As => (),
 364                     A::Bs => (),
 365                     A::Cs(_) => (),
 366                     A::Ds(_, _) => (),
 367                     A::Es { x, y } => (),
 368                 }
 369             }
 370             "#,
 371         );
 372     }
 373
 374     #[test]
 375     fn fill_match_arms_tuple_of_enum() {
 376         check_assist(
 377             fill_match_arms,
 378             r#"
 379             enum A {
 380                 One,
 381                 Two,
 382             }
 383             enum B {
 384                 One,
 385                 Two,
 386             }
 387
 388             fn main() {
 389                 let a = A::One;
 390                 let b = B::One;
 391                 match (a<|>, b) {}
 392             }
 393             "#,
 394             r#"
 395             enum A {
 396                 One,
 397                 Two,
 398             }
 399             enum B {
 400                 One,
 401                 Two,
 402             }
 403
 404             fn main() {
 405                 let a = A::One;
 406                 let b = B::One;
 407                 match <|>(a, b) {
 408                     (A::One, B::One) => (),
 409                     (A::One, B::Two) => (),
 410                     (A::Two, B::One) => (),
 411                     (A::Two, B::Two) => (),
 412                 }
 413             }
 414             "#,
 415         );
 416     }
 417
 418     #[test]
 419     fn fill_match_arms_tuple_of_enum_ref() {
 420         check_assist(
 421             fill_match_arms,
 422             r#"
 423             enum A {
 424                 One,
 425                 Two,
 426             }
 427             enum B {
 428                 One,
 429                 Two,
 430             }
 431
 432             fn main() {
 433                 let a = A::One;
 434                 let b = B::One;
 435                 match (&a<|>, &b) {}
 436             }
 437             "#,
 438             r#"
 439             enum A {
 440                 One,
 441                 Two,
 442             }
 443             enum B {
 444                 One,
 445                 Two,
 446             }
 447
 448             fn main() {
 449                 let a = A::One;
 450                 let b = B::One;
 451                 match <|>(&a, &b) {
 452                     (A::One, B::One) => (),
 453                     (A::One, B::Two) => (),
 454                     (A::Two, B::One) => (),
 455                     (A::Two, B::Two) => (),
 456                 }
 457             }
 458             "#,
 459         );
 460     }
 461
 462     #[test]
 463     fn fill_match_arms_tuple_of_enum_partial() {
 464         check_assist_not_applicable(
 465             fill_match_arms,
 466             r#"
 467             enum A {
 468                 One,
 469                 Two,
 470             }
 471             enum B {
 472                 One,
 473                 Two,
 474             }
 475
 476             fn main() {
 477                 let a = A::One;
 478                 let b = B::One;
 479                 match (a<|>, b) {
 480                     (A::Two, B::One) => (),
 481                 }
 482             }
 483             "#,
 484         );
 485     }
 486
 487     #[test]
 488     fn fill_match_arms_tuple_of_enum_not_applicable() {
 489         check_assist_not_applicable(
 490             fill_match_arms,
 491             r#"
 492             enum A {
 493                 One,
 494                 Two,
 495             }
 496             enum B {
 497                 One,
 498                 Two,
 499             }
 500
 501             fn main() {
 502                 let a = A::One;
 503                 let b = B::One;
 504                 match (a<|>, 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 {
 524                 One,
 525                 Two,
 526             }
 527
 528             fn main() {
 529                 let a = A::One;
 530                 match (a<|>, ) {
 531                 }
 532             }
 533             "#,
 534         );
 535     }
 536
 537     #[test]
 538     fn test_fill_match_arm_refs() {
 539         check_assist(
 540             fill_match_arms,
 541             r#"
 542             enum A {
 543                 As,
 544             }
 545
 546             fn foo(a: &A) {
 547                 match a<|> {
 548                 }
 549             }
 550             "#,
 551             r#"
 552             enum A {
 553                 As,
 554             }
 555
 556             fn foo(a: &A) {
 557                 match <|>a {
 558                     A::As => (),
 559                 }
 560             }
 561             "#,
 562         );
 563
 564         check_assist(
 565             fill_match_arms,
 566             r#"
 567             enum A {
 568                 Es{ x: usize, y: usize }
 569             }
 570
 571             fn foo(a: &mut A) {
 572                 match a<|> {
 573                 }
 574             }
 575             "#,
 576             r#"
 577             enum A {
 578                 Es{ x: usize, y: usize }
 579             }
 580
 581             fn foo(a: &mut A) {
 582                 match <|>a {
 583                     A::Es { x, y } => (),
 584                 }
 585             }
 586             "#,
 587         );
 588     }
 589
 590     #[test]
 591     fn fill_match_arms_target() {
 592         check_assist_target(
 593             fill_match_arms,
 594             r#"
 595             enum E { X, Y }
 596
 597             fn main() {
 598                 match E::X<|> {}
 599             }
 600             "#,
 601             "match E::X {}",
 602         );
 603     }
 604
 605     #[test]
 606     fn fill_match_arms_trivial_arm() {
 607         check_assist(
 608             fill_match_arms,
 609             r#"
 610             enum E { X, Y }
 611
 612             fn main() {
 613                 match E::X {
 614                     <|>_ => {},
 615                 }
 616             }
 617             "#,
 618             r#"
 619             enum E { X, Y }
 620
 621             fn main() {
 622                 match <|>E::X {
 623                     E::X => (),
 624                     E::Y => (),
 625                 }
 626             }
 627             "#,
 628         );
 629     }
 630
 631     #[test]
 632     fn fill_match_arms_qualifies_path() {
 633         check_assist(
 634             fill_match_arms,
 635             r#"
 636             mod foo { pub enum E { X, Y } }
 637             use foo::E::X;
 638
 639             fn main() {
 640                 match X {
 641                     <|>
 642                 }
 643             }
 644             "#,
 645             r#"
 646             mod foo { pub enum E { X, Y } }
 647             use foo::E::X;
 648
 649             fn main() {
 650                 match <|>X {
 651                     X => (),
 652                     foo::E::Y => (),
 653                 }
 654             }
 655             "#,
 656         );
 657     }
 658 }