Unnamed repository; edit this file 'description' to name the repository.
Diffstat (limited to 'crates/ide-assists/src/handlers/generate_enum_is_method.rs')
| -rw-r--r-- | crates/ide-assists/src/handlers/generate_enum_is_method.rs | 92 |
1 files changed, 77 insertions, 15 deletions
diff --git a/crates/ide-assists/src/handlers/generate_enum_is_method.rs b/crates/ide-assists/src/handlers/generate_enum_is_method.rs index 517906b429..b866022a7d 100644 --- a/crates/ide-assists/src/handlers/generate_enum_is_method.rs +++ b/crates/ide-assists/src/handlers/generate_enum_is_method.rs @@ -1,13 +1,12 @@ -use std::slice; - use ide_db::assists::GroupLabel; +use itertools::Itertools; use stdx::to_lower_snake_case; use syntax::ast::HasVisibility; use syntax::ast::{self, AstNode, HasName}; use crate::{ AssistContext, AssistId, Assists, - utils::{add_method_to_adt, find_struct_impl}, + utils::{add_method_to_adt, find_struct_impl, is_selected}, }; // Assist: generate_enum_is_method @@ -41,20 +40,21 @@ use crate::{ // ``` pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { let variant = ctx.find_node_at_offset::<ast::Variant>()?; - let variant_name = variant.name()?; let parent_enum = ast::Adt::Enum(variant.parent_enum()); - let pattern_suffix = match variant.kind() { - ast::StructKind::Record(_) => " { .. }", - ast::StructKind::Tuple(_) => "(..)", - ast::StructKind::Unit => "", - }; - + let variants = variant + .parent_enum() + .variant_list()? + .variants() + .filter(|it| is_selected(it, ctx.selection_trimmed(), true)) + .collect::<Vec<_>>(); + let methods = variants.iter().map(Method::new).collect::<Option<Vec<_>>>()?; let enum_name = parent_enum.name()?; let enum_lowercase_name = to_lower_snake_case(&enum_name.to_string()).replace('_', " "); - let fn_name = format!("is_{}", &to_lower_snake_case(&variant_name.text())); + let fn_names = methods.iter().map(|it| it.fn_name.clone()).collect::<Vec<_>>(); + stdx::never!(variants.is_empty()); // Return early if we've found an existing new fn - let impl_def = find_struct_impl(ctx, &parent_enum, slice::from_ref(&fn_name))?; + let impl_def = find_struct_impl(ctx, &parent_enum, &fn_names)?; let target = variant.syntax().text_range(); acc.add_group( @@ -64,21 +64,47 @@ pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext<'_> target, |builder| { let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{v} ")); - let method = format!( - " /// Returns `true` if the {enum_lowercase_name} is [`{variant_name}`]. + let method = methods + .iter() + .map(|Method { pattern_suffix, fn_name, variant_name }| { + format!( + " \ + /// Returns `true` if the {enum_lowercase_name} is [`{variant_name}`]. /// /// [`{variant_name}`]: {enum_name}::{variant_name} #[must_use] {vis}fn {fn_name}(&self) -> bool {{ matches!(self, Self::{variant_name}{pattern_suffix}) }}", - ); + ) + }) + .join("\n\n"); add_method_to_adt(builder, &parent_enum, impl_def, &method); }, ) } +struct Method { + pattern_suffix: &'static str, + fn_name: String, + variant_name: ast::Name, +} + +impl Method { + fn new(variant: &ast::Variant) -> Option<Self> { + let pattern_suffix = match variant.kind() { + ast::StructKind::Record(_) => " { .. }", + ast::StructKind::Tuple(_) => "(..)", + ast::StructKind::Unit => "", + }; + + let variant_name = variant.name()?; + let fn_name = format!("is_{}", &to_lower_snake_case(&variant_name.text())); + Some(Method { pattern_suffix, fn_name, variant_name }) + } +} + #[cfg(test)] mod tests { use crate::tests::{check_assist, check_assist_not_applicable}; @@ -114,6 +140,42 @@ impl Variant { } #[test] + fn test_generate_enum_is_from_multiple_variant() { + check_assist( + generate_enum_is_method, + r#" +enum Variant { + Undefined, + $0Minor, + M$0ajor, +}"#, + r#"enum Variant { + Undefined, + Minor, + Major, +} + +impl Variant { + /// Returns `true` if the variant is [`Minor`]. + /// + /// [`Minor`]: Variant::Minor + #[must_use] + fn is_minor(&self) -> bool { + matches!(self, Self::Minor) + } + + /// Returns `true` if the variant is [`Major`]. + /// + /// [`Major`]: Variant::Major + #[must_use] + fn is_major(&self) -> bool { + matches!(self, Self::Major) + } +}"#, + ); + } + + #[test] fn test_generate_enum_is_already_implemented() { check_assist_not_applicable( generate_enum_is_method, |