From 2347da8c8d042146a6a0b4cd9570a7454afa25e4 Mon Sep 17 00:00:00 2001 From: Amos Wenger Date: Sat, 21 May 2022 01:18:35 +0200 Subject: [PATCH] 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 --- .../src/handlers/generate_enum_variant.rs | 185 ++++++++++++++++++ .../src/handlers/generate_function.rs | 17 ++ crates/ide-assists/src/lib.rs | 2 + crates/ide-assists/src/tests/generated.rs | 26 +++ 4 files changed, 230 insertions(+) create mode 100644 crates/ide-assists/src/handlers/generate_enum_variant.rs 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 { + 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 @@ pub(crate) fn all() -> &'static [Handler] { 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 @@ -1006,6 +1006,32 @@ fn try_into_text(self) -> Result { ) } +#[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(