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