crates/ra_assists/src/handlers/extract_struct_from_enum_variant.rs

   1 use hir::{EnumVariant, Module, ModuleDef, Name};
   2 use ra_db::FileId;
   3 use ra_fmt::leading_indent;
   4 use ra_ide_db::{defs::Definition, search::Reference, RootDatabase};
   5 use ra_syntax::{
   6     algo::find_node_at_offset,
   7     ast::{self, ArgListOwner, AstNode, NameOwner, VisibilityOwner},
   8     SourceFile, SyntaxNode, TextRange, TextSize,
   9 };
  10 use rustc_hash::FxHashSet;
  11
  12 use crate::{
  13     assist_context::AssistBuilder, utils::insert_use_statement, AssistContext, AssistId, Assists,
  14 };
  15
  16 // Assist: extract_struct_from_enum_variant
  17 //
  18 // Extracts a struct from enum variant.
  19 //
  20 // ```
  21 // enum A { <|>One(u32, u32) }
  22 // ```
  23 // ->
  24 // ```
  25 // struct One(pub u32, pub u32);
  26 //
  27 // enum A { One(One) }
  28 // ```
  29 pub(crate) fn extract_struct_from_enum_variant(
  30     acc: &mut Assists,
  31     ctx: &AssistContext,
  32 ) -> Option<()> {
  33     let variant = ctx.find_node_at_offset::<ast::EnumVariant>()?;
  34     let field_list = match variant.kind() {
  35         ast::StructKind::Tuple(field_list) => field_list,
  36         _ => return None,
  37     };
  38     let variant_name = variant.name()?.to_string();
  39     let variant_hir = ctx.sema.to_def(&variant)?;
  40     if existing_struct_def(ctx.db, &variant_name, &variant_hir) {
  41         return None;
  42     }
  43     let enum_ast = variant.parent_enum();
  44     let visibility = enum_ast.visibility();
  45     let enum_hir = ctx.sema.to_def(&enum_ast)?;
  46     let variant_hir_name = variant_hir.name(ctx.db);
  47     let enum_module_def = ModuleDef::from(enum_hir);
  48     let current_module = enum_hir.module(ctx.db);
  49     let target = variant.syntax().text_range();
  50     acc.add(
  51         AssistId("extract_struct_from_enum_variant"),
  52         "Extract struct from enum variant",
  53         target,
  54         |builder| {
  55             let definition = Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir));
  56             let res = definition.find_usages(&ctx.db, None);
  57             let start_offset = variant.parent_enum().syntax().text_range().start();
  58             let mut visited_modules_set = FxHashSet::default();
  59             visited_modules_set.insert(current_module);
  60             for reference in res {
  61                 let source_file = ctx.sema.parse(reference.file_range.file_id);
  62                 update_reference(
  63                     ctx,
  64                     builder,
  65                     reference,
  66                     &source_file,
  67                     &enum_module_def,
  68                     &variant_hir_name,
  69                     &mut visited_modules_set,
  70                 );
  71             }
  72             extract_struct_def(
  73                 builder,
  74                 enum_ast.syntax(),
  75                 &variant_name,
  76                 &field_list.to_string(),
  77                 start_offset,
  78                 ctx.frange.file_id,
  79                 &visibility,
  80             );
  81             let list_range = field_list.syntax().text_range();
  82             update_variant(builder, &variant_name, ctx.frange.file_id, list_range);
  83         },
  84     )
  85 }
  86
  87 fn existing_struct_def(db: &RootDatabase, variant_name: &str, variant: &EnumVariant) -> bool {
  88     variant
  89         .parent_enum(db)
  90         .module(db)
  91         .scope(db, None)
  92         .into_iter()
  93         .any(|(name, _)| name.to_string() == variant_name.to_string())
  94 }
  95
  96 fn insert_import(
  97     ctx: &AssistContext,
  98     builder: &mut AssistBuilder,
  99     path: &ast::PathExpr,
 100     module: &Module,
 101     enum_module_def: &ModuleDef,
 102     variant_hir_name: &Name,
 103 ) -> Option<()> {
 104     let db = ctx.db;
 105     let mod_path = module.find_use_path(db, enum_module_def.clone());
 106     if let Some(mut mod_path) = mod_path {
 107         mod_path.segments.pop();
 108         mod_path.segments.push(variant_hir_name.clone());
 109         insert_use_statement(path.syntax(), &mod_path, ctx, builder.text_edit_builder());
 110     }
 111     Some(())
 112 }
 113
 114 fn extract_struct_def(
 115     builder: &mut AssistBuilder,
 116     enum_ast: &SyntaxNode,
 117     variant_name: &str,
 118     variant_list: &str,
 119     start_offset: TextSize,
 120     file_id: FileId,
 121     visibility: &Option<ast::Visibility>,
 122 ) -> Option<()> {
 123     let visibility_string = if let Some(visibility) = visibility {
 124         format!("{} ", visibility.to_string())
 125     } else {
 126         "".to_string()
 127     };
 128     let indent = if let Some(indent) = leading_indent(enum_ast) {
 129         indent.to_string()
 130     } else {
 131         "".to_string()
 132     };
 133     let struct_def = format!(
 134         r#"{}struct {}{};
 135
 136 {}"#,
 137         visibility_string,
 138         variant_name,
 139         list_with_visibility(variant_list),
 140         indent
 141     );
 142     builder.edit_file(file_id);
 143     builder.insert(start_offset, struct_def);
 144     Some(())
 145 }
 146
 147 fn update_variant(
 148     builder: &mut AssistBuilder,
 149     variant_name: &str,
 150     file_id: FileId,
 151     list_range: TextRange,
 152 ) -> Option<()> {
 153     let inside_variant_range = TextRange::new(
 154         list_range.start().checked_add(TextSize::from(1))?,
 155         list_range.end().checked_sub(TextSize::from(1))?,
 156     );
 157     builder.edit_file(file_id);
 158     builder.replace(inside_variant_range, variant_name);
 159     Some(())
 160 }
 161
 162 fn update_reference(
 163     ctx: &AssistContext,
 164     builder: &mut AssistBuilder,
 165     reference: Reference,
 166     source_file: &SourceFile,
 167     enum_module_def: &ModuleDef,
 168     variant_hir_name: &Name,
 169     visited_modules_set: &mut FxHashSet<Module>,
 170 ) -> Option<()> {
 171     let path_expr: ast::PathExpr = find_node_at_offset::<ast::PathExpr>(
 172         source_file.syntax(),
 173         reference.file_range.range.start(),
 174     )?;
 175     let call = path_expr.syntax().parent().and_then(ast::CallExpr::cast)?;
 176     let list = call.arg_list()?;
 177     let segment = path_expr.path()?.segment()?;
 178     let module = ctx.sema.scope(&path_expr.syntax()).module()?;
 179     let list_range = list.syntax().text_range();
 180     let inside_list_range = TextRange::new(
 181         list_range.start().checked_add(TextSize::from(1))?,
 182         list_range.end().checked_sub(TextSize::from(1))?,
 183     );
 184     builder.edit_file(reference.file_range.file_id);
 185     if !visited_modules_set.contains(&module) {
 186         if insert_import(ctx, builder, &path_expr, &module, enum_module_def, variant_hir_name)
 187             .is_some()
 188         {
 189             visited_modules_set.insert(module);
 190         }
 191     }
 192     builder.replace(inside_list_range, format!("{}{}", segment, list));
 193     Some(())
 194 }
 195
 196 fn list_with_visibility(list: &str) -> String {
 197     list.split(',')
 198         .map(|part| {
 199             let index = if part.chars().next().unwrap() == '(' { 1usize } else { 0 };
 200             let mut mod_part = part.trim().to_string();
 201             mod_part.insert_str(index, "pub ");
 202             mod_part
 203         })
 204         .collect::<Vec<String>>()
 205         .join(", ")
 206 }
 207
 208 #[cfg(test)]
 209 mod tests {
 210
 211     use crate::{
 212         tests::{check_assist, check_assist_not_applicable},
 213         utils::FamousDefs,
 214     };
 215
 216     use super::*;
 217
 218     #[test]
 219     fn test_extract_struct_several_fields() {
 220         check_assist(
 221             extract_struct_from_enum_variant,
 222             "enum A { <|>One(u32, u32) }",
 223             r#"struct One(pub u32, pub u32);
 224
 225 enum A { One(One) }"#,
 226         );
 227     }
 228
 229     #[test]
 230     fn test_extract_struct_one_field() {
 231         check_assist(
 232             extract_struct_from_enum_variant,
 233             "enum A { <|>One(u32) }",
 234             r#"struct One(pub u32);
 235
 236 enum A { One(One) }"#,
 237         );
 238     }
 239
 240     #[test]
 241     fn test_extract_struct_pub_visibility() {
 242         check_assist(
 243             extract_struct_from_enum_variant,
 244             "pub enum A { <|>One(u32, u32) }",
 245             r#"pub struct One(pub u32, pub u32);
 246
 247 pub enum A { One(One) }"#,
 248         );
 249     }
 250
 251     #[test]
 252     fn test_extract_struct_with_complex_imports() {
 253         check_assist(
 254             extract_struct_from_enum_variant,
 255             r#"mod my_mod {
 256     fn another_fn() {
 257         let m = my_other_mod::MyEnum::MyField(1, 1);
 258     }
 259
 260     pub mod my_other_mod {
 261         fn another_fn() {
 262             let m = MyEnum::MyField(1, 1);
 263         }
 264
 265         pub enum MyEnum {
 266             <|>MyField(u8, u8),
 267         }
 268     }
 269 }
 270
 271 fn another_fn() {
 272     let m = my_mod::my_other_mod::MyEnum::MyField(1, 1);
 273 }"#,
 274             r#"use my_mod::my_other_mod::MyField;
 275
 276 mod my_mod {
 277     use my_other_mod::MyField;
 278
 279     fn another_fn() {
 280         let m = my_other_mod::MyEnum::MyField(MyField(1, 1));
 281     }
 282
 283     pub mod my_other_mod {
 284         fn another_fn() {
 285             let m = MyEnum::MyField(MyField(1, 1));
 286         }
 287
 288         pub struct MyField(pub u8, pub u8);
 289
 290         pub enum MyEnum {
 291             MyField(MyField),
 292         }
 293     }
 294 }
 295
 296 fn another_fn() {
 297     let m = my_mod::my_other_mod::MyEnum::MyField(MyField(1, 1));
 298 }"#,
 299         );
 300     }
 301
 302     fn check_not_applicable(ra_fixture: &str) {
 303         let fixture =
 304             format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
 305         check_assist_not_applicable(extract_struct_from_enum_variant, &fixture)
 306     }
 307
 308     #[test]
 309     fn test_extract_enum_not_applicable_for_element_with_no_fields() {
 310         check_not_applicable("enum A { <|>One }");
 311     }
 312
 313     #[test]
 314     fn test_extract_enum_not_applicable_if_struct_exists() {
 315         check_not_applicable(
 316             r#"struct One;
 317         enum A { <|>One(u8) }"#,
 318         );
 319     }
 320 }