crates/ide_assists/src/handlers/extract_struct_from_enum_variant.rs

   1 use std::iter;
   2
   3 use either::Either;
   4 use hir::{Module, ModuleDef, Name, Variant};
   5 use ide_db::{
   6     defs::Definition,
   7     helpers::{
   8         insert_use::{insert_use, ImportScope, InsertUseConfig},
   9         mod_path_to_ast,
  10     },
  11     search::FileReference,
  12     RootDatabase,
  13 };
  14 use itertools::Itertools;
  15 use rustc_hash::FxHashSet;
  16 use syntax::{
  17     ast::{
  18         self, make, AstNode, AttrsOwner, GenericParamsOwner, NameOwner, TypeBoundsOwner,
  19         VisibilityOwner,
  20     },
  21     match_ast,
  22     ted::{self, Position},
  23     SyntaxNode, T,
  24 };
  25
  26 use crate::{assist_context::AssistBuilder, AssistContext, AssistId, AssistKind, Assists};
  27
  28 // Assist: extract_struct_from_enum_variant
  29 //
  30 // Extracts a struct from enum variant.
  31 //
  32 // ```
  33 // enum A { $0One(u32, u32) }
  34 // ```
  35 // ->
  36 // ```
  37 // struct One(pub u32, pub u32);
  38 //
  39 // enum A { One(One) }
  40 // ```
  41 pub(crate) fn extract_struct_from_enum_variant(
  42     acc: &mut Assists,
  43     ctx: &AssistContext,
  44 ) -> Option<()> {
  45     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
  46     let field_list = extract_field_list_if_applicable(&variant)?;
  47
  48     let variant_name = variant.name()?;
  49     let variant_hir = ctx.sema.to_def(&variant)?;
  50     if existing_definition(ctx.db(), &variant_name, &variant_hir) {
  51         return None;
  52     }
  53
  54     let enum_ast = variant.parent_enum();
  55     let enum_hir = ctx.sema.to_def(&enum_ast)?;
  56     let target = variant.syntax().text_range();
  57     acc.add(
  58         AssistId("extract_struct_from_enum_variant", AssistKind::RefactorRewrite),
  59         "Extract struct from enum variant",
  60         target,
  61         |builder| {
  62             let variant_hir_name = variant_hir.name(ctx.db());
  63             let enum_module_def = ModuleDef::from(enum_hir);
  64             let usages =
  65                 Definition::ModuleDef(ModuleDef::Variant(variant_hir)).usages(&ctx.sema).all();
  66
  67             let mut visited_modules_set = FxHashSet::default();
  68             let current_module = enum_hir.module(ctx.db());
  69             visited_modules_set.insert(current_module);
  70             // record file references of the file the def resides in, we only want to swap to the edited file in the builder once
  71             let mut def_file_references = None;
  72             for (file_id, references) in usages {
  73                 if file_id == ctx.frange.file_id {
  74                     def_file_references = Some(references);
  75                     continue;
  76                 }
  77                 builder.edit_file(file_id);
  78                 let processed = process_references(
  79                     ctx,
  80                     builder,
  81                     &mut visited_modules_set,
  82                     &enum_module_def,
  83                     &variant_hir_name,
  84                     references,
  85                 );
  86                 processed.into_iter().for_each(|(path, node, import)| {
  87                     apply_references(ctx.config.insert_use, path, node, import)
  88                 });
  89             }
  90             builder.edit_file(ctx.frange.file_id);
  91             let variant = builder.make_mut(variant.clone());
  92             if let Some(references) = def_file_references {
  93                 let processed = process_references(
  94                     ctx,
  95                     builder,
  96                     &mut visited_modules_set,
  97                     &enum_module_def,
  98                     &variant_hir_name,
  99                     references,
 100                 );
 101                 processed.into_iter().for_each(|(path, node, import)| {
 102                     apply_references(ctx.config.insert_use, path, node, import)
 103                 });
 104             }
 105
 106             let def = create_struct_def(variant_name.clone(), &field_list, &enum_ast);
 107             let start_offset = &variant.parent_enum().syntax().clone();
 108             ted::insert_raw(ted::Position::before(start_offset), def.syntax());
 109             ted::insert_raw(ted::Position::before(start_offset), &make::tokens::blank_line());
 110
 111             update_variant(&variant, enum_ast.generic_param_list());
 112         },
 113     )
 114 }
 115
 116 fn extract_field_list_if_applicable(
 117     variant: &ast::Variant,
 118 ) -> Option<Either<ast::RecordFieldList, ast::TupleFieldList>> {
 119     match variant.kind() {
 120         ast::StructKind::Record(field_list) if field_list.fields().next().is_some() => {
 121             Some(Either::Left(field_list))
 122         }
 123         ast::StructKind::Tuple(field_list) if field_list.fields().count() > 1 => {
 124             Some(Either::Right(field_list))
 125         }
 126         _ => None,
 127     }
 128 }
 129
 130 fn existing_definition(db: &RootDatabase, variant_name: &ast::Name, variant: &Variant) -> bool {
 131     variant
 132         .parent_enum(db)
 133         .module(db)
 134         .scope(db, None)
 135         .into_iter()
 136         .filter(|(_, def)| match def {
 137             // only check type-namespace
 138             hir::ScopeDef::ModuleDef(def) => matches!(
 139                 def,
 140                 ModuleDef::Module(_)
 141                     | ModuleDef::Adt(_)
 142                     | ModuleDef::Variant(_)
 143                     | ModuleDef::Trait(_)
 144                     | ModuleDef::TypeAlias(_)
 145                     | ModuleDef::BuiltinType(_)
 146             ),
 147             _ => false,
 148         })
 149         .any(|(name, _)| name.to_string() == variant_name.to_string())
 150 }
 151
 152 fn create_struct_def(
 153     variant_name: ast::Name,
 154     field_list: &Either<ast::RecordFieldList, ast::TupleFieldList>,
 155     enum_: &ast::Enum,
 156 ) -> ast::Struct {
 157     let pub_vis = make::visibility_pub();
 158
 159     let insert_pub = |node: &'_ SyntaxNode| {
 160         let pub_vis = pub_vis.clone_for_update();
 161         ted::insert(ted::Position::before(node), pub_vis.syntax());
 162     };
 163
 164     // for fields without any existing visibility, use pub visibility
 165     let field_list = match field_list {
 166         Either::Left(field_list) => {
 167             let field_list = field_list.clone_for_update();
 168
 169             field_list
 170                 .fields()
 171                 .filter(|field| field.visibility().is_none())
 172                 .filter_map(|field| field.name())
 173                 .for_each(|it| insert_pub(it.syntax()));
 174
 175             field_list.into()
 176         }
 177         Either::Right(field_list) => {
 178             let field_list = field_list.clone_for_update();
 179
 180             field_list
 181                 .fields()
 182                 .filter(|field| field.visibility().is_none())
 183                 .filter_map(|field| field.ty())
 184                 .for_each(|it| insert_pub(it.syntax()));
 185
 186             field_list.into()
 187         }
 188     };
 189
 190     // FIXME: This uses all the generic params of the enum, but the variant might not use all of them.
 191     let strukt =
 192         make::struct_(enum_.visibility(), variant_name, enum_.generic_param_list(), field_list)
 193             .clone_for_update();
 194
 195     // copy attributes
 196     ted::insert_all(
 197         Position::first_child_of(strukt.syntax()),
 198         enum_.attrs().map(|it| it.syntax().clone_for_update().into()).collect(),
 199     );
 200     strukt
 201 }
 202
 203 fn update_variant(variant: &ast::Variant, generic: Option<ast::GenericParamList>) -> Option<()> {
 204     let name = variant.name()?;
 205     let ty = match generic {
 206         // FIXME: This uses all the generic params of the enum, but the variant might not use all of them.
 207         Some(gpl) => {
 208             let gpl = gpl.clone_for_update();
 209             gpl.generic_params().for_each(|gp| {
 210                 match gp {
 211                     ast::GenericParam::LifetimeParam(it) => it.type_bound_list(),
 212                     ast::GenericParam::TypeParam(it) => it.type_bound_list(),
 213                     ast::GenericParam::ConstParam(_) => return,
 214                 }
 215                 .map(|it| it.remove());
 216             });
 217             make::ty(&format!("{}<{}>", name.text(), gpl.generic_params().join(", ")))
 218         }
 219         None => make::ty(&name.text()),
 220     };
 221     let tuple_field = make::tuple_field(None, ty);
 222     let replacement = make::variant(
 223         name,
 224         Some(ast::FieldList::TupleFieldList(make::tuple_field_list(iter::once(tuple_field)))),
 225     )
 226     .clone_for_update();
 227     ted::replace(variant.syntax(), replacement.syntax());
 228     Some(())
 229 }
 230
 231 fn apply_references(
 232     insert_use_cfg: InsertUseConfig,
 233     segment: ast::PathSegment,
 234     node: SyntaxNode,
 235     import: Option<(ImportScope, hir::ModPath)>,
 236 ) {
 237     if let Some((scope, path)) = import {
 238         insert_use(&scope, mod_path_to_ast(&path), insert_use_cfg);
 239     }
 240     // deep clone to prevent cycle
 241     let path = make::path_from_segments(iter::once(segment.clone_subtree()), false);
 242     ted::insert_raw(ted::Position::before(segment.syntax()), path.clone_for_update().syntax());
 243     ted::insert_raw(ted::Position::before(segment.syntax()), make::token(T!['(']));
 244     ted::insert_raw(ted::Position::after(&node), make::token(T![')']));
 245 }
 246
 247 fn process_references(
 248     ctx: &AssistContext,
 249     builder: &mut AssistBuilder,
 250     visited_modules: &mut FxHashSet<Module>,
 251     enum_module_def: &ModuleDef,
 252     variant_hir_name: &Name,
 253     refs: Vec<FileReference>,
 254 ) -> Vec<(ast::PathSegment, SyntaxNode, Option<(ImportScope, hir::ModPath)>)> {
 255     // we have to recollect here eagerly as we are about to edit the tree we need to calculate the changes
 256     // and corresponding nodes up front
 257     refs.into_iter()
 258         .flat_map(|reference| {
 259             let (segment, scope_node, module) = reference_to_node(&ctx.sema, reference)?;
 260             let segment = builder.make_mut(segment);
 261             let scope_node = builder.make_syntax_mut(scope_node);
 262             if !visited_modules.contains(&module) {
 263                 let mod_path = module.find_use_path_prefixed(
 264                     ctx.sema.db,
 265                     *enum_module_def,
 266                     ctx.config.insert_use.prefix_kind,
 267                 );
 268                 if let Some(mut mod_path) = mod_path {
 269                     mod_path.pop_segment();
 270                     mod_path.push_segment(variant_hir_name.clone());
 271                     let scope = ImportScope::find_insert_use_container(&scope_node)?;
 272                     visited_modules.insert(module);
 273                     return Some((segment, scope_node, Some((scope, mod_path))));
 274                 }
 275             }
 276             Some((segment, scope_node, None))
 277         })
 278         .collect()
 279 }
 280
 281 fn reference_to_node(
 282     sema: &hir::Semantics<RootDatabase>,
 283     reference: FileReference,
 284 ) -> Option<(ast::PathSegment, SyntaxNode, hir::Module)> {
 285     let segment =
 286         reference.name.as_name_ref()?.syntax().parent().and_then(ast::PathSegment::cast)?;
 287     let parent = segment.parent_path().syntax().parent()?;
 288     let expr_or_pat = match_ast! {
 289         match parent {
 290             ast::PathExpr(_it) => parent.parent()?,
 291             ast::RecordExpr(_it) => parent,
 292             ast::TupleStructPat(_it) => parent,
 293             ast::RecordPat(_it) => parent,
 294             _ => return None,
 295         }
 296     };
 297     let module = sema.scope(&expr_or_pat).module()?;
 298     Some((segment, expr_or_pat, module))
 299 }
 300
 301 #[cfg(test)]
 302 mod tests {
 303     use ide_db::helpers::FamousDefs;
 304
 305     use crate::tests::{check_assist, check_assist_not_applicable};
 306
 307     use super::*;
 308
 309     fn check_not_applicable(ra_fixture: &str) {
 310         let fixture =
 311             format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
 312         check_assist_not_applicable(extract_struct_from_enum_variant, &fixture)
 313     }
 314
 315     #[test]
 316     fn test_extract_struct_several_fields_tuple() {
 317         check_assist(
 318             extract_struct_from_enum_variant,
 319             "enum A { $0One(u32, u32) }",
 320             r#"struct One(pub u32, pub u32);
 321
 322 enum A { One(One) }"#,
 323         );
 324     }
 325
 326     #[test]
 327     fn test_extract_struct_several_fields_named() {
 328         check_assist(
 329             extract_struct_from_enum_variant,
 330             "enum A { $0One { foo: u32, bar: u32 } }",
 331             r#"struct One{ pub foo: u32, pub bar: u32 }
 332
 333 enum A { One(One) }"#,
 334         );
 335     }
 336
 337     #[test]
 338     fn test_extract_struct_one_field_named() {
 339         check_assist(
 340             extract_struct_from_enum_variant,
 341             "enum A { $0One { foo: u32 } }",
 342             r#"struct One{ pub foo: u32 }
 343
 344 enum A { One(One) }"#,
 345         );
 346     }
 347
 348     #[test]
 349     fn test_extract_struct_carries_over_generics() {
 350         check_assist(
 351             extract_struct_from_enum_variant,
 352             r"enum En<T> { Var { a: T$0 } }",
 353             r#"struct Var<T>{ pub a: T }
 354
 355 enum En<T> { Var(Var<T>) }"#,
 356         );
 357     }
 358
 359     #[test]
 360     fn test_extract_struct_carries_over_attributes() {
 361         check_assist(
 362             extract_struct_from_enum_variant,
 363             r#"#[derive(Debug)]
 364 #[derive(Clone)]
 365 enum Enum { Variant{ field: u32$0 } }"#,
 366             r#"#[derive(Debug)]#[derive(Clone)] struct Variant{ pub field: u32 }
 367
 368 #[derive(Debug)]
 369 #[derive(Clone)]
 370 enum Enum { Variant(Variant) }"#,
 371         );
 372     }
 373
 374     #[test]
 375     fn test_extract_struct_keep_comments_and_attrs_one_field_named() {
 376         check_assist(
 377             extract_struct_from_enum_variant,
 378             r#"
 379 enum A {
 380     $0One {
 381         // leading comment
 382         /// doc comment
 383         #[an_attr]
 384         foo: u32
 385         // trailing comment
 386     }
 387 }"#,
 388             r#"
 389 struct One{
 390         // leading comment
 391         /// doc comment
 392         #[an_attr]
 393         pub foo: u32
 394         // trailing comment
 395     }
 396
 397 enum A {
 398     One(One)
 399 }"#,
 400         );
 401     }
 402
 403     #[test]
 404     fn test_extract_struct_keep_comments_and_attrs_several_fields_named() {
 405         check_assist(
 406             extract_struct_from_enum_variant,
 407             r#"
 408 enum A {
 409     $0One {
 410         // comment
 411         /// doc
 412         #[attr]
 413         foo: u32,
 414         // comment
 415         #[attr]
 416         /// doc
 417         bar: u32
 418     }
 419 }"#,
 420             r#"
 421 struct One{
 422         // comment
 423         /// doc
 424         #[attr]
 425         pub foo: u32,
 426         // comment
 427         #[attr]
 428         /// doc
 429         pub bar: u32
 430     }
 431
 432 enum A {
 433     One(One)
 434 }"#,
 435         );
 436     }
 437
 438     #[test]
 439     fn test_extract_struct_keep_comments_and_attrs_several_fields_tuple() {
 440         check_assist(
 441             extract_struct_from_enum_variant,
 442             "enum A { $0One(/* comment */ #[attr] u32, /* another */ u32 /* tail */) }",
 443             r#"
 444 struct One(/* comment */ #[attr] pub u32, /* another */ pub u32 /* tail */);
 445
 446 enum A { One(One) }"#,
 447         );
 448     }
 449
 450     #[test]
 451     fn test_extract_struct_keep_existing_visibility_named() {
 452         check_assist(
 453             extract_struct_from_enum_variant,
 454             "enum A { $0One{ pub a: u32, pub(crate) b: u32, pub(super) c: u32, d: u32 } }",
 455             r#"
 456 struct One{ pub a: u32, pub(crate) b: u32, pub(super) c: u32, pub d: u32 }
 457
 458 enum A { One(One) }"#,
 459         );
 460     }
 461
 462     #[test]
 463     fn test_extract_struct_keep_existing_visibility_tuple() {
 464         check_assist(
 465             extract_struct_from_enum_variant,
 466             "enum A { $0One(pub u32, pub(crate) u32, pub(super) u32, u32) }",
 467             r#"
 468 struct One(pub u32, pub(crate) u32, pub(super) u32, pub u32);
 469
 470 enum A { One(One) }"#,
 471         );
 472     }
 473
 474     #[test]
 475     fn test_extract_enum_variant_name_value_namespace() {
 476         check_assist(
 477             extract_struct_from_enum_variant,
 478             r#"const One: () = ();
 479 enum A { $0One(u32, u32) }"#,
 480             r#"const One: () = ();
 481 struct One(pub u32, pub u32);
 482
 483 enum A { One(One) }"#,
 484         );
 485     }
 486
 487     #[test]
 488     fn test_extract_struct_pub_visibility() {
 489         check_assist(
 490             extract_struct_from_enum_variant,
 491             "pub enum A { $0One(u32, u32) }",
 492             r#"pub struct One(pub u32, pub u32);
 493
 494 pub enum A { One(One) }"#,
 495         );
 496     }
 497
 498     #[test]
 499     fn test_extract_struct_with_complex_imports() {
 500         check_assist(
 501             extract_struct_from_enum_variant,
 502             r#"mod my_mod {
 503     fn another_fn() {
 504         let m = my_other_mod::MyEnum::MyField(1, 1);
 505     }
 506
 507     pub mod my_other_mod {
 508         fn another_fn() {
 509             let m = MyEnum::MyField(1, 1);
 510         }
 511
 512         pub enum MyEnum {
 513             $0MyField(u8, u8),
 514         }
 515     }
 516 }
 517
 518 fn another_fn() {
 519     let m = my_mod::my_other_mod::MyEnum::MyField(1, 1);
 520 }"#,
 521             r#"use my_mod::my_other_mod::MyField;
 522
 523 mod my_mod {
 524     use self::my_other_mod::MyField;
 525
 526     fn another_fn() {
 527         let m = my_other_mod::MyEnum::MyField(MyField(1, 1));
 528     }
 529
 530     pub mod my_other_mod {
 531         fn another_fn() {
 532             let m = MyEnum::MyField(MyField(1, 1));
 533         }
 534
 535         pub struct MyField(pub u8, pub u8);
 536
 537 pub enum MyEnum {
 538             MyField(MyField),
 539         }
 540     }
 541 }
 542
 543 fn another_fn() {
 544     let m = my_mod::my_other_mod::MyEnum::MyField(MyField(1, 1));
 545 }"#,
 546         );
 547     }
 548
 549     #[test]
 550     fn extract_record_fix_references() {
 551         check_assist(
 552             extract_struct_from_enum_variant,
 553             r#"
 554 enum E {
 555     $0V { i: i32, j: i32 }
 556 }
 557
 558 fn f() {
 559     let E::V { i, j } = E::V { i: 9, j: 2 };
 560 }
 561 "#,
 562             r#"
 563 struct V{ pub i: i32, pub j: i32 }
 564
 565 enum E {
 566     V(V)
 567 }
 568
 569 fn f() {
 570     let E::V(V { i, j }) = E::V(V { i: 9, j: 2 });
 571 }
 572 "#,
 573         )
 574     }
 575
 576     #[test]
 577     fn extract_record_fix_references2() {
 578         check_assist(
 579             extract_struct_from_enum_variant,
 580             r#"
 581 enum E {
 582     $0V(i32, i32)
 583 }
 584
 585 fn f() {
 586     let E::V(i, j) = E::V(9, 2);
 587 }
 588 "#,
 589             r#"
 590 struct V(pub i32, pub i32);
 591
 592 enum E {
 593     V(V)
 594 }
 595
 596 fn f() {
 597     let E::V(V(i, j)) = E::V(V(9, 2));
 598 }
 599 "#,
 600         )
 601     }
 602
 603     #[test]
 604     fn test_several_files() {
 605         check_assist(
 606             extract_struct_from_enum_variant,
 607             r#"
 608 //- /main.rs
 609 enum E {
 610     $0V(i32, i32)
 611 }
 612 mod foo;
 613
 614 //- /foo.rs
 615 use crate::E;
 616 fn f() {
 617     let e = E::V(9, 2);
 618 }
 619 "#,
 620             r#"
 621 //- /main.rs
 622 struct V(pub i32, pub i32);
 623
 624 enum E {
 625     V(V)
 626 }
 627 mod foo;
 628
 629 //- /foo.rs
 630 use crate::{E, V};
 631 fn f() {
 632     let e = E::V(V(9, 2));
 633 }
 634 "#,
 635         )
 636     }
 637
 638     #[test]
 639     fn test_several_files_record() {
 640         check_assist(
 641             extract_struct_from_enum_variant,
 642             r#"
 643 //- /main.rs
 644 enum E {
 645     $0V { i: i32, j: i32 }
 646 }
 647 mod foo;
 648
 649 //- /foo.rs
 650 use crate::E;
 651 fn f() {
 652     let e = E::V { i: 9, j: 2 };
 653 }
 654 "#,
 655             r#"
 656 //- /main.rs
 657 struct V{ pub i: i32, pub j: i32 }
 658
 659 enum E {
 660     V(V)
 661 }
 662 mod foo;
 663
 664 //- /foo.rs
 665 use crate::{E, V};
 666 fn f() {
 667     let e = E::V(V { i: 9, j: 2 });
 668 }
 669 "#,
 670         )
 671     }
 672
 673     #[test]
 674     fn test_extract_struct_record_nested_call_exp() {
 675         check_assist(
 676             extract_struct_from_enum_variant,
 677             r#"
 678 enum A { $0One { a: u32, b: u32 } }
 679
 680 struct B(A);
 681
 682 fn foo() {
 683     let _ = B(A::One { a: 1, b: 2 });
 684 }
 685 "#,
 686             r#"
 687 struct One{ pub a: u32, pub b: u32 }
 688
 689 enum A { One(One) }
 690
 691 struct B(A);
 692
 693 fn foo() {
 694     let _ = B(A::One(One { a: 1, b: 2 }));
 695 }
 696 "#,
 697         );
 698     }
 699
 700     #[test]
 701     fn test_extract_enum_not_applicable_for_element_with_no_fields() {
 702         check_not_applicable("enum A { $0One }");
 703     }
 704
 705     #[test]
 706     fn test_extract_enum_not_applicable_if_struct_exists() {
 707         check_not_applicable(
 708             r#"struct One;
 709         enum A { $0One(u8, u32) }"#,
 710         );
 711     }
 712
 713     #[test]
 714     fn test_extract_not_applicable_one_field() {
 715         check_not_applicable(r"enum A { $0One(u32) }");
 716     }
 717
 718     #[test]
 719     fn test_extract_not_applicable_no_field_tuple() {
 720         check_not_applicable(r"enum A { $0None() }");
 721     }
 722
 723     #[test]
 724     fn test_extract_not_applicable_no_field_named() {
 725         check_not_applicable(r"enum A { $0None {} }");
 726     }
 727 }