Merge #7677

7677: More enum matching r=yoshuawuyts a=jDomantas

* Renamed existing `generate_enum_match_method` to `generate_enum_is_variant`
* Added two similar assists to generate `into_` and `as_` methods.
* Made all of them general enough to work on record and tuple variants too.

For `as_` method generation there's room to improve:
* Right now it always returns `Option<&Field>`, even though `Option<Field>` would be nicer when `Field: Copy`. I don't know how to check if the field type implements `Copy`. If given suggestions I could try to fix this in a follow-up pr.
* `&String` could be replaced with `&str`, `&Box<_>` with `&_`, and probably some more. I don't know what would be a good way to do that.

Closes #7604

Co-authored-by: Domantas Jadenkus <[email protected]>

5 files changed, 488 insertions, 64 deletions

diff --git a/crates/ide_assists/src/handlers/generate_enum_match_method.rs b/crates/ide_assists/src/handlers/generate_enum_is_method.rs
index aeb887e71..7e181a480 100644
--- a/crates/ide_assists/src/handlers/generate_enum_match_method.rs
+++ b/crates/ide_assists/src/handlers/generate_enum_is_method.rs
@@ -1,14 +1,13 @@
-use stdx::{format_to, to_lower_snake_case};
+use stdx::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},
+    utils::{add_method_to_adt, find_struct_impl},
     AssistContext, AssistId, AssistKind, Assists,
 };
 
-// Assist: generate_enum_match_method
+// Assist: generate_enum_is_method
 //
 // Generate an `is_` method for an enum variant.
 //
