ide_assists/handlers/
generate_enum_is_method.rs

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

ide_assists/handlers/
generate_enum_is_method.rs
