From 13d663dd16430cec18d7eccd214c3d4891b1a9a1 Mon Sep 17 00:00:00 2001 From: Yoshua Wuyts Date: Fri, 5 Feb 2021 01:57:39 +0100 Subject: add `generate-enum-match` assist --- .../src/handlers/generate_enum_match_method.rs | 293 +++++++++++++++++++++ crates/assists/src/lib.rs | 2 + crates/assists/src/tests/generated.rs | 27 ++ 3 files changed, 322 insertions(+) create mode 100644 crates/assists/src/handlers/generate_enum_match_method.rs diff --git a/crates/assists/src/handlers/generate_enum_match_method.rs b/crates/assists/src/handlers/generate_enum_match_method.rs new file mode 100644 index 000000000..079ed27bd --- /dev/null +++ b/crates/assists/src/handlers/generate_enum_match_method.rs @@ -0,0 +1,293 @@ +use hir::Adt; +use stdx::format_to; +use syntax::ast::{self, AstNode, NameOwner}; +use syntax::{ast::VisibilityOwner, T}; +use test_utils::mark; + +use crate::{AssistContext, AssistId, AssistKind, Assists}; + +// Assist: generate_enum_match_method +// +// Generate an `is_` method for an enum variant. +// +// ``` +// enum Version { +// Undefined, +// Minor$0, +// Major, +// } +// ``` +// -> +// ``` +// enum Version { +// Undefined, +// Minor, +// Major, +// } +// +// impl Version { +// fn is_minor(&self) -> bool { +// matches!(self, Self::Minor) +// } +// } +// ``` +pub(crate) fn generate_enum_match_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { + let variant = ctx.find_node_at_offset::()?; + 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 fn_name = to_lower_snake_case(&format!("{}", variant_name)); + + // Return early if we've found an existing new fn + let impl_def = find_struct_impl(&ctx, &parent_enum, format!("is_{}", fn_name).as_str())?; + + let target = variant.syntax().text_range(); + acc.add( + AssistId("generate_enum_match_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, + " {}fn is_{}(&self) -> bool {{ + matches!(self, Self::{}) + }}", + vis, + fn_name, + variant_name + ); + + let start_offset = impl_def + .and_then(|impl_def| { + buf.push('\n'); + let start = impl_def + .syntax() + .descendants_with_tokens() + .find(|t| t.kind() == T!['{'])? + .text_range() + .end(); + + Some(start) + }) + .unwrap_or_else(|| { + buf = generate_impl_text(&parent_enum, &buf); + parent_enum.syntax().text_range().end() + }); + + builder.insert(start_offset, buf); + }, + ) +} + +// Generates the surrounding `impl Type { }` including type and lifetime +// parameters +fn generate_impl_text(strukt: &ast::Enum, code: &str) -> String { + let mut buf = String::with_capacity(code.len()); + buf.push_str("\n\nimpl"); + buf.push_str(" "); + buf.push_str(strukt.name().unwrap().text()); + format_to!(buf, " {{\n{}\n}}", code); + buf +} + +fn to_lower_snake_case(s: &str) -> String { + let mut buf = String::with_capacity(s.len()); + let mut prev = false; + for c in s.chars() { + if c.is_ascii_uppercase() && prev { + buf.push('_') + } + prev = true; + + buf.push(c.to_ascii_lowercase()); + } + buf +} + +// Uses a syntax-driven approach to find any impl blocks for the struct that +// exist within the module/file +// +// Returns `None` if we've found an existing `new` fn +// +// FIXME: change the new fn checking to a more semantic approach when that's more +// viable (e.g. we process proc macros, etc) +fn find_struct_impl( + ctx: &AssistContext, + strukt: &ast::Enum, + name: &str, +) -> Option> { + let db = ctx.db(); + let module = strukt.syntax().ancestors().find(|node| { + ast::Module::can_cast(node.kind()) || ast::SourceFile::can_cast(node.kind()) + })?; + + let struct_def = ctx.sema.to_def(strukt)?; + + let block = module.descendants().filter_map(ast::Impl::cast).find_map(|impl_blk| { + let blk = ctx.sema.to_def(&impl_blk)?; + + // FIXME: handle e.g. `struct S; impl S {}` + // (we currently use the wrong type parameter) + // also we wouldn't want to use e.g. `impl S` + let same_ty = match blk.target_ty(db).as_adt() { + Some(def) => def == Adt::Enum(struct_def), + None => false, + }; + let not_trait_impl = blk.target_trait(db).is_none(); + + if !(same_ty && not_trait_impl) { + None + } else { + Some(impl_blk) + } + }); + + if let Some(ref impl_blk) = block { + if has_fn(impl_blk, name) { + mark::hit!(test_gen_enum_match_impl_already_exists); + return None; + } + } + + Some(block) +} + +fn has_fn(imp: &ast::Impl, rhs_name: &str) -> bool { + if let Some(il) = imp.assoc_item_list() { + for item in il.assoc_items() { + if let ast::AssocItem::Fn(f) = item { + if let Some(name) = f.name() { + if name.text().eq_ignore_ascii_case(rhs_name) { + return true; + } + } + } + } + } + + false +} + +#[cfg(test)] +mod tests { + use ide_db::helpers::FamousDefs; + 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_enum_match_method, &fixture) + } + + #[test] + fn test_generate_enum_match_from_variant() { + check_assist( + generate_enum_match_method, + r#" +enum Variant { + Undefined, + Minor$0, + Major, +}"#, + r#"enum Variant { + Undefined, + Minor, + Major, +} + +impl Variant { + fn is_minor(&self) -> bool { + matches!(self, Self::Minor) + } +}"#, + ); + } + + #[test] + fn test_generate_enum_match_already_implemented() { + mark::check!(test_gen_enum_match_impl_already_exists); + check_not_applicable( + r#" +enum Variant { + Undefined, + Minor$0, + Major, +} + +impl Variant { + fn is_minor(&self) -> bool { + matches!(self, Self::Minor) + } +}"#, + ); + } + + #[test] + fn test_add_from_impl_no_element() { + mark::check!(test_gen_enum_match_on_non_unit_variant_not_implemented); + check_not_applicable( + r#" +enum Variant { + Undefined, + Minor(u32)$0, + Major, +}"#, + ); + } + + #[test] + fn test_generate_enum_match_from_variant_with_one_variant() { + check_assist( + generate_enum_match_method, + r#"enum Variant { Undefi$0ned }"#, + r#" +enum Variant { Undefined } + +impl Variant { + fn is_undefined(&self) -> bool { + matches!(self, Self::Undefined) + } +}"#, + ); + } + + #[test] + fn test_generate_enum_match_from_variant_with_visibility_marker() { + check_assist( + generate_enum_match_method, + r#" +pub(crate) enum Variant { + Undefined, + Minor$0, + Major, +}"#, + r#"pub(crate) enum Variant { + Undefined, + Minor, + Major, +} + +impl Variant { + pub(crate) fn is_minor(&self) -> bool { + matches!(self, Self::Minor) + } +}"#, + ); + } +} diff --git a/crates/assists/src/lib.rs b/crates/assists/src/lib.rs index 559b9651e..a18232877 100644 --- a/crates/assists/src/lib.rs +++ b/crates/assists/src/lib.rs @@ -126,6 +126,7 @@ mod handlers { mod flip_trait_bound; mod generate_default_from_enum_variant; mod generate_derive; + mod generate_enum_match_method; mod generate_from_impl_for_enum; mod generate_function; mod generate_impl; @@ -183,6 +184,7 @@ 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_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 9aa807f10..ae7b400e2 100644 --- a/crates/assists/src/tests/generated.rs +++ b/crates/assists/src/tests/generated.rs @@ -432,6 +432,33 @@ struct Point { ) } +#[test] +fn doctest_generate_enum_match_method() { + check_doc_test( + "generate_enum_match_method", + r#####" +enum Version { + Undefined, + Minor$0, + Major, +} +"#####, + r#####" +enum Version { + Undefined, + Minor, + Major, +} + +impl Version { + fn is_minor(&self) -> bool { + matches!(self, Self::Minor) + } +} +"#####, + ) +} + #[test] fn doctest_generate_from_impl_for_enum() { check_doc_test( -- cgit v1.2.3