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

   1 use ide_db::assists::GroupLabel;
   2 use stdx::to_lower_snake_case;
   3 use syntax::ast::HasVisibility;
   4 use syntax::ast::{self, AstNode, HasName};
   5
   6 use crate::{
   7     utils::{add_method_to_adt, find_struct_impl},
   8     AssistContext, AssistId, AssistKind, Assists,
   9 };
  10
  11 // Assist: generate_enum_is_method
  12 //
  13 // Generate an `is_` method for this enum variant.
  14 //
  15 // ```
  16 // enum Version {
  17 //  Undefined,
  18 //  Minor$0,
  19 //  Major,
  20 // }
  21 // ```
  22 // ->
  23 // ```
  24 // enum Version {
  25 //  Undefined,
  26 //  Minor,
  27 //  Major,
  28 // }
  29 //
  30 // impl Version {
  31 //     /// Returns `true` if the version is [`Minor`].
  32 //     ///
  33 //     /// [`Minor`]: Version::Minor
  34 //     #[must_use]
  35 //     fn is_minor(&self) -> bool {
  36 //         matches!(self, Self::Minor)
  37 //     }
  38 // }
  39 // ```
  40 pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  41     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
  42     let variant_name = variant.name()?;
  43     let parent_enum = ast::Adt::Enum(variant.parent_enum());
  44     let pattern_suffix = match variant.kind() {
  45         ast::StructKind::Record(_) => " { .. }",
  46         ast::StructKind::Tuple(_) => "(..)",
  47         ast::StructKind::Unit => "",
  48     };
  49
  50     let enum_name = parent_enum.name()?;
  51     let enum_lowercase_name = to_lower_snake_case(&enum_name.to_string()).replace('_', " ");
  52     let fn_name = format!("is_{}", &to_lower_snake_case(&variant_name.text()));
  53
  54     // Return early if we've found an existing new fn
  55     let impl_def = find_struct_impl(ctx, &parent_enum, &[fn_name.clone()])?;
  56
  57     let target = variant.syntax().text_range();
  58     acc.add_group(
  59         &GroupLabel("Generate an `is_`,`as_`, or `try_into_` for this enum variant".to_owned()),
  60         AssistId("generate_enum_is_method", AssistKind::Generate),
  61         "Generate an `is_` method for this enum variant",
  62         target,
  63         |builder| {
  64             let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{v} "));
  65             let method = format!(
  66                 "    /// Returns `true` if the {enum_lowercase_name} is [`{variant_name}`].
  67     ///
  68     /// [`{variant_name}`]: {enum_name}::{variant_name}
  69     #[must_use]
  70     {vis}fn {fn_name}(&self) -> bool {{
  71         matches!(self, Self::{variant_name}{pattern_suffix})
  72     }}",
  73             );
  74
  75             add_method_to_adt(builder, &parent_enum, impl_def, &method);
  76         },
  77     )
  78 }
  79
  80 #[cfg(test)]
  81 mod tests {
  82     use crate::tests::{check_assist, check_assist_not_applicable};
  83
  84     use super::*;
  85
  86     #[test]
  87     fn test_generate_enum_is_from_variant() {
  88         check_assist(
  89             generate_enum_is_method,
  90             r#"
  91 enum Variant {
  92     Undefined,
  93     Minor$0,
  94     Major,
  95 }"#,
  96             r#"enum Variant {
  97     Undefined,
  98     Minor,
  99     Major,
 100 }
 101
 102 impl Variant {
 103     /// Returns `true` if the variant is [`Minor`].
 104     ///
 105     /// [`Minor`]: Variant::Minor
 106     #[must_use]
 107     fn is_minor(&self) -> bool {
 108         matches!(self, Self::Minor)
 109     }
 110 }"#,
 111         );
 112     }
 113
 114     #[test]
 115     fn test_generate_enum_is_already_implemented() {
 116         check_assist_not_applicable(
 117             generate_enum_is_method,
 118             r#"
 119 enum Variant {
 120     Undefined,
 121     Minor$0,
 122     Major,
 123 }
 124
 125 impl Variant {
 126     fn is_minor(&self) -> bool {
 127         matches!(self, Self::Minor)
 128     }
 129 }"#,
 130         );
 131     }
 132
 133     #[test]
 134     fn test_generate_enum_is_from_tuple_variant() {
 135         check_assist(
 136             generate_enum_is_method,
 137             r#"
 138 enum Variant {
 139     Undefined,
 140     Minor(u32)$0,
 141     Major,
 142 }"#,
 143             r#"enum Variant {
 144     Undefined,
 145     Minor(u32),
 146     Major,
 147 }
 148
 149 impl Variant {
 150     /// Returns `true` if the variant is [`Minor`].
 151     ///
 152     /// [`Minor`]: Variant::Minor
 153     #[must_use]
 154     fn is_minor(&self) -> bool {
 155         matches!(self, Self::Minor(..))
 156     }
 157 }"#,
 158         );
 159     }
 160
 161     #[test]
 162     fn test_generate_enum_is_from_record_variant() {
 163         check_assist(
 164             generate_enum_is_method,
 165             r#"
 166 enum Variant {
 167     Undefined,
 168     Minor { foo: i32 }$0,
 169     Major,
 170 }"#,
 171             r#"enum Variant {
 172     Undefined,
 173     Minor { foo: i32 },
 174     Major,
 175 }
 176
 177 impl Variant {
 178     /// Returns `true` if the variant is [`Minor`].
 179     ///
 180     /// [`Minor`]: Variant::Minor
 181     #[must_use]
 182     fn is_minor(&self) -> bool {
 183         matches!(self, Self::Minor { .. })
 184     }
 185 }"#,
 186         );
 187     }
 188
 189     #[test]
 190     fn test_generate_enum_is_from_variant_with_one_variant() {
 191         check_assist(
 192             generate_enum_is_method,
 193             r#"enum Variant { Undefi$0ned }"#,
 194             r#"
 195 enum Variant { Undefined }
 196
 197 impl Variant {
 198     /// Returns `true` if the variant is [`Undefined`].
 199     ///
 200     /// [`Undefined`]: Variant::Undefined
 201     #[must_use]
 202     fn is_undefined(&self) -> bool {
 203         matches!(self, Self::Undefined)
 204     }
 205 }"#,
 206         );
 207     }
 208
 209     #[test]
 210     fn test_generate_enum_is_from_variant_with_visibility_marker() {
 211         check_assist(
 212             generate_enum_is_method,
 213             r#"
 214 pub(crate) enum Variant {
 215     Undefined,
 216     Minor$0,
 217     Major,
 218 }"#,
 219             r#"pub(crate) enum Variant {
 220     Undefined,
 221     Minor,
 222     Major,
 223 }
 224
 225 impl Variant {
 226     /// Returns `true` if the variant is [`Minor`].
 227     ///
 228     /// [`Minor`]: Variant::Minor
 229     #[must_use]
 230     pub(crate) fn is_minor(&self) -> bool {
 231         matches!(self, Self::Minor)
 232     }
 233 }"#,
 234         );
 235     }
 236
 237     #[test]
 238     fn test_multiple_generate_enum_is_from_variant() {
 239         check_assist(
 240             generate_enum_is_method,
 241             r#"
 242 enum Variant {
 243     Undefined,
 244     Minor,
 245     Major$0,
 246 }
 247
 248 impl Variant {
 249     /// Returns `true` if the variant is [`Minor`].
 250     ///
 251     /// [`Minor`]: Variant::Minor
 252     #[must_use]
 253     fn is_minor(&self) -> bool {
 254         matches!(self, Self::Minor)
 255     }
 256 }"#,
 257             r#"enum Variant {
 258     Undefined,
 259     Minor,
 260     Major,
 261 }
 262
 263 impl Variant {
 264     /// Returns `true` if the variant is [`Minor`].
 265     ///
 266     /// [`Minor`]: Variant::Minor
 267     #[must_use]
 268     fn is_minor(&self) -> bool {
 269         matches!(self, Self::Minor)
 270     }
 271
 272     /// Returns `true` if the variant is [`Major`].
 273     ///
 274     /// [`Major`]: Variant::Major
 275     #[must_use]
 276     fn is_major(&self) -> bool {
 277         matches!(self, Self::Major)
 278     }
 279 }"#,
 280         );
 281     }
 282
 283     #[test]
 284     fn test_generate_enum_is_variant_names() {
 285         check_assist(
 286             generate_enum_is_method,
 287             r#"
 288 enum GeneratorState {
 289     Yielded,
 290     Complete$0,
 291     Major,
 292 }"#,
 293             r#"enum GeneratorState {
 294     Yielded,
 295     Complete,
 296     Major,
 297 }
 298
 299 impl GeneratorState {
 300     /// Returns `true` if the generator state is [`Complete`].
 301     ///
 302     /// [`Complete`]: GeneratorState::Complete
 303     #[must_use]
 304     fn is_complete(&self) -> bool {
 305         matches!(self, Self::Complete)
 306     }
 307 }"#,
 308         );
 309     }
 310 }