@@ -34,79 +33,52 @@ use crate::{
 //     }
 // }
 // ```
-pub(crate) fn generate_enum_match_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+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 pattern_suffix = match variant.kind() {
+        ast::StructKind::Record(_) => " { .. }",
+        ast::StructKind::Tuple(_) => "(..)",
+        ast::StructKind::Unit => "",
+    };
 
     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(),
-    )?;
+    let impl_def = find_struct_impl(&ctx, &parent_enum, &fn_name)?;
 
     let target = variant.syntax().text_range();
     acc.add(
-        AssistId("generate_enum_match_method", AssistKind::Generate),
+        AssistId("generate_enum_is_method", AssistKind::Generate),
         "Generate an `is_` method for an enum variant",
         target,
         |builder| {
-            let mut buf = String::with_capacity(512);
-
-            if impl_def.is_some() {
-                buf.push('\n');
-            }
-
             let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
-            format_to!(
-                buf,
+            let method = format!(
                 "    /// 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
+                enum_lowercase_name, variant_name, vis, fn_name, variant_name, 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);
-                    parent_enum.syntax().text_range().end()
-                });
-
-            builder.insert(start_offset, buf);
+            add_method_to_adt(builder, &parent_enum, impl_def, &method);
         },
     )
 }
 
 #[cfg(test)]
 mod tests {
-    use test_utils::mark;
-
     use crate::tests::{check_assist, check_assist_not_applicable};
 
     use super::*;
 
-    fn check_not_applicable(ra_fixture: &str) {
-        check_assist_not_applicable(generate_enum_match_method, ra_fixture)
-    }
-
     #[test]
-    fn test_generate_enum_match_from_variant() {
+    fn test_generate_enum_is_from_variant() {
         check_assist(
-            generate_enum_match_method,
+            generate_enum_is_method,
             r#"
 enum Variant {
     Undefined,
@@ -129,8 +101,9 @@ impl Variant {
     }
 
     #[test]
-    fn test_generate_enum_match_already_implemented() {
-        check_not_applicable(
+    fn test_generate_enum_is_already_implemented() {
+        check_assist_not_applicable(
+            generate_enum_is_method,
             r#"
 enum Variant {
     Undefined,
@@ -147,22 +120,59 @@ 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_is_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_is_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 { .. })
+    }
+}"#,
         );
     }
 
     #[test]
-    fn test_generate_enum_match_from_variant_with_one_variant() {
+    fn test_generate_enum_is_from_variant_with_one_variant() {
         check_assist(
-            generate_enum_match_method,
+            generate_enum_is_method,
             r#"enum Variant { Undefi$0ned }"#,
             r#"
 enum Variant { Undefined }
@@ -177,9 +187,9 @@ impl Variant {
     }
 
     #[test]
-    fn test_generate_enum_match_from_variant_with_visibility_marker() {
+    fn test_generate_enum_is_from_variant_with_visibility_marker() {
         check_assist(
-            generate_enum_match_method,
+            generate_enum_is_method,
             r#"
 pub(crate) enum Variant {
     Undefined,
@@ -202,9 +212,9 @@ impl Variant {
     }
 
     #[test]
-    fn test_multiple_generate_enum_match_from_variant() {
+    fn test_multiple_generate_enum_is_from_variant() {
         check_assist(
-            generate_enum_match_method,
+            generate_enum_is_method,
             r#"
 enum Variant {
     Undefined,
diff --git a/crates/ide_assists/src/handlers/generate_enum_projection_method.rs b/crates/ide_assists/src/handlers/generate_enum_projection_method.rs
new file mode 100644
index 000000000..871bcab50
--- /dev/null
+++ b/crates/ide_assists/src/handlers/generate_enum_projection_method.rs
@@ -0,0 +1,331 @@
+use itertools::Itertools;
+use stdx::to_lower_snake_case;
+use syntax::ast::VisibilityOwner;
+use syntax::ast::{self, AstNode, NameOwner};
+
+use crate::{
+    utils::{add_method_to_adt, find_struct_impl},
+    AssistContext, AssistId, AssistKind, Assists,
+};
+
+// Assist: generate_enum_try_into_method
+//
+// Generate an `try_into_` method for an enum variant.
+//
+// ```
+// enum Value {
+//  Number(i32),
+//  Text(String)$0,
+// }
+// ```
+// ->
+// ```
+// enum Value {
+//  Number(i32),
+//  Text(String),
+// }
+//
+// impl Value {
+//     fn try_into_text(self) -> Result<String, Self> {
+//         if let Self::Text(v) = self {
+//             Ok(v)
+//         } else {
+//             Err(self)
+//         }
+//     }
+// }
+// ```
+pub(crate) fn generate_enum_try_into_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    generate_enum_projection_method(
+        acc,
+        ctx,
+        "generate_enum_try_into_method",
+        "Generate an `try_into_` method for an enum variant",
+        ProjectionProps {
+            fn_name_prefix: "try_into",
+            self_param: "self",
+            return_prefix: "Result<",
+            return_suffix: ", Self>",
+            happy_case: "Ok",
+            sad_case: "Err(self)",
+        },
+    )
+}
+
+// Assist: generate_enum_as_method
+//
+// Generate an `as_` method for an enum variant.
+//
+// ```
+// enum Value {
+//  Number(i32),
+//  Text(String)$0,
+// }
+// ```
+// ->
+// ```
+// enum Value {
+//  Number(i32),
+//  Text(String),
+// }
+//
+// impl Value {
+//     fn as_text(&self) -> Option<&String> {
+//         if let Self::Text(v) = self {
+//             Some(v)
+//         } else {
+//             None
+//         }
+//     }
+// }
+// ```
+pub(crate) fn generate_enum_as_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    generate_enum_projection_method(
+        acc,
+        ctx,
+        "generate_enum_as_method",
+        "Generate an `as_` method for an enum variant",
+        ProjectionProps {
+            fn_name_prefix: "as",
+            self_param: "&self",
+            return_prefix: "Option<&",
+            return_suffix: ">",
+            happy_case: "Some",
+            sad_case: "None",
+        },
+    )
+}
+
+struct ProjectionProps {
+    fn_name_prefix: &'static str,
+    self_param: &'static str,
+    return_prefix: &'static str,
+    return_suffix: &'static str,
+    happy_case: &'static str,
+    sad_case: &'static str,
+}
+
+fn generate_enum_projection_method(
+    acc: &mut Assists,
+    ctx: &AssistContext,
+    assist_id: &'static str,
+    assist_description: &str,
+    props: ProjectionProps,
+) -> 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, field_type, bound_name) = match variant.kind() {
+        ast::StructKind::Record(record) => {
+            let (field,) = record.fields().collect_tuple()?;
+            let name = field.name()?.to_string();
+            let ty = field.ty()?;
+            let pattern_suffix = format!(" {{ {} }}", name);
+            (pattern_suffix, ty, name)
+        }
+        ast::StructKind::Tuple(tuple) => {
+            let (field,) = tuple.fields().collect_tuple()?;
+            let ty = field.ty()?;
+            ("(v)".to_owned(), ty, "v".to_owned())
+        }
+        ast::StructKind::Unit => return None,
+    };
+
+    let fn_name = format!("{}_{}", props.fn_name_prefix, &to_lower_snake_case(variant_name.text()));
+
+    // Return early if we've found an existing new fn
+    let impl_def = find_struct_impl(&ctx, &parent_enum, &fn_name)?;
+
+    let target = variant.syntax().text_range();
+    acc.add(AssistId(assist_id, AssistKind::Generate), assist_description, target, |builder| {
+        let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
+        let method = format!(
+            "    {0}fn {1}({2}) -> {3}{4}{5} {{
+        if let Self::{6}{7} = self {{
+            {8}({9})
+        }} else {{
+            {10}
+        }}
+    }}",
+            vis,
+            fn_name,
+            props.self_param,
+            props.return_prefix,
+            field_type.syntax(),
+            props.return_suffix,
+            variant_name,
+            pattern_suffix,
+            props.happy_case,
+            bound_name,
+            props.sad_case,
+        );
+
+        add_method_to_adt(builder, &parent_enum, impl_def, &method);
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    #[test]
+    fn test_generate_enum_try_into_tuple_variant() {
+        check_assist(
+            generate_enum_try_into_method,
+            r#"
+enum Value {
+    Number(i32),
+    Text(String)$0,
+}"#,
+            r#"enum Value {
+    Number(i32),
+    Text(String),
+}
+
+impl Value {
+    fn try_into_text(self) -> Result<String, Self> {
+        if let Self::Text(v) = self {
+            Ok(v)
+        } else {
+            Err(self)
+        }
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_try_into_already_implemented() {
+        check_assist_not_applicable(
+            generate_enum_try_into_method,
+            r#"enum Value {
+    Number(i32),
+    Text(String)$0,
+}
+
+impl Value {
+    fn try_into_text(self) -> Result<String, Self> {
+        if let Self::Text(v) = self {
+            Ok(v)
+        } else {
+            Err(self)
+        }
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_try_into_unit_variant() {
+        check_assist_not_applicable(
+            generate_enum_try_into_method,
+            r#"enum Value {
+    Number(i32),
+    Text(String),
+    Unit$0,
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_try_into_record_with_multiple_fields() {
+        check_assist_not_applicable(
+            generate_enum_try_into_method,
+            r#"enum Value {
+    Number(i32),
+    Text(String),
+    Both { first: i32, second: String }$0,
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_try_into_tuple_with_multiple_fields() {
+        check_assist_not_applicable(
+            generate_enum_try_into_method,
+            r#"enum Value {
+    Number(i32),
+    Text(String, String)$0,
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_try_into_record_variant() {
+        check_assist(
+            generate_enum_try_into_method,
+            r#"enum Value {
+    Number(i32),
+    Text { text: String }$0,
+}"#,
+            r#"enum Value {
+    Number(i32),
+    Text { text: String },
+}
+
+impl Value {
+    fn try_into_text(self) -> Result<String, Self> {
+        if let Self::Text { text } = self {
+            Ok(text)
+        } else {
+            Err(self)
+        }
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_as_tuple_variant() {
+        check_assist(
+            generate_enum_as_method,
+            r#"
+enum Value {
+    Number(i32),
+    Text(String)$0,
+}"#,
+            r#"enum Value {
+    Number(i32),
+    Text(String),
+}
+
+impl Value {
+    fn as_text(&self) -> Option<&String> {
+        if let Self::Text(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_as_record_variant() {
+        check_assist(
+            generate_enum_as_method,
+            r#"enum Value {
+    Number(i32),
+    Text { text: String }$0,
+}"#,
+            r#"enum Value {
+    Number(i32),
+    Text { text: String },
+}
+
+impl Value {
+    fn as_text(&self) -> Option<&String> {
+        if let Self::Text { text } = self {
+            Some(text)
+        } else {
+            None
+        }
+    }
+}"#,
+        );
+    }
+}
diff --git a/crates/ide_assists/src/lib.rs b/crates/ide_assists/src/lib.rs
index f4c7e6fbf..4c067d451 100644
--- a/crates/ide_assists/src/lib.rs
+++ b/crates/ide_assists/src/lib.rs
@@ -128,7 +128,8 @@ mod handlers {
     mod flip_trait_bound;
     mod generate_default_from_enum_variant;
     mod generate_derive;
-    mod generate_enum_match_method;
+    mod generate_enum_is_method;
+    mod generate_enum_projection_method;
     mod generate_from_impl_for_enum;
     mod generate_function;
     mod generate_getter;
@@ -189,7 +190,9 @@ mod handlers {
             flip_trait_bound::flip_trait_bound,
             generate_default_from_enum_variant::generate_default_from_enum_variant,
             generate_derive::generate_derive,
-            generate_enum_match_method::generate_enum_match_method,
+            generate_enum_is_method::generate_enum_is_method,
+            generate_enum_projection_method::generate_enum_try_into_method,
+            generate_enum_projection_method::generate_enum_as_method,
             generate_from_impl_for_enum::generate_from_impl_for_enum,
             generate_function::generate_function,
             generate_getter::generate_getter,
diff --git a/crates/ide_assists/src/tests/generated.rs b/crates/ide_assists/src/tests/generated.rs
index d42875822..7f6dbbccf 100644
--- a/crates/ide_assists/src/tests/generated.rs
+++ b/crates/ide_assists/src/tests/generated.rs
@@ -483,9 +483,38 @@ struct Point {
 }
 
 #[test]
-fn doctest_generate_enum_match_method() {
+fn doctest_generate_enum_as_method() {
     check_doc_test(
-        "generate_enum_match_method",
+        "generate_enum_as_method",
+        r#####"
+enum Value {
+ Number(i32),
+ Text(String)$0,
+}
+"#####,
+        r#####"
+enum Value {
+ Number(i32),
+ Text(String),
+}
+
+impl Value {
+    fn as_text(&self) -> Option<&String> {
+        if let Self::Text(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+}
+"#####,
+    )
+}
+
+#[test]
+fn doctest_generate_enum_is_method() {
+    check_doc_test(
+        "generate_enum_is_method",
         r#####"
 enum Version {
  Undefined,
@@ -511,6 +540,35 @@ impl Version {
 }
 
 #[test]
+fn doctest_generate_enum_try_into_method() {
+    check_doc_test(
+        "generate_enum_try_into_method",
+        r#####"
+enum Value {
+ Number(i32),
+ Text(String)$0,
+}
+"#####,
+        r#####"
+enum Value {
+ Number(i32),
+ Text(String),
+}
+
+impl Value {
+    fn try_into_text(self) -> Result<String, Self> {
+        if let Self::Text(v) = self {
+            Ok(v)
+        } else {
+            Err(self)
+        }
+    }
+}
+"#####,
+    )
+}
+
+#[test]
 fn doctest_generate_from_impl_for_enum() {
     check_doc_test(
         "generate_from_impl_for_enum",
diff --git a/crates/ide_assists/src/utils.rs b/crates/ide_assists/src/utils.rs
index 276792bc1..880ab6fe3 100644
--- a/crates/ide_assists/src/utils.rs
+++ b/crates/ide_assists/src/utils.rs
@@ -21,7 +21,7 @@ use syntax::{
 };
 
 use crate::{
-    assist_context::AssistContext,
+    assist_context::{AssistBuilder, AssistContext},
     ast_transform::{self, AstTransform, QualifyPaths, SubstituteTypeParams},
 };
 
@@ -464,3 +464,25 @@ fn generate_impl_text_inner(adt: &ast::Adt, trait_text: Option<&str>, code: &str
 
     buf
 }
+
+pub(crate) fn add_method_to_adt(
+    builder: &mut AssistBuilder,
+    adt: &ast::Adt,
+    impl_def: Option<ast::Impl>,
+    method: &str,
+) {
+    let mut buf = String::with_capacity(method.len() + 2);
+    if impl_def.is_some() {
+        buf.push('\n');
+    }
+    buf.push_str(method);
+
+    let start_offset = impl_def
+        .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
+        .unwrap_or_else(|| {
+            buf = generate_impl_text(&adt, &buf);
+            adt.syntax().text_range().end()
+        });
+
+    builder.insert(start_offset, buf);
+}