generate default implementation for an enum from an enum variant #6860

Signed-off-by: Benjamin Coenen <[email protected]>

3 files changed, 200 insertions, 0 deletions

diff --git a/crates/assists/src/handlers/generate_default_from_enum_variant.rs b/crates/assists/src/handlers/generate_default_from_enum_variant.rs
new file mode 100644
index 000000000..6521cb543
--- /dev/null
+++ b/crates/assists/src/handlers/generate_default_from_enum_variant.rs
@@ -0,0 +1,171 @@
+use ide_db::helpers::FamousDefs;
+use ide_db::RootDatabase;
+use syntax::ast::{self, AstNode, NameOwner};
+use test_utils::mark;
+
+use crate::{AssistContext, AssistId, AssistKind, Assists};
+
+// Assist: generate_default_from_enum_variant
+//
+// Adds a Default impl for an enum using a variant.
+//
+// ```
+// enum Version {
+//  Undefined,
+//  Minor<|>,
+//  Major,
+// }
+// ```
+// ->
+// ```
+// enum Version {
+//  Undefined,
+//  Minor,
+//  Major,
+// }
+//
+// impl Default for Version {
+//     fn default() -> Self {
+//         Self::Minor
+//     }
+// }
+// ```
+pub(crate) fn generate_default_from_enum_variant(
+    acc: &mut Assists,
+    ctx: &AssistContext,
+) -> Option<()> {
+    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
+    let variant_name = variant.name()?;
+    let enum_name = variant.parent_enum().name()?;
+    if !matches!(variant.kind(), ast::StructKind::Unit) {
+        mark::hit!(test_gen_default_on_non_unit_variant_not_implemented);
+        return None;
+    }
+
+    if existing_default_impl(&ctx.sema, &variant).is_some() {
+        mark::hit!(test_gen_default_impl_already_exists);
+        return None;
+    }
+
+    let target = variant.syntax().text_range();
+    acc.add(
+        AssistId("generate_default_from_enum_variant", AssistKind::Generate),
+        "Generate `Default` impl from this enum variant",
+        target,
+        |edit| {
+            let start_offset = variant.parent_enum().syntax().text_range().end();
+            let buf = format!(
+                r#"
+
+impl Default for {0} {{
+    fn default() -> Self {{
+        Self::{1}
+    }}
+}}"#,
+                enum_name, variant_name
+            );
+            edit.insert(start_offset, buf);
+        },
+    )
+}
+
+fn existing_default_impl(
+    sema: &'_ hir::Semantics<'_, RootDatabase>,
+    variant: &ast::Variant,
+) -> Option<()> {
+    let variant = sema.to_def(variant)?;
+    let enum_ = variant.parent_enum(sema.db);
+    let krate = enum_.module(sema.db).krate();
+
+    let default_trait = FamousDefs(sema, Some(krate)).core_default_Default()?;
+    let enum_type = enum_.ty(sema.db);
+
+    if enum_type.impls_trait(sema.db, default_trait, &[]) {
+        Some(())
+    } else {
+        None
+    }
+}
+
+#[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) {
+        let fixture =
+            format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
+        check_assist_not_applicable(generate_default_from_enum_variant, &fixture)
+    }
+
+    #[test]
+    fn test_generate_default_from_variant() {
+        check_assist(
+            generate_default_from_enum_variant,
+            r#"enum Variant {
+                Undefined,
+                Minor<|>,
+                Major,
+            }"#,
+            r#"enum Variant {
+                Undefined,
+                Minor,
+                Major,
+            }
+
+impl Default for Variant {
+    fn default() -> Self {
+        Self::Minor
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_default_already_implemented() {
+        mark::check!(test_gen_default_impl_already_exists);
+        check_not_applicable(
+            r#"enum Variant {
+                Undefined,
+                Minor<|>,
+                Major,
+            }
+
+            impl Default for Variant {
+                fn default() -> Self {
+                    Self::Minor
+                }
+            }"#,
+        );
+    }
+
+    #[test]
+    fn test_add_from_impl_no_element() {
+        mark::check!(test_gen_default_on_non_unit_variant_not_implemented);
+        check_not_applicable(
+            r#"enum Variant {
+            Undefined,
+            Minor(u32)<|>,
+            Major,
+        }"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_default_from_variant_with_one_variant() {
+        check_assist(
+            generate_default_from_enum_variant,
+            r#"enum Variant { Undefi<|>ned }"#,
+            r#"enum Variant { Undefined }
+
+impl Default for Variant {
+    fn default() -> Self {
+        Self::Undefined
+    }
+}"#,
+        );
+    }
+}
diff --git a/crates/assists/src/lib.rs b/crates/assists/src/lib.rs
index b8ce7418d..6e736ccb3 100644
--- a/crates/assists/src/lib.rs
+++ b/crates/assists/src/lib.rs
@@ -137,6 +137,7 @@ mod handlers {
     mod flip_comma;
     mod flip_trait_bound;
     mod generate_derive;
+    mod generate_default_from_enum_variant;
     mod generate_from_impl_for_enum;
     mod generate_function;
     mod generate_impl;
@@ -186,6 +187,7 @@ mod handlers {
             flip_comma::flip_comma,
             flip_trait_bound::flip_trait_bound,
             generate_derive::generate_derive,
+            generate_default_from_enum_variant::generate_default_from_enum_variant,
             generate_from_impl_for_enum::generate_from_impl_for_enum,
             generate_function::generate_function,
             generate_impl::generate_impl,
diff --git a/crates/assists/src/tests/generated.rs b/crates/assists/src/tests/generated.rs
index 853bde09c..cc7c4a343 100644
--- a/crates/assists/src/tests/generated.rs
+++ b/crates/assists/src/tests/generated.rs
@@ -365,6 +365,33 @@ fn foo<T: Copy + Clone>() { }
 }
 
 #[test]
+fn doctest_generate_default_from_enum_variant() {
+    check_doc_test(
+        "generate_default_from_enum_variant",
+        r#####"
+enum Version {
+ Undefined,
+ Minor<|>,
+ Major,
+}
+"#####,
+        r#####"
+enum Version {
+ Undefined,
+ Minor,
+ Major,
+}
+
+impl Default for Version {
+    fn default() -> Self {
+        Self::Minor
+    }
+}
+"#####,
+    )
+}
+
+#[test]
 fn doctest_generate_derive() {
     check_doc_test(
         "generate_derive",