diff options
Diffstat (limited to 'crates/ide_assists/src/handlers/generate_enum_projection_method.rs')
-rw-r--r-- | crates/ide_assists/src/handlers/generate_enum_projection_method.rs | 331 |
1 files changed, 331 insertions, 0 deletions
diff --git a/crates/ide_assists/src/handlers/generate_enum_projection_method.rs b/crates/ide_assists/src/handlers/generate_enum_projection_method.rs new file mode 100644 index 000000000..871bcab50 --- /dev/null +++ b/crates/ide_assists/src/handlers/generate_enum_projection_method.rs | |||
@@ -0,0 +1,331 @@ | |||
1 | use itertools::Itertools; | ||
2 | use stdx::to_lower_snake_case; | ||
3 | use syntax::ast::VisibilityOwner; | ||
4 | use syntax::ast::{self, AstNode, NameOwner}; | ||
5 | |||
6 | use crate::{ | ||
7 | utils::{add_method_to_adt, find_struct_impl}, | ||
8 | AssistContext, AssistId, AssistKind, Assists, | ||
9 | }; | ||
10 | |||
11 | // Assist: generate_enum_try_into_method | ||
12 | // | ||
13 | // Generate an `try_into_` method for an enum variant. | ||
14 | // | ||
15 | // ``` | ||
16 | // enum Value { | ||
17 | // Number(i32), | ||
18 | // Text(String)$0, | ||
19 | // } | ||
20 | // ``` | ||
21 | // -> | ||
22 | // ``` | ||
23 | // enum Value { | ||
24 | // Number(i32), | ||
25 | // Text(String), | ||
26 | // } | ||
27 | // | ||
28 | // impl Value { | ||
29 | // fn try_into_text(self) -> Result<String, Self> { | ||
30 | // if let Self::Text(v) = self { | ||
31 | // Ok(v) | ||
32 | // } else { | ||
33 | // Err(self) | ||
34 | // } | ||
35 | // } | ||
36 | // } | ||
37 | // ``` | ||
38 | pub(crate) fn generate_enum_try_into_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { | ||
39 | generate_enum_projection_method( | ||
40 | acc, | ||
41 | ctx, | ||
42 | "generate_enum_try_into_method", | ||
43 | "Generate an `try_into_` method for an enum variant", | ||
44 | ProjectionProps { | ||
45 | fn_name_prefix: "try_into", | ||
46 | self_param: "self", | ||
47 | return_prefix: "Result<", | ||
48 | return_suffix: ", Self>", | ||
49 | happy_case: "Ok", | ||
50 | sad_case: "Err(self)", | ||
51 | }, | ||
52 | ) | ||
53 | } | ||
54 | |||
55 | // Assist: generate_enum_as_method | ||
56 | // | ||
57 | // Generate an `as_` method for an enum variant. | ||
58 | // | ||
59 | // ``` | ||
60 | // enum Value { | ||
61 | // Number(i32), | ||
62 | // Text(String)$0, | ||
63 | // } | ||
64 | // ``` | ||
65 | // -> | ||
66 | // ``` | ||
67 | // enum Value { | ||
68 | // Number(i32), | ||
69 | // Text(String), | ||
70 | // } | ||
71 | // | ||
72 | // impl Value { | ||
73 | // fn as_text(&self) -> Option<&String> { | ||
74 | // if let Self::Text(v) = self { | ||
75 | // Some(v) | ||
76 | // } else { | ||
77 | // None | ||
78 | // } | ||
79 | // } | ||
80 | // } | ||
81 | // ``` | ||
82 | pub(crate) fn generate_enum_as_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { | ||
83 | generate_enum_projection_method( | ||
84 | acc, | ||
85 | ctx, | ||
86 | "generate_enum_as_method", | ||
87 | "Generate an `as_` method for an enum variant", | ||
88 | ProjectionProps { | ||
89 | fn_name_prefix: "as", | ||
90 | self_param: "&self", | ||
91 | return_prefix: "Option<&", | ||
92 | return_suffix: ">", | ||
93 | happy_case: "Some", | ||
94 | sad_case: "None", | ||
95 | }, | ||
96 | ) | ||
97 | } | ||
98 | |||
99 | struct ProjectionProps { | ||
100 | fn_name_prefix: &'static str, | ||
101 | self_param: &'static str, | ||
102 | return_prefix: &'static str, | ||
103 | return_suffix: &'static str, | ||
104 | happy_case: &'static str, | ||
105 | sad_case: &'static str, | ||
106 | } | ||
107 | |||
108 | fn generate_enum_projection_method( | ||
109 | acc: &mut Assists, | ||
110 | ctx: &AssistContext, | ||
111 | assist_id: &'static str, | ||
112 | assist_description: &str, | ||
113 | props: ProjectionProps, | ||
114 | ) -> Option<()> { | ||
115 | let variant = ctx.find_node_at_offset::<ast::Variant>()?; | ||
116 | let variant_name = variant.name()?; | ||
117 | let parent_enum = ast::Adt::Enum(variant.parent_enum()); | ||
118 | |||
119 | let (pattern_suffix, field_type, bound_name) = match variant.kind() { | ||
120 | ast::StructKind::Record(record) => { | ||
121 | let (field,) = record.fields().collect_tuple()?; | ||
122 | let name = field.name()?.to_string(); | ||
123 | let ty = field.ty()?; | ||
124 | let pattern_suffix = format!(" {{ {} }}", name); | ||
125 | (pattern_suffix, ty, name) | ||
126 | } | ||
127 | ast::StructKind::Tuple(tuple) => { | ||
128 | let (field,) = tuple.fields().collect_tuple()?; | ||
129 | let ty = field.ty()?; | ||
130 | ("(v)".to_owned(), ty, "v".to_owned()) | ||
131 | } | ||
132 | ast::StructKind::Unit => return None, | ||
133 | }; | ||
134 | |||
135 | let fn_name = format!("{}_{}", props.fn_name_prefix, &to_lower_snake_case(variant_name.text())); | ||
136 | |||
137 | // Return early if we've found an existing new fn | ||
138 | let impl_def = find_struct_impl(&ctx, &parent_enum, &fn_name)?; | ||
139 | |||
140 | let target = variant.syntax().text_range(); | ||
141 | acc.add(AssistId(assist_id, AssistKind::Generate), assist_description, target, |builder| { | ||
142 | let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v)); | ||
143 | let method = format!( | ||
144 | " {0}fn {1}({2}) -> {3}{4}{5} {{ | ||
145 | if let Self::{6}{7} = self {{ | ||
146 | {8}({9}) | ||
147 | }} else {{ | ||
148 | {10} | ||
149 | }} | ||
150 | }}", | ||
151 | vis, | ||
152 | fn_name, | ||
153 | props.self_param, | ||
154 | props.return_prefix, | ||
155 | field_type.syntax(), | ||
156 | props.return_suffix, | ||
157 | variant_name, | ||
158 | pattern_suffix, | ||
159 | props.happy_case, | ||
160 | bound_name, | ||
161 | props.sad_case, | ||
162 | ); | ||
163 | |||
164 | add_method_to_adt(builder, &parent_enum, impl_def, &method); | ||
165 | }) | ||
166 | } | ||
167 | |||
168 | #[cfg(test)] | ||
169 | mod tests { | ||
170 | use crate::tests::{check_assist, check_assist_not_applicable}; | ||
171 | |||
172 | use super::*; | ||
173 | |||
174 | #[test] | ||
175 | fn test_generate_enum_try_into_tuple_variant() { | ||
176 | check_assist( | ||
177 | generate_enum_try_into_method, | ||
178 | r#" | ||
179 | enum Value { | ||
180 | Number(i32), | ||
181 | Text(String)$0, | ||
182 | }"#, | ||
183 | r#"enum Value { | ||
184 | Number(i32), | ||
185 | Text(String), | ||
186 | } | ||
187 | |||
188 | impl Value { | ||
189 | fn try_into_text(self) -> Result<String, Self> { | ||
190 | if let Self::Text(v) = self { | ||
191 | Ok(v) | ||
192 | } else { | ||
193 | Err(self) | ||
194 | } | ||
195 | } | ||
196 | }"#, | ||
197 | ); | ||
198 | } | ||
199 | |||
200 | #[test] | ||
201 | fn test_generate_enum_try_into_already_implemented() { | ||
202 | check_assist_not_applicable( | ||
203 | generate_enum_try_into_method, | ||
204 | r#"enum Value { | ||
205 | Number(i32), | ||
206 | Text(String)$0, | ||
207 | } | ||
208 | |||
209 | impl Value { | ||
210 | fn try_into_text(self) -> Result<String, Self> { | ||
211 | if let Self::Text(v) = self { | ||
212 | Ok(v) | ||
213 | } else { | ||
214 | Err(self) | ||
215 | } | ||
216 | } | ||
217 | }"#, | ||
218 | ); | ||
219 | } | ||
220 | |||
221 | #[test] | ||
222 | fn test_generate_enum_try_into_unit_variant() { | ||
223 | check_assist_not_applicable( | ||
224 | generate_enum_try_into_method, | ||
225 | r#"enum Value { | ||
226 | Number(i32), | ||
227 | Text(String), | ||
228 | Unit$0, | ||
229 | }"#, | ||
230 | ); | ||
231 | } | ||
232 | |||
233 | #[test] | ||
234 | fn test_generate_enum_try_into_record_with_multiple_fields() { | ||
235 | check_assist_not_applicable( | ||
236 | generate_enum_try_into_method, | ||
237 | r#"enum Value { | ||
238 | Number(i32), | ||
239 | Text(String), | ||
240 | Both { first: i32, second: String }$0, | ||
241 | }"#, | ||
242 | ); | ||
243 | } | ||
244 | |||
245 | #[test] | ||
246 | fn test_generate_enum_try_into_tuple_with_multiple_fields() { | ||
247 | check_assist_not_applicable( | ||
248 | generate_enum_try_into_method, | ||
249 | r#"enum Value { | ||
250 | Number(i32), | ||
251 | Text(String, String)$0, | ||
252 | }"#, | ||
253 | ); | ||
254 | } | ||
255 | |||
256 | #[test] | ||
257 | fn test_generate_enum_try_into_record_variant() { | ||
258 | check_assist( | ||
259 | generate_enum_try_into_method, | ||
260 | r#"enum Value { | ||
261 | Number(i32), | ||
262 | Text { text: String }$0, | ||
263 | }"#, | ||
264 | r#"enum Value { | ||
265 | Number(i32), | ||
266 | Text { text: String }, | ||
267 | } | ||
268 | |||
269 | impl Value { | ||
270 | fn try_into_text(self) -> Result<String, Self> { | ||
271 | if let Self::Text { text } = self { | ||
272 | Ok(text) | ||
273 | } else { | ||
274 | Err(self) | ||
275 | } | ||
276 | } | ||
277 | }"#, | ||
278 | ); | ||
279 | } | ||
280 | |||
281 | #[test] | ||
282 | fn test_generate_enum_as_tuple_variant() { | ||
283 | check_assist( | ||
284 | generate_enum_as_method, | ||
285 | r#" | ||
286 | enum Value { | ||
287 | Number(i32), | ||
288 | Text(String)$0, | ||
289 | }"#, | ||
290 | r#"enum Value { | ||
291 | Number(i32), | ||
292 | Text(String), | ||
293 | } | ||
294 | |||
295 | impl Value { | ||
296 | fn as_text(&self) -> Option<&String> { | ||
297 | if let Self::Text(v) = self { | ||
298 | Some(v) | ||
299 | } else { | ||
300 | None | ||
301 | } | ||
302 | } | ||
303 | }"#, | ||
304 | ); | ||
305 | } | ||
306 | |||
307 | #[test] | ||
308 | fn test_generate_enum_as_record_variant() { | ||
309 | check_assist( | ||
310 | generate_enum_as_method, | ||
311 | r#"enum Value { | ||
312 | Number(i32), | ||
313 | Text { text: String }$0, | ||
314 | }"#, | ||
315 | r#"enum Value { | ||
316 | Number(i32), | ||
317 | Text { text: String }, | ||
318 | } | ||
319 | |||
320 | impl Value { | ||
321 | fn as_text(&self) -> Option<&String> { | ||
322 | if let Self::Text { text } = self { | ||
323 | Some(text) | ||
324 | } else { | ||
325 | None | ||
326 | } | ||
327 | } | ||
328 | }"#, | ||
329 | ); | ||
330 | } | ||
331 | } | ||