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 | 320 |
1 files changed, 320 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..44db7917a --- /dev/null +++ b/crates/ra_assists/src/handlers/extract_struct_from_enum_variant.rs | |||
@@ -0,0 +1,320 @@ | |||
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 | } | ||