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