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)?;
  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 {} is [`{variant}`].
  67     ///
  68     /// [`{variant}`]: {}::{variant}
  69     #[must_use]
  70     {}fn {}(&self) -> bool {{
  71         matches!(self, Self::{variant}{})
  72     }}",
  73                 enum_lowercase_name,
  74                 enum_name,
  75                 vis,
  76                 fn_name,
  77                 pattern_suffix,
  78                 variant = variant_name
  79             );
  80
  81             add_method_to_adt(builder, &parent_enum, impl_def, &method);
  82         },
  83     )
  84 }
  85
  86 #[cfg(test)]
  87 mod tests {
  88     use crate::tests::{check_assist, check_assist_not_applicable};
  89
  90     use super::*;
  91
  92     #[test]
  93     fn test_generate_enum_is_from_variant() {
  94         check_assist(
  95             generate_enum_is_method,
  96             r#"
  97 enum Variant {
  98     Undefined,
  99     Minor$0,
 100     Major,
 101 }"#,
 102             r#"enum Variant {
 103     Undefined,
 104     Minor,
 105     Major,
 106 }
 107
 108 impl Variant {
 109     /// Returns `true` if the variant is [`Minor`].
 110     ///
 111     /// [`Minor`]: Variant::Minor
 112     #[must_use]
 113     fn is_minor(&self) -> bool {
 114         matches!(self, Self::Minor)
 115     }
 116 }"#,
 117         );
 118     }
 119
 120     #[test]
 121     fn test_generate_enum_is_already_implemented() {
 122         check_assist_not_applicable(
 123             generate_enum_is_method,
 124             r#"
 125 enum Variant {
 126     Undefined,
 127     Minor$0,
 128     Major,
 129 }
 130
 131 impl Variant {
 132     fn is_minor(&self) -> bool {
 133         matches!(self, Self::Minor)
 134     }
 135 }"#,
 136         );
 137     }
 138
 139     #[test]
 140     fn test_generate_enum_is_from_tuple_variant() {
 141         check_assist(
 142             generate_enum_is_method,
 143             r#"
 144 enum Variant {
 145     Undefined,
 146     Minor(u32)$0,
 147     Major,
 148 }"#,
 149             r#"enum Variant {
 150     Undefined,
 151     Minor(u32),
 152     Major,
 153 }
 154
 155 impl Variant {
 156     /// Returns `true` if the variant is [`Minor`].
 157     ///
 158     /// [`Minor`]: Variant::Minor
 159     #[must_use]
 160     fn is_minor(&self) -> bool {
 161         matches!(self, Self::Minor(..))
 162     }
 163 }"#,
 164         );
 165     }
 166
 167     #[test]
 168     fn test_generate_enum_is_from_record_variant() {
 169         check_assist(
 170             generate_enum_is_method,
 171             r#"
 172 enum Variant {
 173     Undefined,
 174     Minor { foo: i32 }$0,
 175     Major,
 176 }"#,
 177             r#"enum Variant {
 178     Undefined,
 179     Minor { foo: i32 },
 180     Major,
 181 }
 182
 183 impl Variant {
 184     /// Returns `true` if the variant is [`Minor`].
 185     ///
 186     /// [`Minor`]: Variant::Minor
 187     #[must_use]
 188     fn is_minor(&self) -> bool {
 189         matches!(self, Self::Minor { .. })
 190     }
 191 }"#,
 192         );
 193     }
 194
 195     #[test]
 196     fn test_generate_enum_is_from_variant_with_one_variant() {
 197         check_assist(
 198             generate_enum_is_method,
 199             r#"enum Variant { Undefi$0ned }"#,
 200             r#"
 201 enum Variant { Undefined }
 202
 203 impl Variant {
 204     /// Returns `true` if the variant is [`Undefined`].
 205     ///
 206     /// [`Undefined`]: Variant::Undefined
 207     #[must_use]
 208     fn is_undefined(&self) -> bool {
 209         matches!(self, Self::Undefined)
 210     }
 211 }"#,
 212         );
 213     }
 214
 215     #[test]
 216     fn test_generate_enum_is_from_variant_with_visibility_marker() {
 217         check_assist(
 218             generate_enum_is_method,
 219             r#"
 220 pub(crate) enum Variant {
 221     Undefined,
 222     Minor$0,
 223     Major,
 224 }"#,
 225             r#"pub(crate) enum Variant {
 226     Undefined,
 227     Minor,
 228     Major,
 229 }
 230
 231 impl Variant {
 232     /// Returns `true` if the variant is [`Minor`].
 233     ///
 234     /// [`Minor`]: Variant::Minor
 235     #[must_use]
 236     pub(crate) fn is_minor(&self) -> bool {
 237         matches!(self, Self::Minor)
 238     }
 239 }"#,
 240         );
 241     }
 242
 243     #[test]
 244     fn test_multiple_generate_enum_is_from_variant() {
 245         check_assist(
 246             generate_enum_is_method,
 247             r#"
 248 enum Variant {
 249     Undefined,
 250     Minor,
 251     Major$0,
 252 }
 253
 254 impl Variant {
 255     /// Returns `true` if the variant is [`Minor`].
 256     ///
 257     /// [`Minor`]: Variant::Minor
 258     #[must_use]
 259     fn is_minor(&self) -> bool {
 260         matches!(self, Self::Minor)
 261     }
 262 }"#,
 263             r#"enum Variant {
 264     Undefined,
 265     Minor,
 266     Major,
 267 }
 268
 269 impl Variant {
 270     /// Returns `true` if the variant is [`Minor`].
 271     ///
 272     /// [`Minor`]: Variant::Minor
 273     #[must_use]
 274     fn is_minor(&self) -> bool {
 275         matches!(self, Self::Minor)
 276     }
 277
 278     /// Returns `true` if the variant is [`Major`].
 279     ///
 280     /// [`Major`]: Variant::Major
 281     #[must_use]
 282     fn is_major(&self) -> bool {
 283         matches!(self, Self::Major)
 284     }
 285 }"#,
 286         );
 287     }
 288
 289     #[test]
 290     fn test_generate_enum_is_variant_names() {
 291         check_assist(
 292             generate_enum_is_method,
 293             r#"
 294 enum GeneratorState {
 295     Yielded,
 296     Complete$0,
 297     Major,
 298 }"#,
 299             r#"enum GeneratorState {
 300     Yielded,
 301     Complete,
 302     Major,
 303 }
 304
 305 impl GeneratorState {
 306     /// Returns `true` if the generator state is [`Complete`].
 307     ///
 308     /// [`Complete`]: GeneratorState::Complete
 309     #[must_use]
 310     fn is_complete(&self) -> bool {
 311         matches!(self, Self::Complete)
 312     }
 313 }"#,
 314         );
 315     }
 316 }