Generate enum variant assist

This also disables "generate function" when what we clearly want is to
generate an enum variant.

Co-authored-by: Maarten Flippo <maartenflippo@outlook.com>

4 files changed, 230 insertions, 0 deletions

diff --git a/crates/ide-assists/src/handlers/generate_enum_variant.rs b/crates/ide-assists/src/handlers/generate_enum_variant.rs
new file mode 100644
index 00000000000..98e56a9633e
--- /dev/null
+++ b/crates/ide-assists/src/handlers/generate_enum_variant.rs
@@ -0,0 +1,185 @@
+use hir::HasSource;
+use ide_db::assists::{AssistId, AssistKind};
+use syntax::{
+    ast::{self, edit::IndentLevel},
+    AstNode, TextSize,
+};
+
+use crate::assist_context::{AssistContext, Assists};
+
+// Assist: generate_enum_variant
+//
+// Adds a variant to an enum.
+//
+// ```
+// enum Countries {
+//     Ghana,
+// }
+//
+// fn main() {
+//     let country = Countries::Lesotho$0;
+// }
+// ```
+// ->
+// ```
+// enum Countries {
+//     Ghana,
+//     Lesotho,
+// }
+//
+// fn main() {
+//     let country = Countries::Lesotho;
+// }
+// ```
+pub(crate) fn generate_enum_variant(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    let path_expr: ast::PathExpr = ctx.find_node_at_offset()?;
+    let path = path_expr.path()?;
+
+    if ctx.sema.resolve_path(&path).is_some() {
+        // No need to generate anything if the path resolves
+        return None;
+    }
+
+    let name_ref = path.segment()?.name_ref()?;
+
+    if let Some(hir::PathResolution::Def(hir::ModuleDef::Adt(hir::Adt::Enum(e)))) =
+        ctx.sema.resolve_path(&path.qualifier()?)
+    {
+        let target = path.syntax().text_range();
+        return add_variant_to_accumulator(acc, ctx, target, e, &name_ref);
+    }
+
+    None
+}
+
+fn add_variant_to_accumulator(
+    acc: &mut Assists,
+    ctx: &AssistContext,
+    target: syntax::TextRange,
+    adt: hir::Enum,
+    name_ref: &ast::NameRef,
+) -> Option<()> {
+    let adt_ast = get_enum_ast(ctx, adt)?;
+
+    let enum_indent_level = IndentLevel::from_node(&adt_ast.syntax());
+
+    let offset = adt_ast.variant_list()?.syntax().text_range().end() - TextSize::of('}');
+
+    let prefix = if adt_ast.variant_list()?.variants().next().is_none() {
+        format!("\n{}", IndentLevel(1))
+    } else {
+        format!("{}", IndentLevel(1))
+    };
+    let text = format!("{}{},\n{}", prefix, name_ref, enum_indent_level);
+
+    acc.add(
+        AssistId("generate_enum_variant", AssistKind::Generate),
+        "Generate variant",
+        target,
+        |builder| builder.insert(offset, text),
+    )
+}
+
+fn get_enum_ast(ctx: &AssistContext, adt: hir::Enum) -> Option<ast::Enum> {
+    let range = adt.source(ctx.db())?.syntax().original_file_range(ctx.db());
+    let file = ctx.sema.parse(range.file_id);
+    let adt_ast: ast::Enum =
+        ctx.sema.find_node_at_offset_with_macros(file.syntax(), range.range.start())?;
+
+    Some(adt_ast)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    #[test]
+    fn generate_basic_enum_variant_in_empty_enum() {
+        check_assist(
+            generate_enum_variant,
+            r"
+enum Foo {}
+fn main() {
+    Foo::Bar$0
+}
+",
+            r"
+enum Foo {
+    Bar,
+}
+fn main() {
+    Foo::Bar
+}
+",
+        )
+    }
+
+    #[test]
+    fn generate_basic_enum_variant_in_non_empty_enum() {
+        check_assist(
+            generate_enum_variant,
+            r"
+enum Foo {
+    Bar,
+}
+fn main() {
+    Foo::Baz$0
+}
+",
+            r"
+enum Foo {
+    Bar,
+    Baz,
+}
+fn main() {
+    Foo::Baz
+}
+",
+        )
+    }
+
+    #[test]
+    fn not_applicable_for_existing_variant() {
+        check_assist_not_applicable(
+            generate_enum_variant,
+            r"
+enum Foo {
+    Bar,
+}
+fn main() {
+    Foo::Bar$0
+}
+",
+        )
+    }
+
+    #[test]
+    fn indentation_level_is_correct() {
+        check_assist(
+            generate_enum_variant,
+            r"
+mod m {
+    enum Foo {
+        Bar,
+    }
+}
+fn main() {
+    m::Foo::Baz$0
+}
+",
+            r"
+mod m {
+    enum Foo {
+        Bar,
+        Baz,
+    }
+}
+fn main() {
+    m::Foo::Baz
+}
+",
+        )
+    }
+}
diff --git a/crates/ide-assists/src/handlers/generate_function.rs b/crates/ide-assists/src/handlers/generate_function.rs
index 24cef5630f1..e3ef321d4b2 100644
--- a/crates/ide-assists/src/handlers/generate_function.rs
+++ b/crates/ide-assists/src/handlers/generate_function.rs
@@ -71,6 +71,10 @@ fn gen_fn(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
                 get_fn_target(ctx, &target_module, call.clone())?
             }
             Some(hir::PathResolution::Def(hir::ModuleDef::Adt(adt))) => {
+                if let hir::Adt::Enum(_) = adt {
+                    return None;
+                }
+
                 let current_module = ctx.sema.scope(call.syntax())?.module();
                 let module = adt.module(ctx.sema.db);
                 target_module = if current_module == module { None } else { Some(module) };
@@ -1738,4 +1742,17 @@ fn foo(value: usize) ${0:-> _} {
 ",
         )
     }
