From 5cd4eb6dd6d8c733077a6aeea5d2cc0812ded096 Mon Sep 17 00:00:00 2001 From: Mikhail Rakhmanov Date: Fri, 22 May 2020 22:28:30 +0200 Subject: Add preliminary implementation of extract struct from enum variant --- crates/ra_assists/Cargo.toml | 1 + crates/ra_assists/src/assist_context.rs | 62 +++- .../handlers/extract_struct_from_enum_variant.rs | 338 +++++++++++++++++++++ crates/ra_assists/src/lib.rs | 2 + 4 files changed, 402 insertions(+), 1 deletion(-) create mode 100644 crates/ra_assists/src/handlers/extract_struct_from_enum_variant.rs (limited to 'crates/ra_assists') diff --git a/crates/ra_assists/Cargo.toml b/crates/ra_assists/Cargo.toml index 3bcf58ba4..f3481bdeb 100644 --- a/crates/ra_assists/Cargo.toml +++ b/crates/ra_assists/Cargo.toml @@ -20,5 +20,6 @@ ra_fmt = { path = "../ra_fmt" } ra_prof = { path = "../ra_prof" } ra_db = { path = "../ra_db" } ra_ide_db = { path = "../ra_ide_db" } +hir_expand = { path = "../ra_hir_expand", package = "ra_hir_expand" } hir = { path = "../ra_hir", package = "ra_hir" } test_utils = { path = "../test_utils" } diff --git a/crates/ra_assists/src/assist_context.rs b/crates/ra_assists/src/assist_context.rs index 5b1a4680b..6291c68de 100644 --- a/crates/ra_assists/src/assist_context.rs +++ b/crates/ra_assists/src/assist_context.rs @@ -2,7 +2,7 @@ use algo::find_covering_element; use hir::Semantics; -use ra_db::{FileId, FileRange}; +use ra_db::{FileId, FileRange, FilePosition}; use ra_fmt::{leading_indent, reindent}; use ra_ide_db::{ source_change::{SourceChange, SourceFileEdit}, @@ -19,6 +19,7 @@ use crate::{ assist_config::{AssistConfig, SnippetCap}, Assist, AssistId, GroupLabel, ResolvedAssist, }; +use rustc_hash::FxHashMap; /// `AssistContext` allows to apply an assist or check if it could be applied. /// @@ -138,6 +139,16 @@ impl Assists { let label = Assist::new(id, label.into(), None, target); self.add_impl(label, f) } + pub(crate) fn add_in_multiple_files( + &mut self, + id: AssistId, + label: impl Into, + target: TextRange, + f: impl FnOnce(&mut AssistDirector), + ) -> Option<()> { + let label = Assist::new(id, label.into(), None, target); + self.add_impl_multiple_files(label, f) + } pub(crate) fn add_group( &mut self, group: &GroupLabel, @@ -162,6 +173,27 @@ impl Assists { Some(()) } + fn add_impl_multiple_files(&mut self, label: Assist, f: impl FnOnce(&mut AssistDirector)) -> Option<()> { + let change_label = label.label.clone(); + if !self.resolve { + return None + } + let mut director = AssistDirector::new(change_label.clone()); + f(&mut director); + let changes = director.finish(); + let file_edits: Vec = changes.into_iter() + .map(|mut change| change.source_file_edits.pop().unwrap()).collect(); + + let source_change = SourceChange { + source_file_edits: file_edits, + file_system_edits: vec![], + is_snippet: false, + }; + + self.buf.push((label, Some(source_change))); + Some(()) + } + fn finish(mut self) -> Vec<(Assist, Option)> { self.buf.sort_by_key(|(label, _edit)| label.target.len()); self.buf @@ -255,3 +287,31 @@ impl AssistBuilder { res } } + +pub(crate) struct AssistDirector { + source_changes: Vec, + builders: FxHashMap, + change_label: String +} + +impl AssistDirector { + fn new(change_label: String) -> AssistDirector { + AssistDirector { + source_changes: vec![], + builders: FxHashMap::default(), + change_label + } + } + + pub(crate) fn perform(&mut self, file_id: FileId, f: impl FnOnce(&mut AssistBuilder)) { + let mut builder = self.builders.entry(file_id).or_insert(AssistBuilder::new(file_id)); + f(&mut builder); + } + + fn finish(mut self) -> Vec { + for (file_id, builder) in self.builders.into_iter().collect::>() { + self.source_changes.push(builder.finish()); + } + self.source_changes + } +} diff --git a/crates/ra_assists/src/handlers/extract_struct_from_enum_variant.rs b/crates/ra_assists/src/handlers/extract_struct_from_enum_variant.rs new file mode 100644 index 000000000..6e19a6feb --- /dev/null +++ b/crates/ra_assists/src/handlers/extract_struct_from_enum_variant.rs @@ -0,0 +1,338 @@ +use hir_expand::name::AsName; +use ra_ide_db::{ + defs::Definition, imports_locator::ImportsLocator, search::Reference, RootDatabase, +}; +use ra_syntax::{ + algo::find_node_at_offset, + ast::{self, AstNode, NameOwner}, + SourceFile, SyntaxNode, TextRange, TextSize, +}; +use stdx::format_to; + +use crate::{ + assist_context::{AssistBuilder, AssistDirector}, + utils::insert_use_statement, + AssistContext, AssistId, Assists, +}; +use ast::{ArgListOwner, VisibilityOwner}; +use hir::{EnumVariant, Module, ModuleDef}; +use ra_fmt::leading_indent; +use rustc_hash::FxHashSet; +use ra_db::FileId; + +// Assist extract_struct_from_enum +// +// Extracts a from struct from enum variant +// +// ``` +// enum A { <|>One(u32, u32) } +// ``` +// -> +// ``` +// struct One(pub u32, pub u32); +// +// enum A { One(One) }" +// ``` +pub(crate) fn extract_struct_from_enum(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { + let variant = ctx.find_node_at_offset::()?; + let field_list = match variant.kind() { + ast::StructKind::Tuple(field_list) => field_list, + _ => return None, + }; + let variant_name = variant.name()?.to_string(); + let enum_ast = variant.parent_enum(); + let enum_name = enum_ast.name().unwrap().to_string(); + let visibility = enum_ast.visibility(); + let variant_hir = ctx.sema.to_def(&variant)?; + + if existing_struct_def(ctx.db, &variant_name, &variant_hir) { + return None; + } + + let target = variant.syntax().text_range(); + return acc.add_in_multiple_files( + AssistId("extract_struct_from_enum_variant"), + "Extract struct from enum variant", + target, + |edit| { + let definition = Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir)); + let res = definition.find_usages(&ctx.db, None); + let module_def = mod_def_for_target_module(ctx, &enum_name); + let start_offset = variant.parent_enum().syntax().text_range().start(); + let mut seen_files_map: FxHashSet = FxHashSet::default(); + seen_files_map.insert(module_def.module(ctx.db).unwrap()); + for reference in res { + let source_file = ctx.sema.parse(reference.file_range.file_id); + update_reference( + ctx, + edit, + reference, + &source_file, + &module_def, + &mut seen_files_map, + ); + } + extract_struct_def( + edit, + enum_ast.syntax(), + &variant_name, + &field_list.to_string(), + start_offset, + ctx.frange.file_id, + &visibility, + ); + let list_range = field_list.syntax().text_range(); + update_variant(edit, &variant_name, ctx.frange.file_id, list_range); + }, + ); +} + +fn existing_struct_def(db: &RootDatabase, variant_name: &str, variant: &EnumVariant) -> bool { + let module_defs = variant.parent_enum(db).module(db).scope(db, None); + for (name, _) in module_defs { + if name.to_string() == variant_name.to_string() { + return true; + } + } + false +} + +fn mod_def_for_target_module(ctx: &AssistContext, enum_name: &str) -> ModuleDef { + ImportsLocator::new(ctx.db).find_imports(enum_name).first().unwrap().left().unwrap() +} + +fn insert_use_import( + ctx: &AssistContext, + builder: &mut AssistBuilder, + path: &ast::PathExpr, + module: &Module, + module_def: &ModuleDef, + path_segment: ast::NameRef, +) -> Option<()> { + let db = ctx.db; + let mod_path = module.find_use_path(db, module_def.clone()); + if let Some(mut mod_path) = mod_path { + mod_path.segments.pop(); + mod_path.segments.push(path_segment.as_name()); + insert_use_statement(path.syntax(), &mod_path, ctx, builder.text_edit_builder()); + } + Some(()) +} + +fn extract_struct_def( + edit: &mut AssistDirector, + enum_ast: &SyntaxNode, + variant_name: &str, + variant_list: &str, + start_offset: TextSize, + file_id: FileId, + visibility: &Option, +) -> Option<()> { + let visibility_string = if let Some(visibility) = visibility { + format!("{} ", visibility.to_string()) + } else { + "".to_string() + }; + let mut buf = String::new(); + let indent = if let Some(indent) = leading_indent(enum_ast) { + indent.to_string() + } else { + "".to_string() + }; + + format_to!( + buf, + r#"{}struct {}{}; + +{}"#, + visibility_string, + variant_name, + list_with_visibility(variant_list), + indent + ); + edit.perform(file_id, |builder| { + builder.insert(start_offset, buf); + }); + Some(()) +} + +fn update_variant( + edit: &mut AssistDirector, + variant_name: &str, + file_id: FileId, + list_range: TextRange, +) -> Option<()> { + let inside_variant_range = TextRange::new( + list_range.start().checked_add(TextSize::from(1))?, + list_range.end().checked_sub(TextSize::from(1))?, + ); + edit.perform(file_id, |builder| { + builder.set_file(file_id); + builder.replace(inside_variant_range, variant_name); + }); + Some(()) +} + +fn update_reference( + ctx: &AssistContext, + edit: &mut AssistDirector, + reference: Reference, + source_file: &SourceFile, + module_def: &ModuleDef, + seen_files_map: &mut FxHashSet, +) -> Option<()> { + let path_expr: ast::PathExpr = find_node_at_offset::( + source_file.syntax(), + reference.file_range.range.start(), + )?; + let call = path_expr.syntax().parent().and_then(ast::CallExpr::cast)?; + let list = call.arg_list()?; + let segment = path_expr.path()?.segment()?; + let list_range = list.syntax().text_range(); + let inside_list_range = TextRange::new( + list_range.start().checked_add(TextSize::from(1))?, + list_range.end().checked_sub(TextSize::from(1))?, + ); + edit.perform(reference.file_range.file_id, |builder| { + let module = ctx.sema.scope(&path_expr.syntax()).module().unwrap(); + if !seen_files_map.contains(&module) { + if insert_use_import( + ctx, + builder, + &path_expr, + &module, + module_def, + segment.name_ref().unwrap(), + ) + .is_some() + { + seen_files_map.insert(module); + } + } + builder.replace(inside_list_range, format!("{}{}", segment, list)); + }); + Some(()) +} + +fn list_with_visibility(list: &str) -> String { + list.split(',') + .map(|part| { + let index = if part.chars().next().unwrap() == '(' { 1usize } else { 0 }; + let mut mod_part = part.trim().to_string(); + mod_part.insert_str(index, "pub "); + mod_part + }) + .collect::>() + .join(", ") +} + +#[cfg(test)] +mod tests { + + use crate::{utils::FamousDefs, tests::{check_assist, check_assist_not_applicable}}; + + use super::*; + + #[test] + fn test_extract_struct_several_fields() { + check_assist( + extract_struct_from_enum, + "enum A { <|>One(u32, u32) }", + r#"struct One(pub u32, pub u32); + +enum A { One(One) }"#, + ); + } + + #[test] + fn test_extract_struct_one_field() { + check_assist( + extract_struct_from_enum, + "enum A { <|>One(u32) }", + r#"struct One(pub u32); + +enum A { One(One) }"#, + ); + } + + #[test] + fn test_extract_struct_pub_visibility() { + check_assist( + extract_struct_from_enum, + "pub enum A { <|>One(u32, u32) }", + r#"pub struct One(pub u32, pub u32); + +pub enum A { One(One) }"#, + ); + } + + #[test] + fn test_extract_struct_with_complex_imports() { + check_assist( + extract_struct_from_enum, + r#"mod my_mod { + fn another_fn() { + let m = my_other_mod::MyEnum::MyField(1, 1); + } + + pub mod my_other_mod { + fn another_fn() { + let m = MyEnum::MyField(1, 1); + } + + pub enum MyEnum { + <|>MyField(u8, u8), + } + } +} + +fn another_fn() { + let m = my_mod::my_other_mod::MyEnum::MyField(1, 1); +}"#, + r#"use my_mod::my_other_mod::MyField; + +mod my_mod { + use my_other_mod::MyField; + + fn another_fn() { + let m = my_other_mod::MyEnum::MyField(MyField(1, 1)); + } + + pub mod my_other_mod { + fn another_fn() { + let m = MyEnum::MyField(MyField(1, 1)); + } + + pub struct MyField(pub u8, pub u8); + + pub enum MyEnum { + MyField(MyField), + } + } +} + +fn another_fn() { + let m = my_mod::my_other_mod::MyEnum::MyField(MyField(1, 1)); +}"#, + ); + } + + fn check_not_applicable(ra_fixture: &str) { + let fixture = + format!("//- main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE); + check_assist_not_applicable(extract_struct_from_enum, &fixture) + } + + #[test] + fn test_extract_enum_not_applicable_for_element_with_no_fields() { + check_not_applicable("enum A { <|>One }"); + } + + #[test] + fn test_extract_enum_not_applicable_if_struct_exists() { + check_not_applicable( + r#"struct One; + enum A { <|>One(u8) }"#, + ); + } +} diff --git a/crates/ra_assists/src/lib.rs b/crates/ra_assists/src/lib.rs index 464bc03dd..9933f7a50 100644 --- a/crates/ra_assists/src/lib.rs +++ b/crates/ra_assists/src/lib.rs @@ -115,6 +115,7 @@ mod handlers { mod change_return_type_to_result; mod change_visibility; mod early_return; + mod extract_struct_from_enum_variant; mod fill_match_arms; mod fix_visibility; mod flip_binexpr; @@ -154,6 +155,7 @@ mod handlers { change_return_type_to_result::change_return_type_to_result, change_visibility::change_visibility, early_return::convert_to_guarded_return, + extract_struct_from_enum_variant::extract_struct_from_enum, fill_match_arms::fill_match_arms, fix_visibility::fix_visibility, flip_binexpr::flip_binexpr, -- cgit v1.2.3