src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_enum_projection_method.rs

   1 use ide_db::assists::GroupLabel;
   2 use itertools::Itertools;
   3 use stdx::to_lower_snake_case;
   4 use syntax::ast::HasVisibility;
   5 use syntax::ast::{self, AstNode, HasName};
   6
   7 use crate::{
   8     utils::{add_method_to_adt, find_struct_impl},
   9     AssistContext, AssistId, AssistKind, Assists,
  10 };
  11
  12 // Assist: generate_enum_try_into_method
  13 //
  14 // Generate a `try_into_` method for this enum variant.
  15 //
  16 // ```
  17 // enum Value {
  18 //  Number(i32),
  19 //  Text(String)$0,
  20 // }
  21 // ```
  22 // ->
  23 // ```
  24 // enum Value {
  25 //  Number(i32),
  26 //  Text(String),
  27 // }
  28 //
  29 // impl Value {
  30 //     fn try_into_text(self) -> Result<String, Self> {
  31 //         if let Self::Text(v) = self {
  32 //             Ok(v)
  33 //         } else {
  34 //             Err(self)
  35 //         }
  36 //     }
  37 // }
  38 // ```
  39 pub(crate) fn generate_enum_try_into_method(
  40     acc: &mut Assists,
  41     ctx: &AssistContext<'_>,
  42 ) -> Option<()> {
  43     generate_enum_projection_method(
  44         acc,
  45         ctx,
  46         "generate_enum_try_into_method",
  47         "Generate a `try_into_` method for this enum variant",
  48         ProjectionProps {
  49             fn_name_prefix: "try_into",
  50             self_param: "self",
  51             return_prefix: "Result<",
  52             return_suffix: ", Self>",
  53             happy_case: "Ok",
  54             sad_case: "Err(self)",
  55         },
  56     )
  57 }
  58
  59 // Assist: generate_enum_as_method
  60 //
  61 // Generate an `as_` method for this enum variant.
  62 //
  63 // ```
  64 // enum Value {
  65 //  Number(i32),
  66 //  Text(String)$0,
  67 // }
  68 // ```
  69 // ->
  70 // ```
  71 // enum Value {
  72 //  Number(i32),
  73 //  Text(String),
  74 // }
  75 //
  76 // impl Value {
  77 //     fn as_text(&self) -> Option<&String> {
  78 //         if let Self::Text(v) = self {
  79 //             Some(v)
  80 //         } else {
  81 //             None
  82 //         }
  83 //     }
  84 // }
  85 // ```
  86 pub(crate) fn generate_enum_as_method(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  87     generate_enum_projection_method(
  88         acc,
  89         ctx,
  90         "generate_enum_as_method",
  91         "Generate an `as_` method for this enum variant",
  92         ProjectionProps {
  93             fn_name_prefix: "as",
  94             self_param: "&self",
  95             return_prefix: "Option<&",
  96             return_suffix: ">",
  97             happy_case: "Some",
  98             sad_case: "None",
  99         },
 100     )
 101 }
 102
 103 struct ProjectionProps {
 104     fn_name_prefix: &'static str,
 105     self_param: &'static str,
 106     return_prefix: &'static str,
 107     return_suffix: &'static str,
 108     happy_case: &'static str,
 109     sad_case: &'static str,
 110 }
 111
 112 fn generate_enum_projection_method(
 113     acc: &mut Assists,
 114     ctx: &AssistContext<'_>,
 115     assist_id: &'static str,
 116     assist_description: &str,
 117     props: ProjectionProps,
 118 ) -> Option<()> {
 119     let ProjectionProps {
 120         fn_name_prefix,
 121         self_param,
 122         return_prefix,
 123         return_suffix,
 124         happy_case,
 125         sad_case,
 126     } = props;
 127     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
 128     let variant_name = variant.name()?;
 129     let parent_enum = ast::Adt::Enum(variant.parent_enum());
 130
 131     let (pattern_suffix, field_type, bound_name) = match variant.kind() {
 132         ast::StructKind::Record(record) => {
 133             let (field,) = record.fields().collect_tuple()?;
 134             let name = field.name()?.to_string();
 135             let ty = field.ty()?;
 136             let pattern_suffix = format!(" {{ {name} }}");
 137             (pattern_suffix, ty, name)
 138         }
 139         ast::StructKind::Tuple(tuple) => {
 140             let (field,) = tuple.fields().collect_tuple()?;
 141             let ty = field.ty()?;
 142             ("(v)".to_owned(), ty, "v".to_owned())
 143         }
 144         ast::StructKind::Unit => return None,
 145     };
 146
 147     let fn_name = format!("{}_{}", fn_name_prefix, &to_lower_snake_case(&variant_name.text()));
 148
 149     // Return early if we've found an existing new fn
 150     let impl_def = find_struct_impl(ctx, &parent_enum, &[fn_name.clone()])?;
 151
 152     let target = variant.syntax().text_range();
 153     acc.add_group(
 154         &GroupLabel("Generate an `is_`,`as_`, or `try_into_` for this enum variant".to_owned()),
 155         AssistId(assist_id, AssistKind::Generate),
 156         assist_description,
 157         target,
 158         |builder| {
 159             let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{v} "));
 160             let method = format!(
 161                 "    {vis}fn {fn_name}({self_param}) -> {return_prefix}{field_type}{return_suffix} {{
 162         if let Self::{variant_name}{pattern_suffix} = self {{
 163             {happy_case}({bound_name})
 164         }} else {{
 165             {sad_case}
 166         }}
 167     }}");
 168
 169             add_method_to_adt(builder, &parent_enum, impl_def, &method);
 170         },
 171     )
 172 }
 173
 174 #[cfg(test)]
 175 mod tests {
 176     use crate::tests::{check_assist, check_assist_not_applicable};
 177
 178     use super::*;
 179
 180     #[test]
 181     fn test_generate_enum_try_into_tuple_variant() {
 182         check_assist(
 183             generate_enum_try_into_method,
 184             r#"
 185 enum Value {
 186     Number(i32),
 187     Text(String)$0,
 188 }"#,
 189             r#"enum Value {
 190     Number(i32),
 191     Text(String),
 192 }
 193
 194 impl Value {
 195     fn try_into_text(self) -> Result<String, Self> {
 196         if let Self::Text(v) = self {
 197             Ok(v)
 198         } else {
 199             Err(self)
 200         }
 201     }
 202 }"#,
 203         );
 204     }
 205
 206     #[test]
 207     fn test_generate_enum_try_into_already_implemented() {
 208         check_assist_not_applicable(
 209             generate_enum_try_into_method,
 210             r#"enum Value {
 211     Number(i32),
 212     Text(String)$0,
 213 }
 214
 215 impl Value {
 216     fn try_into_text(self) -> Result<String, Self> {
 217         if let Self::Text(v) = self {
 218             Ok(v)
 219         } else {
 220             Err(self)
 221         }
 222     }
 223 }"#,
 224         );
 225     }
 226
 227     #[test]
 228     fn test_generate_enum_try_into_unit_variant() {
 229         check_assist_not_applicable(
 230             generate_enum_try_into_method,
 231             r#"enum Value {
 232     Number(i32),
 233     Text(String),
 234     Unit$0,
 235 }"#,
 236         );
 237     }
 238
 239     #[test]
 240     fn test_generate_enum_try_into_record_with_multiple_fields() {
 241         check_assist_not_applicable(
 242             generate_enum_try_into_method,
 243             r#"enum Value {
 244     Number(i32),
 245     Text(String),
 246     Both { first: i32, second: String }$0,
 247 }"#,
 248         );
 249     }
 250
 251     #[test]
 252     fn test_generate_enum_try_into_tuple_with_multiple_fields() {
 253         check_assist_not_applicable(
 254             generate_enum_try_into_method,
 255             r#"enum Value {
 256     Number(i32),
 257     Text(String, String)$0,
 258 }"#,
 259         );
 260     }
 261
 262     #[test]
 263     fn test_generate_enum_try_into_record_variant() {
 264         check_assist(
 265             generate_enum_try_into_method,
 266             r#"enum Value {
 267     Number(i32),
 268     Text { text: String }$0,
 269 }"#,
 270             r#"enum Value {
 271     Number(i32),
 272     Text { text: String },
 273 }
 274
 275 impl Value {
 276     fn try_into_text(self) -> Result<String, Self> {
 277         if let Self::Text { text } = self {
 278             Ok(text)
 279         } else {
 280             Err(self)
 281         }
 282     }
 283 }"#,
 284         );
 285     }
 286
 287     #[test]
 288     fn test_generate_enum_as_tuple_variant() {
 289         check_assist(
 290             generate_enum_as_method,
 291             r#"
 292 enum Value {
 293     Number(i32),
 294     Text(String)$0,
 295 }"#,
 296             r#"enum Value {
 297     Number(i32),
 298     Text(String),
 299 }
 300
 301 impl Value {
 302     fn as_text(&self) -> Option<&String> {
 303         if let Self::Text(v) = self {
 304             Some(v)
 305         } else {
 306             None
 307         }
 308     }
 309 }"#,
 310         );
 311     }
 312
 313     #[test]
 314     fn test_generate_enum_as_record_variant() {
 315         check_assist(
 316             generate_enum_as_method,
 317             r#"enum Value {
 318     Number(i32),
 319     Text { text: String }$0,
 320 }"#,
 321             r#"enum Value {
 322     Number(i32),
 323     Text { text: String },
 324 }
 325
 326 impl Value {
 327     fn as_text(&self) -> Option<&String> {
 328         if let Self::Text { text } = self {
 329             Some(text)
 330         } else {
 331             None
 332         }
 333     }
 334 }"#,
 335         );
 336     }
 337 }