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
 128     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
 129     let variant_name = variant.name()?;
 130     let parent_enum = ast::Adt::Enum(variant.parent_enum());
 131
 132     let (pattern_suffix, field_type, bound_name) = match variant.kind() {
 133         ast::StructKind::Record(record) => {
 134             let (field,) = record.fields().collect_tuple()?;
 135             let name = field.name()?.to_string();
 136             let ty = field.ty()?;
 137             let pattern_suffix = format!(" {{ {name} }}");
 138             (pattern_suffix, ty, name)
 139         }
 140         ast::StructKind::Tuple(tuple) => {
 141             let (field,) = tuple.fields().collect_tuple()?;
 142             let ty = field.ty()?;
 143             ("(v)".to_owned(), ty, "v".to_owned())
 144         }
 145         ast::StructKind::Unit => return None,
 146     };
 147
 148     let fn_name = format!("{fn_name_prefix}_{}", &to_lower_snake_case(&variant_name.text()));
 149
 150     // Return early if we've found an existing new fn
 151     let impl_def = find_struct_impl(ctx, &parent_enum, &[fn_name.clone()])?;
 152
 153     let target = variant.syntax().text_range();
 154     acc.add_group(
 155         &GroupLabel("Generate an `is_`,`as_`, or `try_into_` for this enum variant".to_owned()),
 156         AssistId(assist_id, AssistKind::Generate),
 157         assist_description,
 158         target,
 159         |builder| {
 160             let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
 161
 162             let field_type_syntax = field_type.syntax();
 163
 164             let must_use = if ctx.config.assist_emit_must_use {
 165                 "#[must_use]\n    "
 166             } else {
 167                 ""
 168             };
 169
 170             let method = format!(
 171                 "    {must_use}{vis}fn {fn_name}({self_param}) -> {return_prefix}{field_type_syntax}{return_suffix} {{
 172         if let Self::{variant_name}{pattern_suffix} = self {{
 173             {happy_case}({bound_name})
 174         }} else {{
 175             {sad_case}
 176         }}
 177     }}"
 178             );
 179
 180             add_method_to_adt(builder, &parent_enum, impl_def, &method);
 181         },
 182     )
 183 }
 184
 185 #[cfg(test)]
 186 mod tests {
 187     use crate::tests::{check_assist, check_assist_not_applicable};
 188
 189     use super::*;
 190
 191     #[test]
 192     fn test_generate_enum_try_into_tuple_variant() {
 193         check_assist(
 194             generate_enum_try_into_method,
 195             r#"
 196 enum Value {
 197     Number(i32),
 198     Text(String)$0,
 199 }"#,
 200             r#"enum Value {
 201     Number(i32),
 202     Text(String),
 203 }
 204
 205 impl Value {
 206     fn try_into_text(self) -> Result<String, Self> {
 207         if let Self::Text(v) = self {
 208             Ok(v)
 209         } else {
 210             Err(self)
 211         }
 212     }
 213 }"#,
 214         );
 215     }
 216
 217     #[test]
 218     fn test_generate_enum_try_into_already_implemented() {
 219         check_assist_not_applicable(
 220             generate_enum_try_into_method,
 221             r#"enum Value {
 222     Number(i32),
 223     Text(String)$0,
 224 }
 225
 226 impl Value {
 227     fn try_into_text(self) -> Result<String, Self> {
 228         if let Self::Text(v) = self {
 229             Ok(v)
 230         } else {
 231             Err(self)
 232         }
 233     }
 234 }"#,
 235         );
 236     }
 237
 238     #[test]
 239     fn test_generate_enum_try_into_unit_variant() {
 240         check_assist_not_applicable(
 241             generate_enum_try_into_method,
 242             r#"enum Value {
 243     Number(i32),
 244     Text(String),
 245     Unit$0,
 246 }"#,
 247         );
 248     }
 249
 250     #[test]
 251     fn test_generate_enum_try_into_record_with_multiple_fields() {
 252         check_assist_not_applicable(
 253             generate_enum_try_into_method,
 254             r#"enum Value {
 255     Number(i32),
 256     Text(String),
 257     Both { first: i32, second: String }$0,
 258 }"#,
 259         );
 260     }
 261
 262     #[test]
 263     fn test_generate_enum_try_into_tuple_with_multiple_fields() {
 264         check_assist_not_applicable(
 265             generate_enum_try_into_method,
 266             r#"enum Value {
 267     Number(i32),
 268     Text(String, String)$0,
 269 }"#,
 270         );
 271     }
 272
 273     #[test]
 274     fn test_generate_enum_try_into_record_variant() {
 275         check_assist(
 276             generate_enum_try_into_method,
 277             r#"enum Value {
 278     Number(i32),
 279     Text { text: String }$0,
 280 }"#,
 281             r#"enum Value {
 282     Number(i32),
 283     Text { text: String },
 284 }
 285
 286 impl Value {
 287     fn try_into_text(self) -> Result<String, Self> {
 288         if let Self::Text { text } = self {
 289             Ok(text)
 290         } else {
 291             Err(self)
 292         }
 293     }
 294 }"#,
 295         );
 296     }
 297
 298     #[test]
 299     fn test_generate_enum_as_tuple_variant() {
 300         check_assist(
 301             generate_enum_as_method,
 302             r#"
 303 enum Value {
 304     Number(i32),
 305     Text(String)$0,
 306 }"#,
 307             r#"enum Value {
 308     Number(i32),
 309     Text(String),
 310 }
 311
 312 impl Value {
 313     fn as_text(&self) -> Option<&String> {
 314         if let Self::Text(v) = self {
 315             Some(v)
 316         } else {
 317             None
 318         }
 319     }
 320 }"#,
 321         );
 322     }
 323
 324     #[test]
 325     fn test_generate_enum_as_record_variant() {
 326         check_assist(
 327             generate_enum_as_method,
 328             r#"enum Value {
 329     Number(i32),
 330     Text { text: String }$0,
 331 }"#,
 332             r#"enum Value {
 333     Number(i32),
 334     Text { text: String },
 335 }
 336
 337 impl Value {
 338     fn as_text(&self) -> Option<&String> {
 339         if let Self::Text { text } = self {
 340             Some(text)
 341         } else {
 342             None
 343         }
 344     }
 345 }"#,
 346         );
 347     }
 348 }