handle generic enums

1 files changed, 58 insertions, 8 deletions

diff --git a/crates/assists/src/handlers/generate_from_impl_for_enum.rs b/crates/assists/src/handlers/generate_from_impl_for_enum.rs
index ac9806368..8e68ebc7a 100644
--- a/crates/assists/src/handlers/generate_from_impl_for_enum.rs
+++ b/crates/assists/src/handlers/generate_from_impl_for_enum.rs
@@ -1,6 +1,9 @@
+use ast::GenericParamsOwner;
 use ide_db::helpers::FamousDefs;
 use ide_db::RootDatabase;
-use syntax::ast::{self, AstNode, NameOwner};
+use itertools::Itertools;
+use stdx::format_to;
+use syntax::{SmolStr, ast::{self, AstNode, NameOwner}};
 use test_utils::mark;
 
 use crate::{AssistContext, AssistId, AssistKind, Assists};
@@ -26,6 +29,7 @@ pub(crate) fn generate_from_impl_for_enum(acc: &mut Assists, ctx: &AssistContext
     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
     let variant_name = variant.name()?;
     let enum_name = variant.parent_enum().name()?;
+    let enum_type_params = variant.parent_enum().generic_param_list();
     let field_list = match variant.kind() {
         ast::StructKind::Tuple(field_list) => field_list,
         _ => return None,
@@ -47,17 +51,33 @@ pub(crate) fn generate_from_impl_for_enum(acc: &mut Assists, ctx: &AssistContext
         target,
         |edit| {
             let start_offset = variant.parent_enum().syntax().text_range().end();
-            let buf = format!(
-                r#"
+            let mut buf = String::from("\n\nimpl");
+            if let Some(type_params) = &enum_type_params {
+                format_to!(buf, "{}", type_params.syntax());
+            }
+            format_to!(buf, " From<{}> for {}", field_type.syntax(), enum_name);
+            if let Some(type_params) = enum_type_params {
+                let lifetime_params = type_params
+                    .lifetime_params()
+                    .filter_map(|it| it.lifetime())
+                    .map(|it| SmolStr::from(it.text()));
+                let type_params = type_params
+                    .type_params()
+                    .filter_map(|it| it.name())
+                    .map(|it| SmolStr::from(it.text()));
 
-impl From<{0}> for {1} {{
-    fn from(v: {0}) -> Self {{
-        Self::{2}(v)
+                let generic_params = lifetime_params.chain(type_params).format(", ");
+                format_to!(buf, "<{}>", generic_params)
+            }
+            format_to!(
+                buf,
+                r#" {{
+    fn from(v: {}) -> Self {{
+        Self::{}(v)
     }}
 }}"#,
                 field_type.syntax(),
-                enum_name,
-                variant_name
+                variant_name,
             );
             edit.insert(start_offset, buf);
         },
@@ -209,4 +229,34 @@ impl From<&'static str> for A {
 }"#,
         );
     }
+
+    #[test]
+    fn test_add_from_impl_generic_enum() {
+        check_assist(
+            generate_from_impl_for_enum,
+            "enum Generic<T, U: Clone> { $0One(T), Two(U) }",
+            r#"enum Generic<T, U: Clone> { One(T), Two(U) }
+
+impl<T, U: Clone> From<T> for Generic<T, U> {
+    fn from(v: T) -> Self {
+        Self::One(v)
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_add_from_impl_with_lifetime() {
+        check_assist(
+            generate_from_impl_for_enum,
+            "enum Generic<'a> { $0One(&'a i32) }",
+            r#"enum Generic<'a> { One(&'a i32) }
+
+impl<'a> From<&'a i32> for Generic<'a> {
+    fn from(v: &'a i32) -> Self {
+        Self::One(v)
+    }
+}"#,
+        );
+    }
 }