make generate_enum_is_method work on any variants

Author: jDomantas

crates/ide_assists/src/handlers/generate_enum_match_method.rs

@@ -1,7 +1,6 @@
 use stdx::{format_to, to_lower_snake_case};
 use syntax::ast::VisibilityOwner;
 use syntax::ast::{self, AstNode, NameOwner};
-use test_utils::mark;
 
 use crate::{
     utils::{find_impl_block_end, find_struct_impl, generate_impl_text},
@@ -37,20 +36,17 @@ 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 = variant.parent_enum();
-    if !matches!(variant.kind(), ast::StructKind::Unit) {
-        mark::hit!(test_gen_enum_match_on_non_unit_variant_not_implemented);
-        return None;
-    }
+    let parent_enum = ast::Adt::Enum(variant.parent_enum());
+    let variant_kind = variant_kind(&variant);
 
     let enum_lowercase_name = to_lower_snake_case(&parent_enum.name()?.to_string());
-    let fn_name = to_lower_snake_case(&variant_name.to_string());
+    let fn_name = format!("is_{}", &to_lower_snake_case(variant_name.text()));
 
     // Return early if we've found an existing new fn
     let impl_def = find_struct_impl(
         &ctx,
-        &ast::Adt::Enum(parent_enum.clone()),
-        format!("is_{}", fn_name).as_str(),
+        &parent_enum,
+        &fn_name,
     )?;
 
     let target = variant.syntax().text_range();
@@ -69,20 +65,21 @@ pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext) ->
             format_to!(
                 buf,
                 "    /// Returns `true` if the {} is [`{}`].
-    {}fn is_{}(&self) -> bool {{
-        matches!(self, Self::{})
+    {}fn {}(&self) -> bool {{
+        matches!(self, Self::{}{})
     }}",
                 enum_lowercase_name,
                 variant_name,
                 vis,
                 fn_name,
-                variant_name
+                variant_name,
+                variant_kind.pattern_suffix(),
             );
 
             let start_offset = impl_def
                 .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
                 .unwrap_or_else(|| {
-                    buf = generate_impl_text(&ast::Adt::Enum(parent_enum.clone()), &buf);
+                    buf = generate_impl_text(&parent_enum, &buf);
                     parent_enum.syntax().text_range().end()
                 });
 
@@ -91,10 +88,53 @@ pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext) ->
     )
 }
 
+enum VariantKind {
+    Unit,
+    /// Tuple with a single field
+    NewtypeTuple,
+    /// Tuple with 0 or more than 2 fields
+    Tuple,
+    /// Record with a single field
+    NewtypeRecord { field_name: Option<ast::Name> },
+    /// Record with 0 or more than 2 fields
+    Record,
+}
+
+impl VariantKind {
+    fn pattern_suffix(&self) -> &'static str {
+        match self {
+            VariantKind::Unit => "",
+            VariantKind::NewtypeTuple |
+            VariantKind::Tuple => "(..)",
+            VariantKind::NewtypeRecord { .. } |
+            VariantKind::Record => " { .. }",
+        }
+    }
+}
+
+fn variant_kind(variant: &ast::Variant) -> VariantKind {
+    match variant.kind() {
+        ast::StructKind::Record(record) => {
+            if record.fields().count() == 1 {
+                let field_name = record.fields().nth(0).unwrap().name();
+                VariantKind::NewtypeRecord { field_name }
+            } else {
+                VariantKind::Record
+            }
+        }
+        ast::StructKind::Tuple(tuple) => {
+            if tuple.fields().count() == 1 {
+                VariantKind::NewtypeTuple
+            } else {
+                VariantKind::Tuple
+            }
+        }
+        ast::StructKind::Unit => VariantKind::Unit,
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use test_utils::mark;
-
     use crate::tests::{check_assist, check_assist_not_applicable};
 
     use super::*;
@@ -147,14 +187,51 @@ impl Variant {
     }
 
     #[test]
-    fn test_add_from_impl_no_element() {
-        mark::check!(test_gen_enum_match_on_non_unit_variant_not_implemented);
-        check_not_applicable(
+    fn test_generate_enum_match_from_tuple_variant() {
+        check_assist(
+            generate_enum_is_method,
             r#"
 enum Variant {
     Undefined,
     Minor(u32)$0,
     Major,
+}"#,
+            r#"enum Variant {
+    Undefined,
+    Minor(u32),
+    Major,
+}
+
+impl Variant {
+    /// Returns `true` if the variant is [`Minor`].
+    fn is_minor(&self) -> bool {
+        matches!(self, Self::Minor(..))
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_match_from_record_variant() {
+        check_assist(
+            generate_enum_is_method,
+            r#"
+enum Variant {
+    Undefined,
+    Minor { foo: i32 }$0,
+    Major,
+}"#,
+            r#"enum Variant {
+    Undefined,
+    Minor { foo: i32 },
+    Major,
+}
+
+impl Variant {
+    /// Returns `true` if the variant is [`Minor`].
+    fn is_minor(&self) -> bool {
+        matches!(self, Self::Minor { .. })
+    }
 }"#,
         );
     }