1 use itertools::Itertools;
2 use stdx::to_lower_snake_case;
3 use syntax::ast::VisibilityOwner;
4 use syntax::ast::{self, AstNode, NameOwner};
7 utils::{add_method_to_adt, find_struct_impl},
8 AssistContext, AssistId, AssistKind, Assists,
11 // Assist: generate_enum_try_into_method
13 // Generate an `try_into_` method for an enum variant.
29 // fn try_into_text(self) -> Result<String, Self> {
30 // if let Self::Text(v) = self {
38 pub(crate) fn generate_enum_try_into_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
39 generate_enum_projection_method(
42 "generate_enum_try_into_method",
43 "Generate an `try_into_` method for an enum variant",
45 fn_name_prefix: "try_into",
47 return_prefix: "Result<",
48 return_suffix: ", Self>",
50 sad_case: "Err(self)",
55 // Assist: generate_enum_as_method
57 // Generate an `as_` method for an enum variant.
73 // fn as_text(&self) -> Option<&String> {
74 // if let Self::Text(v) = self {
82 pub(crate) fn generate_enum_as_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
83 generate_enum_projection_method(
86 "generate_enum_as_method",
87 "Generate an `as_` method for an enum variant",
91 return_prefix: "Option<&",
99 struct ProjectionProps {
100 fn_name_prefix: &'static str,
101 self_param: &'static str,
102 return_prefix: &'static str,
103 return_suffix: &'static str,
104 happy_case: &'static str,
105 sad_case: &'static str,
108 fn generate_enum_projection_method(
111 assist_id: &'static str,
112 assist_description: &str,
113 props: ProjectionProps,
115 let variant = ctx.find_node_at_offset::<ast::Variant>()?;
116 let variant_name = variant.name()?;
117 let parent_enum = ast::Adt::Enum(variant.parent_enum());
119 let (pattern_suffix, field_type, bound_name) = match variant.kind() {
120 ast::StructKind::Record(record) => {
121 let (field,) = record.fields().collect_tuple()?;
122 let name = field.name()?.to_string();
123 let ty = field.ty()?;
124 let pattern_suffix = format!(" {{ {} }}", name);
125 (pattern_suffix, ty, name)
127 ast::StructKind::Tuple(tuple) => {
128 let (field,) = tuple.fields().collect_tuple()?;
129 let ty = field.ty()?;
130 ("(v)".to_owned(), ty, "v".to_owned())
132 ast::StructKind::Unit => return None,
136 format!("{}_{}", props.fn_name_prefix, &to_lower_snake_case(&variant_name.text()));
138 // Return early if we've found an existing new fn
139 let impl_def = find_struct_impl(ctx, &parent_enum, &fn_name)?;
141 let target = variant.syntax().text_range();
142 acc.add(AssistId(assist_id, AssistKind::Generate), assist_description, target, |builder| {
143 let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
144 let method = format!(
145 " {0}fn {1}({2}) -> {3}{4}{5} {{
146 if let Self::{6}{7} = self {{
165 add_method_to_adt(builder, &parent_enum, impl_def, &method);
171 use crate::tests::{check_assist, check_assist_not_applicable};
176 fn test_generate_enum_try_into_tuple_variant() {
178 generate_enum_try_into_method,
190 fn try_into_text(self) -> Result<String, Self> {
191 if let Self::Text(v) = self {
202 fn test_generate_enum_try_into_already_implemented() {
203 check_assist_not_applicable(
204 generate_enum_try_into_method,
211 fn try_into_text(self) -> Result<String, Self> {
212 if let Self::Text(v) = self {
223 fn test_generate_enum_try_into_unit_variant() {
224 check_assist_not_applicable(
225 generate_enum_try_into_method,
235 fn test_generate_enum_try_into_record_with_multiple_fields() {
236 check_assist_not_applicable(
237 generate_enum_try_into_method,
241 Both { first: i32, second: String }$0,
247 fn test_generate_enum_try_into_tuple_with_multiple_fields() {
248 check_assist_not_applicable(
249 generate_enum_try_into_method,
252 Text(String, String)$0,
258 fn test_generate_enum_try_into_record_variant() {
260 generate_enum_try_into_method,
263 Text { text: String }$0,
267 Text { text: String },
271 fn try_into_text(self) -> Result<String, Self> {
272 if let Self::Text { text } = self {
283 fn test_generate_enum_as_tuple_variant() {
285 generate_enum_as_method,
297 fn as_text(&self) -> Option<&String> {
298 if let Self::Text(v) = self {
309 fn test_generate_enum_as_record_variant() {
311 generate_enum_as_method,
314 Text { text: String }$0,
318 Text { text: String },
322 fn as_text(&self) -> Option<&String> {
323 if let Self::Text { text } = self {