diff options
Diffstat (limited to 'crates/assists/src/handlers/generate_enum_match_method.rs')
-rw-r--r-- | crates/assists/src/handlers/generate_enum_match_method.rs | 293 |
1 files changed, 293 insertions, 0 deletions
diff --git a/crates/assists/src/handlers/generate_enum_match_method.rs b/crates/assists/src/handlers/generate_enum_match_method.rs new file mode 100644 index 000000000..079ed27bd --- /dev/null +++ b/crates/assists/src/handlers/generate_enum_match_method.rs | |||
@@ -0,0 +1,293 @@ | |||
1 | use hir::Adt; | ||
2 | use stdx::format_to; | ||
3 | use syntax::ast::{self, AstNode, NameOwner}; | ||
4 | use syntax::{ast::VisibilityOwner, T}; | ||
5 | use test_utils::mark; | ||
6 | |||
7 | use crate::{AssistContext, AssistId, AssistKind, Assists}; | ||
8 | |||
9 | // Assist: generate_enum_match_method | ||
10 | // | ||
11 | // Generate an `is_` method for an enum variant. | ||
12 | // | ||
13 | // ``` | ||
14 | // enum Version { | ||
15 | // Undefined, | ||
16 | // Minor$0, | ||
17 | // Major, | ||
18 | // } | ||
19 | // ``` | ||
20 | // -> | ||
21 | // ``` | ||
22 | // enum Version { | ||
23 | // Undefined, | ||
24 | // Minor, | ||
25 | // Major, | ||
26 | // } | ||
27 | // | ||
28 | // impl Version { | ||
29 | // fn is_minor(&self) -> bool { | ||
30 | // matches!(self, Self::Minor) | ||
31 | // } | ||
32 | // } | ||
33 | // ``` | ||
34 | pub(crate) fn generate_enum_match_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { | ||
35 | let variant = ctx.find_node_at_offset::<ast::Variant>()?; | ||
36 | let variant_name = variant.name()?; | ||
37 | let parent_enum = variant.parent_enum(); | ||
38 | if !matches!(variant.kind(), ast::StructKind::Unit) { | ||
39 | mark::hit!(test_gen_enum_match_on_non_unit_variant_not_implemented); | ||
40 | return None; | ||
41 | } | ||
42 | |||
43 | let fn_name = to_lower_snake_case(&format!("{}", variant_name)); | ||
44 | |||
45 | // Return early if we've found an existing new fn | ||
46 | let impl_def = find_struct_impl(&ctx, &parent_enum, format!("is_{}", fn_name).as_str())?; | ||
47 | |||
48 | let target = variant.syntax().text_range(); | ||
49 | acc.add( | ||
50 | AssistId("generate_enum_match_method", AssistKind::Generate), | ||
51 | "Generate an `is_` method for an enum variant", | ||
52 | target, | ||
53 | |builder| { | ||
54 | let mut buf = String::with_capacity(512); | ||
55 | |||
56 | if impl_def.is_some() { | ||
57 | buf.push('\n'); | ||
58 | } | ||
59 | |||
60 | let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v)); | ||
61 | |||
62 | format_to!( | ||
63 | buf, | ||
64 | " {}fn is_{}(&self) -> bool {{ | ||
65 | matches!(self, Self::{}) | ||
66 | }}", | ||
67 | vis, | ||
68 | fn_name, | ||
69 | variant_name | ||
70 | ); | ||
71 | |||
72 | let start_offset = impl_def | ||
73 | .and_then(|impl_def| { | ||
74 | buf.push('\n'); | ||
75 | let start = impl_def | ||
76 | .syntax() | ||
77 | .descendants_with_tokens() | ||
78 | .find(|t| t.kind() == T!['{'])? | ||
79 | .text_range() | ||
80 | .end(); | ||
81 | |||
82 | Some(start) | ||
83 | }) | ||
84 | .unwrap_or_else(|| { | ||
85 | buf = generate_impl_text(&parent_enum, &buf); | ||
86 | parent_enum.syntax().text_range().end() | ||
87 | }); | ||
88 | |||
89 | builder.insert(start_offset, buf); | ||
90 | }, | ||
91 | ) | ||
92 | } | ||
93 | |||
94 | // Generates the surrounding `impl Type { <code> }` including type and lifetime | ||
95 | // parameters | ||
96 | fn generate_impl_text(strukt: &ast::Enum, code: &str) -> String { | ||
97 | let mut buf = String::with_capacity(code.len()); | ||
98 | buf.push_str("\n\nimpl"); | ||
99 | buf.push_str(" "); | ||
100 | buf.push_str(strukt.name().unwrap().text()); | ||
101 | format_to!(buf, " {{\n{}\n}}", code); | ||
102 | buf | ||
103 | } | ||
104 | |||
105 | fn to_lower_snake_case(s: &str) -> String { | ||
106 | let mut buf = String::with_capacity(s.len()); | ||
107 | let mut prev = false; | ||
108 | for c in s.chars() { | ||
109 | if c.is_ascii_uppercase() && prev { | ||
110 | buf.push('_') | ||
111 | } | ||
112 | prev = true; | ||
113 | |||
114 | buf.push(c.to_ascii_lowercase()); | ||
115 | } | ||
116 | buf | ||
117 | } | ||
118 | |||
119 | // Uses a syntax-driven approach to find any impl blocks for the struct that | ||
120 | // exist within the module/file | ||
121 | // | ||
122 | // Returns `None` if we've found an existing `new` fn | ||
123 | // | ||
124 | // FIXME: change the new fn checking to a more semantic approach when that's more | ||
125 | // viable (e.g. we process proc macros, etc) | ||
126 | fn find_struct_impl( | ||
127 | ctx: &AssistContext, | ||
128 | strukt: &ast::Enum, | ||
129 | name: &str, | ||
130 | ) -> Option<Option<ast::Impl>> { | ||
131 | let db = ctx.db(); | ||
132 | let module = strukt.syntax().ancestors().find(|node| { | ||
133 | ast::Module::can_cast(node.kind()) || ast::SourceFile::can_cast(node.kind()) | ||
134 | })?; | ||
135 | |||
136 | let struct_def = ctx.sema.to_def(strukt)?; | ||
137 | |||
138 | let block = module.descendants().filter_map(ast::Impl::cast).find_map(|impl_blk| { | ||
139 | let blk = ctx.sema.to_def(&impl_blk)?; | ||
140 | |||
141 | // FIXME: handle e.g. `struct S<T>; impl<U> S<U> {}` | ||
142 | // (we currently use the wrong type parameter) | ||
143 | // also we wouldn't want to use e.g. `impl S<u32>` | ||
144 | let same_ty = match blk.target_ty(db).as_adt() { | ||
145 | Some(def) => def == Adt::Enum(struct_def), | ||
146 | None => false, | ||
147 | }; | ||
148 | let not_trait_impl = blk.target_trait(db).is_none(); | ||
149 | |||
150 | if !(same_ty && not_trait_impl) { | ||
151 | None | ||
152 | } else { | ||
153 | Some(impl_blk) | ||
154 | } | ||
155 | }); | ||
156 | |||
157 | if let Some(ref impl_blk) = block { | ||
158 | if has_fn(impl_blk, name) { | ||
159 | mark::hit!(test_gen_enum_match_impl_already_exists); | ||
160 | return None; | ||
161 | } | ||
162 | } | ||
163 | |||
164 | Some(block) | ||
165 | } | ||
166 | |||
167 | fn has_fn(imp: &ast::Impl, rhs_name: &str) -> bool { | ||
168 | if let Some(il) = imp.assoc_item_list() { | ||
169 | for item in il.assoc_items() { | ||
170 | if let ast::AssocItem::Fn(f) = item { | ||
171 | if let Some(name) = f.name() { | ||
172 | if name.text().eq_ignore_ascii_case(rhs_name) { | ||
173 | return true; | ||
174 | } | ||
175 | } | ||
176 | } | ||
177 | } | ||
178 | } | ||
179 | |||
180 | false | ||
181 | } | ||
182 | |||
183 | #[cfg(test)] | ||
184 | mod tests { | ||
185 | use ide_db::helpers::FamousDefs; | ||
186 | use test_utils::mark; | ||
187 | |||
188 | use crate::tests::{check_assist, check_assist_not_applicable}; | ||
189 | |||
190 | use super::*; | ||
191 | |||
192 | fn check_not_applicable(ra_fixture: &str) { | ||
193 | let fixture = | ||
194 | format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE); | ||
195 | check_assist_not_applicable(generate_enum_match_method, &fixture) | ||
196 | } | ||
197 | |||
198 | #[test] | ||
199 | fn test_generate_enum_match_from_variant() { | ||
200 | check_assist( | ||
201 | generate_enum_match_method, | ||
202 | r#" | ||
203 | enum Variant { | ||
204 | Undefined, | ||
205 | Minor$0, | ||
206 | Major, | ||
207 | }"#, | ||
208 | r#"enum Variant { | ||
209 | Undefined, | ||
210 | Minor, | ||
211 | Major, | ||
212 | } | ||
213 | |||
214 | impl Variant { | ||
215 | fn is_minor(&self) -> bool { | ||
216 | matches!(self, Self::Minor) | ||
217 | } | ||
218 | }"#, | ||
219 | ); | ||
220 | } | ||
221 | |||
222 | #[test] | ||
223 | fn test_generate_enum_match_already_implemented() { | ||
224 | mark::check!(test_gen_enum_match_impl_already_exists); | ||
225 | check_not_applicable( | ||
226 | r#" | ||
227 | enum Variant { | ||
228 | Undefined, | ||
229 | Minor$0, | ||
230 | Major, | ||
231 | } | ||
232 | |||
233 | impl Variant { | ||
234 | fn is_minor(&self) -> bool { | ||
235 | matches!(self, Self::Minor) | ||
236 | } | ||
237 | }"#, | ||
238 | ); | ||
239 | } | ||
240 | |||
241 | #[test] | ||
242 | fn test_add_from_impl_no_element() { | ||
243 | mark::check!(test_gen_enum_match_on_non_unit_variant_not_implemented); | ||
244 | check_not_applicable( | ||
245 | r#" | ||
246 | enum Variant { | ||
247 | Undefined, | ||
248 | Minor(u32)$0, | ||
249 | Major, | ||
250 | }"#, | ||
251 | ); | ||
252 | } | ||
253 | |||
254 | #[test] | ||
255 | fn test_generate_enum_match_from_variant_with_one_variant() { | ||
256 | check_assist( | ||
257 | generate_enum_match_method, | ||
258 | r#"enum Variant { Undefi$0ned }"#, | ||
259 | r#" | ||
260 | enum Variant { Undefined } | ||
261 | |||
262 | impl Variant { | ||
263 | fn is_undefined(&self) -> bool { | ||
264 | matches!(self, Self::Undefined) | ||
265 | } | ||
266 | }"#, | ||
267 | ); | ||
268 | } | ||
269 | |||
270 | #[test] | ||
271 | fn test_generate_enum_match_from_variant_with_visibility_marker() { | ||
272 | check_assist( | ||
273 | generate_enum_match_method, | ||
274 | r#" | ||
275 | pub(crate) enum Variant { | ||
276 | Undefined, | ||
277 | Minor$0, | ||
278 | Major, | ||
279 | }"#, | ||
280 | r#"pub(crate) enum Variant { | ||
281 | Undefined, | ||
282 | Minor, | ||
283 | Major, | ||
284 | } | ||
285 | |||
286 | impl Variant { | ||
287 | pub(crate) fn is_minor(&self) -> bool { | ||
288 | matches!(self, Self::Minor) | ||
289 | } | ||
290 | }"#, | ||
291 | ); | ||
292 | } | ||
293 | } | ||