+
+    #[test]
+    fn not_applicable_for_enum_variant() {
+        check_assist_not_applicable(
+            generate_function,
+            r"
+enum Foo {}
+fn main() {
+    Foo::Bar$0(true)
+}
+",
+        );
+    }
 }
diff --git a/crates/ide-assists/src/lib.rs b/crates/ide-assists/src/lib.rs
index b8607579a32..94fbaff400b 100644
--- a/crates/ide-assists/src/lib.rs
+++ b/crates/ide-assists/src/lib.rs
@@ -139,6 +139,7 @@ mod handlers {
     mod generate_documentation_template;
     mod generate_enum_is_method;
     mod generate_enum_projection_method;
+    mod generate_enum_variant;
     mod generate_from_impl_for_enum;
     mod generate_function;
     mod generate_getter;
@@ -227,6 +228,7 @@ mod handlers {
             generate_enum_is_method::generate_enum_is_method,
             generate_enum_projection_method::generate_enum_as_method,
             generate_enum_projection_method::generate_enum_try_into_method,
+            generate_enum_variant::generate_enum_variant,
             generate_from_impl_for_enum::generate_from_impl_for_enum,
             generate_function::generate_function,
             generate_impl::generate_impl,
diff --git a/crates/ide-assists/src/tests/generated.rs b/crates/ide-assists/src/tests/generated.rs
index c9b0b1ce82e..2d571014090 100644
--- a/crates/ide-assists/src/tests/generated.rs
+++ b/crates/ide-assists/src/tests/generated.rs
@@ -1007,6 +1007,32 @@ impl Value {
 }
 
 #[test]
+fn doctest_generate_enum_variant() {
+    check_doc_test(
+        "generate_enum_variant",
+        r#####"
+enum Countries {
+    Ghana,
+}
+
+fn main() {
+    let country = Countries::Lesotho$0;
+}
+"#####,
+        r#####"
+enum Countries {
+    Ghana,
+    Lesotho,
+}
+
+fn main() {
+    let country = Countries::Lesotho;
+}
+"#####,
+    )
+}
+
+#[test]
 fn doctest_generate_from_impl_for_enum() {
     check_doc_test(
         "generate_from_impl_for_enum",