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