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 = Some(make::visibility_pub());
 155     let field_list = match field_list {
 156         Either::Left(field_list) => {
 157             make::record_field_list(field_list.fields().flat_map(|field| {
 158                 Some(make::record_field(pub_vis.clone(), field.name()?, field.ty()?))
 159             }))
 160             .into()
 161         }
 162         Either::Right(field_list) => make::tuple_field_list(
 163             field_list
 164                 .fields()
 165                 .flat_map(|field| Some(make::tuple_field(pub_vis.clone(), field.ty()?))),
 166         )
 167         .into(),
 168     };
 169
 170     make::struct_(visibility, variant_name, None, field_list).clone_for_update()
 171 }
 172
 173 fn update_variant(variant: &ast::Variant) -> Option<()> {
 174     let name = variant.name()?;
 175     let tuple_field = make::tuple_field(None, make::ty(&name.text()));
 176     let replacement = make::variant(
 177         name,
 178         Some(ast::FieldList::TupleFieldList(make::tuple_field_list(iter::once(tuple_field)))),
 179     )
 180     .clone_for_update();
 181     ted::replace(variant.syntax(), replacement.syntax());
 182     Some(())
 183 }
 184
 185 fn apply_references(
 186     insert_use_cfg: InsertUseConfig,
 187     segment: ast::PathSegment,
 188     node: SyntaxNode,
 189     import: Option<(ImportScope, hir::ModPath)>,
 190 ) {
 191     if let Some((scope, path)) = import {
 192         insert_use(&scope, mod_path_to_ast(&path), insert_use_cfg);
 193     }
 194     ted::insert_raw(
 195         ted::Position::before(segment.syntax()),
 196         make::path_from_text(&format!("{}", segment)).clone_for_update().syntax(),
 197     );
 198     ted::insert_raw(ted::Position::before(segment.syntax()), make::token(T!['(']));
 199     ted::insert_raw(ted::Position::after(&node), make::token(T![')']));
 200 }
 201
 202 fn process_references(
 203     ctx: &AssistContext,
 204     visited_modules: &mut FxHashSet<Module>,
 205     source_file: &SyntaxNode,
 206     enum_module_def: &ModuleDef,
 207     variant_hir_name: &Name,
 208     refs: Vec<FileReference>,
 209 ) -> Vec<(ast::PathSegment, SyntaxNode, Option<(ImportScope, hir::ModPath)>)> {
 210     // we have to recollect here eagerly as we are about to edit the tree we need to calculate the changes
 211     // and corresponding nodes up front
 212     refs.into_iter()
 213         .flat_map(|reference| {
 214             let (segment, scope_node, module) =
 215                 reference_to_node(&ctx.sema, source_file, reference)?;
 216             if !visited_modules.contains(&module) {
 217                 let mod_path = module.find_use_path_prefixed(
 218                     ctx.sema.db,
 219                     *enum_module_def,
 220                     ctx.config.insert_use.prefix_kind,
 221                 );
 222                 if let Some(mut mod_path) = mod_path {
 223                     mod_path.pop_segment();
 224                     mod_path.push_segment(variant_hir_name.clone());
 225                     let scope = ImportScope::find_insert_use_container(&scope_node)?;
 226                     visited_modules.insert(module);
 227                     return Some((segment, scope_node, Some((scope, mod_path))));
 228                 }
 229             }
 230             Some((segment, scope_node, None))
 231         })
 232         .collect()
 233 }
 234
 235 fn reference_to_node(
 236     sema: &hir::Semantics<RootDatabase>,
 237     source_file: &SyntaxNode,
 238     reference: FileReference,
 239 ) -> Option<(ast::PathSegment, SyntaxNode, hir::Module)> {
 240     let offset = reference.range.start();
 241     if let Some(path_expr) = find_node_at_offset::<ast::PathExpr>(source_file, offset) {
 242         // tuple variant
 243         Some((path_expr.path()?.segment()?, path_expr.syntax().parent()?))
 244     } else if let Some(record_expr) = find_node_at_offset::<ast::RecordExpr>(source_file, offset) {
 245         // record variant
 246         Some((record_expr.path()?.segment()?, record_expr.syntax().clone()))
 247     } else {
 248         None
 249     }
 250     .and_then(|(segment, expr)| {
 251         let module = sema.scope(&expr).module()?;
 252         Some((segment, expr, module))
 253     })
 254 }
 255
 256 #[cfg(test)]
 257 mod tests {
 258     use ide_db::helpers::FamousDefs;
 259
 260     use crate::tests::{check_assist, check_assist_not_applicable};
 261
 262     use super::*;
 263
 264     #[test]
 265     fn test_extract_struct_several_fields_tuple() {
 266         check_assist(
 267             extract_struct_from_enum_variant,
 268             "enum A { $0One(u32, u32) }",
 269             r#"struct One(pub u32, pub u32);
 270
 271 enum A { One(One) }"#,
 272         );
 273     }
 274
 275     #[test]
 276     fn test_extract_struct_several_fields_named() {
 277         check_assist(
 278             extract_struct_from_enum_variant,
 279             "enum A { $0One { foo: u32, bar: u32 } }",
 280             r#"struct One{ pub foo: u32, pub bar: u32 }
 281
 282 enum A { One(One) }"#,
 283         );
 284     }
 285
 286     #[test]
 287     fn test_extract_struct_one_field_named() {
 288         check_assist(
 289             extract_struct_from_enum_variant,
 290             "enum A { $0One { foo: u32 } }",
 291             r#"struct One{ pub foo: u32 }
 292
 293 enum A { One(One) }"#,
 294         );
 295     }
 296
 297     #[test]
 298     fn test_extract_enum_variant_name_value_namespace() {
 299         check_assist(
 300             extract_struct_from_enum_variant,
 301             r#"const One: () = ();
 302 enum A { $0One(u32, u32) }"#,
 303             r#"const One: () = ();
 304 struct One(pub u32, pub u32);
 305
 306 enum A { One(One) }"#,
 307         );
 308     }
 309
 310     #[test]
 311     fn test_extract_struct_pub_visibility() {
 312         check_assist(
 313             extract_struct_from_enum_variant,
 314             "pub enum A { $0One(u32, u32) }",
 315             r#"pub struct One(pub u32, pub u32);
 316
 317 pub enum A { One(One) }"#,
 318         );
 319     }
 320
 321     #[test]
 322     fn test_extract_struct_with_complex_imports() {
 323         check_assist(
 324             extract_struct_from_enum_variant,
 325             r#"mod my_mod {
 326     fn another_fn() {
 327         let m = my_other_mod::MyEnum::MyField(1, 1);
 328     }
 329
 330     pub mod my_other_mod {
 331         fn another_fn() {
 332             let m = MyEnum::MyField(1, 1);
 333         }
 334
 335         pub enum MyEnum {
 336             $0MyField(u8, u8),
 337         }
 338     }
 339 }
 340
 341 fn another_fn() {
 342     let m = my_mod::my_other_mod::MyEnum::MyField(1, 1);
 343 }"#,
 344             r#"use my_mod::my_other_mod::MyField;
 345
 346 mod my_mod {
 347     use self::my_other_mod::MyField;
 348
 349     fn another_fn() {
 350         let m = my_other_mod::MyEnum::MyField(MyField(1, 1));
 351     }
 352
 353     pub mod my_other_mod {
 354         fn another_fn() {
 355             let m = MyEnum::MyField(MyField(1, 1));
 356         }
 357
 358         pub struct MyField(pub u8, pub u8);
 359
 360 pub enum MyEnum {
 361             MyField(MyField),
 362         }
 363     }
 364 }
 365
 366 fn another_fn() {
 367     let m = my_mod::my_other_mod::MyEnum::MyField(MyField(1, 1));
 368 }"#,
 369         );
 370     }
 371
 372     #[test]
 373     fn extract_record_fix_references() {
 374         check_assist(
 375             extract_struct_from_enum_variant,
 376             r#"
 377 enum E {
 378     $0V { i: i32, j: i32 }
 379 }
 380
 381 fn f() {
 382     let e = E::V { i: 9, j: 2 };
 383 }
 384 "#,
 385             r#"
 386 struct V{ pub i: i32, pub j: i32 }
 387
 388 enum E {
 389     V(V)
 390 }
 391
 392 fn f() {
 393     let e = E::V(V { i: 9, j: 2 });
 394 }
 395 "#,
 396         )
 397     }
 398
 399     #[test]
 400     fn test_several_files() {
 401         check_assist(
 402             extract_struct_from_enum_variant,
 403             r#"
 404 //- /main.rs
 405 enum E {
 406     $0V(i32, i32)
 407 }
 408 mod foo;
 409
 410 //- /foo.rs
 411 use crate::E;
 412 fn f() {
 413     let e = E::V(9, 2);
 414 }
 415 "#,
 416             r#"
 417 //- /main.rs
 418 struct V(pub i32, pub i32);
 419
 420 enum E {
 421     V(V)
 422 }
 423 mod foo;
 424
 425 //- /foo.rs
 426 use crate::{E, V};
 427 fn f() {
 428     let e = E::V(V(9, 2));
 429 }
 430 "#,
 431         )
 432     }
 433
 434     #[test]
 435     fn test_several_files_record() {
 436         check_assist(
 437             extract_struct_from_enum_variant,
 438             r#"
 439 //- /main.rs
 440 enum E {
 441     $0V { i: i32, j: i32 }
 442 }
 443 mod foo;
 444
 445 //- /foo.rs
 446 use crate::E;
 447 fn f() {
 448     let e = E::V { i: 9, j: 2 };
 449 }
 450 "#,
 451             r#"
 452 //- /main.rs
 453 struct V{ pub i: i32, pub j: i32 }
 454
 455 enum E {
 456     V(V)
 457 }
 458 mod foo;
 459
 460 //- /foo.rs
 461 use crate::{E, V};
 462 fn f() {
 463     let e = E::V(V { i: 9, j: 2 });
 464 }
 465 "#,
 466         )
 467     }
 468
 469     #[test]
 470     fn test_extract_struct_record_nested_call_exp() {
 471         check_assist(
 472             extract_struct_from_enum_variant,
 473             r#"
 474 enum A { $0One { a: u32, b: u32 } }
 475
 476 struct B(A);
 477
 478 fn foo() {
 479     let _ = B(A::One { a: 1, b: 2 });
 480 }
 481 "#,
 482             r#"
 483 struct One{ pub a: u32, pub b: u32 }
 484
 485 enum A { One(One) }
 486
 487 struct B(A);
 488
 489 fn foo() {
 490     let _ = B(A::One(One { a: 1, b: 2 }));
 491 }
 492 "#,
 493         );
 494     }
 495
 496     fn check_not_applicable(ra_fixture: &str) {
 497         let fixture =
 498             format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
 499         check_assist_not_applicable(extract_struct_from_enum_variant, &fixture)
 500     }
 501
 502     #[test]
 503     fn test_extract_enum_not_applicable_for_element_with_no_fields() {
 504         check_not_applicable("enum A { $0One }");
 505     }
 506
 507     #[test]
 508     fn test_extract_enum_not_applicable_if_struct_exists() {
 509         check_not_applicable(
 510             r#"struct One;
 511         enum A { $0One(u8, u32) }"#,
 512         );
 513     }
 514
 515     #[test]
 516     fn test_extract_not_applicable_one_field() {
 517         check_not_applicable(r"enum A { $0One(u32) }");
 518     }
 519
 520     #[test]
 521     fn test_extract_not_applicable_no_field_tuple() {
 522         check_not_applicable(r"enum A { $0None() }");
 523     }
 524
 525     #[test]
 526     fn test_extract_not_applicable_no_field_named() {
 527         check_not_applicable(r"enum A { $0None {} }");
 528     }
 529 }