ide_assists/handlers/
generate_enum_is_method.rs

1use ide_db::assists::GroupLabel;
2use itertools::Itertools;
3use stdx::to_lower_snake_case;
4use syntax::ast::HasVisibility;
5use syntax::ast::{self, AstNode, HasName};
6
7use crate::{
8    AssistContext, AssistId, Assists,
9    utils::{add_method_to_adt, find_struct_impl, is_selected},
10};
11
12// Assist: generate_enum_is_method
13//
14// Generate an `is_` method for this enum variant.
15//
16// ```
17// enum Version {
18//  Undefined,
19//  Minor$0,
20//  Major,
21// }
22// ```
23// ->
24// ```
25// enum Version {
26//  Undefined,
27//  Minor,
28//  Major,
29// }
30//
31// impl Version {
32//     /// Returns `true` if the version is [`Minor`].
33//     ///
34//     /// [`Minor`]: Version::Minor
35//     #[must_use]
36//     fn is_minor(&self) -> bool {
37//         matches!(self, Self::Minor)
38//     }
39// }
40// ```
41pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
42    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
43    let parent_enum = ast::Adt::Enum(variant.parent_enum());
44    let variants = variant
45        .parent_enum()
46        .variant_list()?
47        .variants()
48        .filter(|it| is_selected(it, ctx.selection_trimmed(), true))
49        .collect::<Vec<_>>();
50    let methods = variants.iter().map(Method::new).collect::<Option<Vec<_>>>()?;
51    let enum_name = parent_enum.name()?;
52    let enum_lowercase_name = to_lower_snake_case(&enum_name.to_string()).replace('_', " ");
53    let fn_names = methods.iter().map(|it| it.fn_name.clone()).collect::<Vec<_>>();
54    stdx::never!(variants.is_empty());
55
56    // Return early if we've found an existing new fn
57    let impl_def = find_struct_impl(ctx, &parent_enum, &fn_names)?;
58
59    let target = variant.syntax().text_range();
60    acc.add_group(
61        &GroupLabel("Generate an `is_`,`as_`, or `try_into_` for this enum variant".to_owned()),
62        AssistId::generate("generate_enum_is_method"),
63        "Generate an `is_` method for this enum variant",
64        target,
65        |builder| {
66            let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{v} "));
67            let method = methods
68                .iter()
69                .map(|Method { pattern_suffix, fn_name, variant_name }| {
70                    format!(
71                        "    \
72    /// Returns `true` if the {enum_lowercase_name} is [`{variant_name}`].
73    ///
74    /// [`{variant_name}`]: {enum_name}::{variant_name}
75    #[must_use]
76    {vis}fn {fn_name}(&self) -> bool {{
77        matches!(self, Self::{variant_name}{pattern_suffix})
78    }}",
79                    )
80                })
81                .join("\n\n");
82
83            add_method_to_adt(builder, &parent_enum, impl_def, &method);
84        },
85    )
86}
87
88struct Method {
89    pattern_suffix: &'static str,
90    fn_name: String,
91    variant_name: ast::Name,
92}
93
94impl Method {
95    fn new(variant: &ast::Variant) -> Option<Self> {
96        let pattern_suffix = match variant.kind() {
97            ast::StructKind::Record(_) => " { .. }",
98            ast::StructKind::Tuple(_) => "(..)",
99            ast::StructKind::Unit => "",
100        };
101
102        let variant_name = variant.name()?;
103        let fn_name = format!("is_{}", &to_lower_snake_case(&variant_name.text()));
104        Some(Method { pattern_suffix, fn_name, variant_name })
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use crate::tests::{check_assist, check_assist_not_applicable};
111
112    use super::*;
113
114    #[test]
115    fn test_generate_enum_is_from_variant() {
116        check_assist(
117            generate_enum_is_method,
118            r#"
119enum Variant {
120    Undefined,
121    Minor$0,
122    Major,
123}"#,
124            r#"enum Variant {
125    Undefined,
126    Minor,
127    Major,
128}
129
130impl Variant {
131    /// Returns `true` if the variant is [`Minor`].
132    ///
133    /// [`Minor`]: Variant::Minor
134    #[must_use]
135    fn is_minor(&self) -> bool {
136        matches!(self, Self::Minor)
137    }
138}"#,
139        );
140    }
141
142    #[test]
143    fn test_generate_enum_is_from_multiple_variant() {
144        check_assist(
145            generate_enum_is_method,
146            r#"
147enum Variant {
148    Undefined,
149    $0Minor,
150    M$0ajor,
151}"#,
152            r#"enum Variant {
153    Undefined,
154    Minor,
155    Major,
156}
157
158impl Variant {
159    /// Returns `true` if the variant is [`Minor`].
160    ///
161    /// [`Minor`]: Variant::Minor
162    #[must_use]
163    fn is_minor(&self) -> bool {
164        matches!(self, Self::Minor)
165    }
166
167    /// Returns `true` if the variant is [`Major`].
168    ///
169    /// [`Major`]: Variant::Major
170    #[must_use]
171    fn is_major(&self) -> bool {
172        matches!(self, Self::Major)
173    }
174}"#,
175        );
176    }
177
178    #[test]
179    fn test_generate_enum_is_already_implemented() {
180        check_assist_not_applicable(
181            generate_enum_is_method,
182            r#"
183enum Variant {
184    Undefined,
185    Minor$0,
186    Major,
187}
188
189impl Variant {
190    fn is_minor(&self) -> bool {
191        matches!(self, Self::Minor)
192    }
193}"#,
194        );
195    }
196
197    #[test]
198    fn test_generate_enum_is_from_tuple_variant() {
199        check_assist(
200            generate_enum_is_method,
201            r#"
202enum Variant {
203    Undefined,
204    Minor(u32)$0,
205    Major,
206}"#,
207            r#"enum Variant {
208    Undefined,
209    Minor(u32),
210    Major,
211}
212
213impl Variant {
214    /// Returns `true` if the variant is [`Minor`].
215    ///
216    /// [`Minor`]: Variant::Minor
217    #[must_use]
218    fn is_minor(&self) -> bool {
219        matches!(self, Self::Minor(..))
220    }
221}"#,
222        );
223    }
224
225    #[test]
226    fn test_generate_enum_is_from_record_variant() {
227        check_assist(
228            generate_enum_is_method,
229            r#"
230enum Variant {
231    Undefined,
232    Minor { foo: i32 }$0,
233    Major,
234}"#,
235            r#"enum Variant {
236    Undefined,
237    Minor { foo: i32 },
238    Major,
239}
240
241impl Variant {
242    /// Returns `true` if the variant is [`Minor`].
243    ///
244    /// [`Minor`]: Variant::Minor
245    #[must_use]
246    fn is_minor(&self) -> bool {
247        matches!(self, Self::Minor { .. })
248    }
249}"#,
250        );
251    }
252
253    #[test]
254    fn test_generate_enum_is_from_variant_with_one_variant() {
255        check_assist(
256            generate_enum_is_method,
257            r#"enum Variant { Undefi$0ned }"#,
258            r#"
259enum Variant { Undefined }
260
261impl Variant {
262    /// Returns `true` if the variant is [`Undefined`].
263    ///
264    /// [`Undefined`]: Variant::Undefined
265    #[must_use]
266    fn is_undefined(&self) -> bool {
267        matches!(self, Self::Undefined)
268    }
269}"#,
270        );
271    }
272
273    #[test]
274    fn test_generate_enum_is_from_variant_with_visibility_marker() {
275        check_assist(
276            generate_enum_is_method,
277            r#"
278pub(crate) enum Variant {
279    Undefined,
280    Minor$0,
281    Major,
282}"#,
283            r#"pub(crate) enum Variant {
284    Undefined,
285    Minor,
286    Major,
287}
288
289impl Variant {
290    /// Returns `true` if the variant is [`Minor`].
291    ///
292    /// [`Minor`]: Variant::Minor
293    #[must_use]
294    pub(crate) fn is_minor(&self) -> bool {
295        matches!(self, Self::Minor)
296    }
297}"#,
298        );
299    }
300
301    #[test]
302    fn test_multiple_generate_enum_is_from_variant() {
303        check_assist(
304            generate_enum_is_method,
305            r#"
306enum Variant {
307    Undefined,
308    Minor,
309    Major$0,
310}
311
312impl Variant {
313    /// Returns `true` if the variant is [`Minor`].
314    ///
315    /// [`Minor`]: Variant::Minor
316    #[must_use]
317    fn is_minor(&self) -> bool {
318        matches!(self, Self::Minor)
319    }
320}"#,
321            r#"enum Variant {
322    Undefined,
323    Minor,
324    Major,
325}
326
327impl Variant {
328    /// Returns `true` if the variant is [`Minor`].
329    ///
330    /// [`Minor`]: Variant::Minor
331    #[must_use]
332    fn is_minor(&self) -> bool {
333        matches!(self, Self::Minor)
334    }
335
336    /// Returns `true` if the variant is [`Major`].
337    ///
338    /// [`Major`]: Variant::Major
339    #[must_use]
340    fn is_major(&self) -> bool {
341        matches!(self, Self::Major)
342    }
343}"#,
344        );
345    }
346
347    #[test]
348    fn test_generate_enum_is_variant_names() {
349        check_assist(
350            generate_enum_is_method,
351            r#"
352enum CoroutineState {
353    Yielded,
354    Complete$0,
355    Major,
356}"#,
357            r#"enum CoroutineState {
358    Yielded,
359    Complete,
360    Major,
361}
362
363impl CoroutineState {
364    /// Returns `true` if the coroutine state is [`Complete`].
365    ///
366    /// [`Complete`]: CoroutineState::Complete
367    #[must_use]
368    fn is_complete(&self) -> bool {
369        matches!(self, Self::Complete)
370    }
371}"#,
372        );
373    }
374}
ide_assists/handlers/generate_enum_is_method.rs

ide_assists/handlers/
generate_enum_is_method.rs
