diff --git a/crates/ide_assists/src/handlers/generate_enum_match_method.rs b/crates/ide_assists/src/handlers/generate_enum_is_method.rs similarity index 58% rename from crates/ide_assists/src/handlers/generate_enum_match_method.rs rename to crates/ide_assists/src/handlers/generate_enum_is_method.rs index aeb887e717..7e181a4801 100644 --- a/crates/ide_assists/src/handlers/generate_enum_match_method.rs +++ b/crates/ide_assists/src/handlers/generate_enum_is_method.rs @@ -1,14 +1,13 @@ -use stdx::{format_to, to_lower_snake_case}; +use stdx::to_lower_snake_case; use syntax::ast::VisibilityOwner; use syntax::ast::{self, AstNode, NameOwner}; -use test_utils::mark; use crate::{ - utils::{find_impl_block_end, find_struct_impl, generate_impl_text}, + utils::{add_method_to_adt, find_struct_impl}, AssistContext, AssistId, AssistKind, Assists, }; -// Assist: generate_enum_match_method +// Assist: generate_enum_is_method // // Generate an `is_` method for an enum variant. // @@ -34,79 +33,52 @@ use crate::{ // } // } // ``` -pub(crate) fn generate_enum_match_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { +pub(crate) fn generate_enum_is_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 parent_enum = ast::Adt::Enum(variant.parent_enum()); + let pattern_suffix = match variant.kind() { + ast::StructKind::Record(_) => " { .. }", + ast::StructKind::Tuple(_) => "(..)", + ast::StructKind::Unit => "", + }; let enum_lowercase_name = to_lower_snake_case(&parent_enum.name()?.to_string()); - let fn_name = to_lower_snake_case(&variant_name.to_string()); + let fn_name = format!("is_{}", &to_lower_snake_case(variant_name.text())); // Return early if we've found an existing new fn - let impl_def = find_struct_impl( - &ctx, - &ast::Adt::Enum(parent_enum.clone()), - format!("is_{}", fn_name).as_str(), - )?; + let impl_def = find_struct_impl(&ctx, &parent_enum, &fn_name)?; let target = variant.syntax().text_range(); acc.add( - AssistId("generate_enum_match_method", AssistKind::Generate), + AssistId("generate_enum_is_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, + let method = format!( " /// Returns `true` if the {} is [`{}`]. - {}fn is_{}(&self) -> bool {{ - matches!(self, Self::{}) + {}fn {}(&self) -> bool {{ + matches!(self, Self::{}{}) }}", - enum_lowercase_name, - variant_name, - vis, - fn_name, - variant_name + enum_lowercase_name, variant_name, vis, fn_name, variant_name, pattern_suffix, ); - let start_offset = impl_def - .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf)) - .unwrap_or_else(|| { - buf = generate_impl_text(&ast::Adt::Enum(parent_enum.clone()), &buf); - parent_enum.syntax().text_range().end() - }); - - builder.insert(start_offset, buf); + add_method_to_adt(builder, &parent_enum, impl_def, &method); }, ) } #[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) { - check_assist_not_applicable(generate_enum_match_method, ra_fixture) - } - #[test] - fn test_generate_enum_match_from_variant() { + fn test_generate_enum_is_from_variant() { check_assist( - generate_enum_match_method, + generate_enum_is_method, r#" enum Variant { Undefined, @@ -129,8 +101,9 @@ impl Variant { } #[test] - fn test_generate_enum_match_already_implemented() { - check_not_applicable( + fn test_generate_enum_is_already_implemented() { + check_assist_not_applicable( + generate_enum_is_method, r#" enum Variant { Undefined, @@ -147,22 +120,59 @@ impl Variant { } #[test] - fn test_add_from_impl_no_element() { - mark::check!(test_gen_enum_match_on_non_unit_variant_not_implemented); - check_not_applicable( + fn test_generate_enum_is_from_tuple_variant() { + check_assist( + generate_enum_is_method, r#" enum Variant { Undefined, Minor(u32)$0, Major, +}"#, + r#"enum Variant { + Undefined, + Minor(u32), + Major, +} + +impl Variant { + /// Returns `true` if the variant is [`Minor`]. + fn is_minor(&self) -> bool { + matches!(self, Self::Minor(..)) + } }"#, ); } #[test] - fn test_generate_enum_match_from_variant_with_one_variant() { + fn test_generate_enum_is_from_record_variant() { check_assist( - generate_enum_match_method, + generate_enum_is_method, + r#" +enum Variant { + Undefined, + Minor { foo: i32 }$0, + Major, +}"#, + r#"enum Variant { + Undefined, + Minor { foo: i32 }, + Major, +} + +impl Variant { + /// Returns `true` if the variant is [`Minor`]. + fn is_minor(&self) -> bool { + matches!(self, Self::Minor { .. }) + } +}"#, + ); + } + + #[test] + fn test_generate_enum_is_from_variant_with_one_variant() { + check_assist( + generate_enum_is_method, r#"enum Variant { Undefi$0ned }"#, r#" enum Variant { Undefined } @@ -177,9 +187,9 @@ impl Variant { } #[test] - fn test_generate_enum_match_from_variant_with_visibility_marker() { + fn test_generate_enum_is_from_variant_with_visibility_marker() { check_assist( - generate_enum_match_method, + generate_enum_is_method, r#" pub(crate) enum Variant { Undefined, @@ -202,9 +212,9 @@ impl Variant { } #[test] - fn test_multiple_generate_enum_match_from_variant() { + fn test_multiple_generate_enum_is_from_variant() { check_assist( - generate_enum_match_method, + generate_enum_is_method, r#" enum Variant { Undefined, diff --git a/crates/ide_assists/src/handlers/generate_enum_projection_method.rs b/crates/ide_assists/src/handlers/generate_enum_projection_method.rs new file mode 100644 index 0000000000..871bcab50b --- /dev/null +++ b/crates/ide_assists/src/handlers/generate_enum_projection_method.rs @@ -0,0 +1,331 @@ +use itertools::Itertools; +use stdx::to_lower_snake_case; +use syntax::ast::VisibilityOwner; +use syntax::ast::{self, AstNode, NameOwner}; + +use crate::{ + utils::{add_method_to_adt, find_struct_impl}, + AssistContext, AssistId, AssistKind, Assists, +}; + +// Assist: generate_enum_try_into_method +// +// Generate an `try_into_` method for an enum variant. +// +// ``` +// enum Value { +// Number(i32), +// Text(String)$0, +// } +// ``` +// -> +// ``` +// enum Value { +// Number(i32), +// Text(String), +// } +// +// impl Value { +// fn try_into_text(self) -> Result { +// if let Self::Text(v) = self { +// Ok(v) +// } else { +// Err(self) +// } +// } +// } +// ``` +pub(crate) fn generate_enum_try_into_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { + generate_enum_projection_method( + acc, + ctx, + "generate_enum_try_into_method", + "Generate an `try_into_` method for an enum variant", + ProjectionProps { + fn_name_prefix: "try_into", + self_param: "self", + return_prefix: "Result<", + return_suffix: ", Self>", + happy_case: "Ok", + sad_case: "Err(self)", + }, + ) +} + +// Assist: generate_enum_as_method +// +// Generate an `as_` method for an enum variant. +// +// ``` +// enum Value { +// Number(i32), +// Text(String)$0, +// } +// ``` +// -> +// ``` +// enum Value { +// Number(i32), +// Text(String), +// } +// +// impl Value { +// fn as_text(&self) -> Option<&String> { +// if let Self::Text(v) = self { +// Some(v) +// } else { +// None +// } +// } +// } +// ``` +pub(crate) fn generate_enum_as_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { + generate_enum_projection_method( + acc, + ctx, + "generate_enum_as_method", + "Generate an `as_` method for an enum variant", + ProjectionProps { + fn_name_prefix: "as", + self_param: "&self", + return_prefix: "Option<&", + return_suffix: ">", + happy_case: "Some", + sad_case: "None", + }, + ) +} + +struct ProjectionProps { + fn_name_prefix: &'static str, + self_param: &'static str, + return_prefix: &'static str, + return_suffix: &'static str, + happy_case: &'static str, + sad_case: &'static str, +} + +fn generate_enum_projection_method( + acc: &mut Assists, + ctx: &AssistContext, + assist_id: &'static str, + assist_description: &str, + props: ProjectionProps, +) -> Option<()> { + let variant = ctx.find_node_at_offset::()?; + let variant_name = variant.name()?; + let parent_enum = ast::Adt::Enum(variant.parent_enum()); + + let (pattern_suffix, field_type, bound_name) = match variant.kind() { + ast::StructKind::Record(record) => { + let (field,) = record.fields().collect_tuple()?; + let name = field.name()?.to_string(); + let ty = field.ty()?; + let pattern_suffix = format!(" {{ {} }}", name); + (pattern_suffix, ty, name) + } + ast::StructKind::Tuple(tuple) => { + let (field,) = tuple.fields().collect_tuple()?; + let ty = field.ty()?; + ("(v)".to_owned(), ty, "v".to_owned()) + } + ast::StructKind::Unit => return None, + }; + + let fn_name = format!("{}_{}", props.fn_name_prefix, &to_lower_snake_case(variant_name.text())); + + // Return early if we've found an existing new fn + let impl_def = find_struct_impl(&ctx, &parent_enum, &fn_name)?; + + let target = variant.syntax().text_range(); + acc.add(AssistId(assist_id, AssistKind::Generate), assist_description, target, |builder| { + let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v)); + let method = format!( + " {0}fn {1}({2}) -> {3}{4}{5} {{ + if let Self::{6}{7} = self {{ + {8}({9}) + }} else {{ + {10} + }} + }}", + vis, + fn_name, + props.self_param, + props.return_prefix, + field_type.syntax(), + props.return_suffix, + variant_name, + pattern_suffix, + props.happy_case, + bound_name, + props.sad_case, + ); + + add_method_to_adt(builder, &parent_enum, impl_def, &method); + }) +} + +#[cfg(test)] +mod tests { + use crate::tests::{check_assist, check_assist_not_applicable}; + + use super::*; + + #[test] + fn test_generate_enum_try_into_tuple_variant() { + check_assist( + generate_enum_try_into_method, + r#" +enum Value { + Number(i32), + Text(String)$0, +}"#, + r#"enum Value { + Number(i32), + Text(String), +} + +impl Value { + fn try_into_text(self) -> Result { + if let Self::Text(v) = self { + Ok(v) + } else { + Err(self) + } + } +}"#, + ); + } + + #[test] + fn test_generate_enum_try_into_already_implemented() { + check_assist_not_applicable( + generate_enum_try_into_method, + r#"enum Value { + Number(i32), + Text(String)$0, +} + +impl Value { + fn try_into_text(self) -> Result { + if let Self::Text(v) = self { + Ok(v) + } else { + Err(self) + } + } +}"#, + ); + } + + #[test] + fn test_generate_enum_try_into_unit_variant() { + check_assist_not_applicable( + generate_enum_try_into_method, + r#"enum Value { + Number(i32), + Text(String), + Unit$0, +}"#, + ); + } + + #[test] + fn test_generate_enum_try_into_record_with_multiple_fields() { + check_assist_not_applicable( + generate_enum_try_into_method, + r#"enum Value { + Number(i32), + Text(String), + Both { first: i32, second: String }$0, +}"#, + ); + } + + #[test] + fn test_generate_enum_try_into_tuple_with_multiple_fields() { + check_assist_not_applicable( + generate_enum_try_into_method, + r#"enum Value { + Number(i32), + Text(String, String)$0, +}"#, + ); + } + + #[test] + fn test_generate_enum_try_into_record_variant() { + check_assist( + generate_enum_try_into_method, + r#"enum Value { + Number(i32), + Text { text: String }$0, +}"#, + r#"enum Value { + Number(i32), + Text { text: String }, +} + +impl Value { + fn try_into_text(self) -> Result { + if let Self::Text { text } = self { + Ok(text) + } else { + Err(self) + } + } +}"#, + ); + } + + #[test] + fn test_generate_enum_as_tuple_variant() { + check_assist( + generate_enum_as_method, + r#" +enum Value { + Number(i32), + Text(String)$0, +}"#, + r#"enum Value { + Number(i32), + Text(String), +} + +impl Value { + fn as_text(&self) -> Option<&String> { + if let Self::Text(v) = self { + Some(v) + } else { + None + } + } +}"#, + ); + } + + #[test] + fn test_generate_enum_as_record_variant() { + check_assist( + generate_enum_as_method, + r#"enum Value { + Number(i32), + Text { text: String }$0, +}"#, + r#"enum Value { + Number(i32), + Text { text: String }, +} + +impl Value { + fn as_text(&self) -> Option<&String> { + if let Self::Text { text } = self { + Some(text) + } else { + None + } + } +}"#, + ); + } +} diff --git a/crates/ide_assists/src/lib.rs b/crates/ide_assists/src/lib.rs index f4c7e6fbf7..4c067d4511 100644 --- a/crates/ide_assists/src/lib.rs +++ b/crates/ide_assists/src/lib.rs @@ -128,7 +128,8 @@ mod handlers { mod flip_trait_bound; mod generate_default_from_enum_variant; mod generate_derive; - mod generate_enum_match_method; + mod generate_enum_is_method; + mod generate_enum_projection_method; mod generate_from_impl_for_enum; mod generate_function; mod generate_getter; @@ -189,7 +190,9 @@ 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_enum_is_method::generate_enum_is_method, + generate_enum_projection_method::generate_enum_try_into_method, + generate_enum_projection_method::generate_enum_as_method, generate_from_impl_for_enum::generate_from_impl_for_enum, generate_function::generate_function, generate_getter::generate_getter, diff --git a/crates/ide_assists/src/tests/generated.rs b/crates/ide_assists/src/tests/generated.rs index d428758222..7f6dbbccf2 100644 --- a/crates/ide_assists/src/tests/generated.rs +++ b/crates/ide_assists/src/tests/generated.rs @@ -483,9 +483,38 @@ struct Point { } #[test] -fn doctest_generate_enum_match_method() { +fn doctest_generate_enum_as_method() { check_doc_test( - "generate_enum_match_method", + "generate_enum_as_method", + r#####" +enum Value { + Number(i32), + Text(String)$0, +} +"#####, + r#####" +enum Value { + Number(i32), + Text(String), +} + +impl Value { + fn as_text(&self) -> Option<&String> { + if let Self::Text(v) = self { + Some(v) + } else { + None + } + } +} +"#####, + ) +} + +#[test] +fn doctest_generate_enum_is_method() { + check_doc_test( + "generate_enum_is_method", r#####" enum Version { Undefined, @@ -510,6 +539,35 @@ impl Version { ) } +#[test] +fn doctest_generate_enum_try_into_method() { + check_doc_test( + "generate_enum_try_into_method", + r#####" +enum Value { + Number(i32), + Text(String)$0, +} +"#####, + r#####" +enum Value { + Number(i32), + Text(String), +} + +impl Value { + fn try_into_text(self) -> Result { + if let Self::Text(v) = self { + Ok(v) + } else { + Err(self) + } + } +} +"#####, + ) +} + #[test] fn doctest_generate_from_impl_for_enum() { check_doc_test( diff --git a/crates/ide_assists/src/utils.rs b/crates/ide_assists/src/utils.rs index 276792bc1a..880ab6fe3e 100644 --- a/crates/ide_assists/src/utils.rs +++ b/crates/ide_assists/src/utils.rs @@ -21,7 +21,7 @@ use syntax::{ }; use crate::{ - assist_context::AssistContext, + assist_context::{AssistBuilder, AssistContext}, ast_transform::{self, AstTransform, QualifyPaths, SubstituteTypeParams}, }; @@ -464,3 +464,25 @@ fn generate_impl_text_inner(adt: &ast::Adt, trait_text: Option<&str>, code: &str buf } + +pub(crate) fn add_method_to_adt( + builder: &mut AssistBuilder, + adt: &ast::Adt, + impl_def: Option, + method: &str, +) { + let mut buf = String::with_capacity(method.len() + 2); + if impl_def.is_some() { + buf.push('\n'); + } + buf.push_str(method); + + let start_offset = impl_def + .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf)) + .unwrap_or_else(|| { + buf = generate_impl_text(&adt, &buf); + adt.syntax().text_range().end() + }); + + builder.insert(start_offset, buf); +}