Unnamed repository; edit this file 'description' to name the repository.
Diffstat (limited to 'crates/ide-assists/src/handlers/generate_enum_projection_method.rs')
| -rw-r--r-- | crates/ide-assists/src/handlers/generate_enum_projection_method.rs | 339 |
1 files changed, 339 insertions, 0 deletions
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..fe7a3cd535 --- /dev/null +++ b/crates/ide-assists/src/handlers/generate_enum_projection_method.rs @@ -0,0 +1,339 @@ +use ide_db::assists::GroupLabel; +use itertools::Itertools; +use stdx::to_lower_snake_case; +use syntax::ast::HasVisibility; +use syntax::ast::{self, AstNode, HasName}; + +use crate::{ + utils::{add_method_to_adt, find_struct_impl}, + AssistContext, AssistId, AssistKind, Assists, +}; + +// Assist: generate_enum_try_into_method +// +// Generate a `try_into_` method for this enum variant. +// +// ``` +// enum Value { +// Number(i32), +// Text(String)$0, +// } +// ``` +// -> +// ``` +// enum Value { +// Number(i32), +// Text(String), +// } +// +// impl Value { +// fn try_into_text(self) -> Result<String, Self> { +// 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 a `try_into_` method for this 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 this 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 this 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::<ast::Variant>()?; + 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_group( + &GroupLabel("Generate an `is_`,`as_`, or `try_into_` for this enum variant".to_owned()), + 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<String, Self> { + 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<String, Self> { + 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<String, Self> { + 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 + } + } +}"#, + ); + } +} |