crates/ide_assists/src/handlers/generate_enum_match_method.rs

   1 use stdx::{format_to, to_lower_snake_case};
   2 use syntax::ast::VisibilityOwner;
   3 use syntax::ast::{self, AstNode, NameOwner};
   4 use test_utils::mark;
   5
   6 use crate::{
   7     utils::{find_impl_block_end, find_struct_impl, generate_impl_text},
   8     AssistContext, AssistId, AssistKind, Assists,
   9 };
  10
  11 // Assist: generate_enum_is_method
  12 //
  13 // Generate an `is_` method for an 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 //     fn is_minor(&self) -> bool {
  33 //         matches!(self, Self::Minor)
  34 //     }
  35 // }
  36 // ```
  37 pub(crate) fn generate_enum_is_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  38     let variant = ctx.find_node_at_offset::<ast::Variant>()?;
  39     let variant_name = variant.name()?;
  40     let parent_enum = variant.parent_enum();
  41     if !matches!(variant.kind(), ast::StructKind::Unit) {
  42         mark::hit!(test_gen_enum_match_on_non_unit_variant_not_implemented);
  43         return None;
  44     }
  45
  46     let enum_lowercase_name = to_lower_snake_case(&parent_enum.name()?.to_string());
  47     let fn_name = to_lower_snake_case(&variant_name.to_string());
  48
  49     // Return early if we've found an existing new fn
  50     let impl_def = find_struct_impl(
  51         &ctx,
  52         &ast::Adt::Enum(parent_enum.clone()),
  53         format!("is_{}", fn_name).as_str(),
  54     )?;
  55
  56     let target = variant.syntax().text_range();
  57     acc.add(
  58         AssistId("generate_enum_is_method", AssistKind::Generate),
  59         "Generate an `is_` method for an enum variant",
  60         target,
  61         |builder| {
  62             let mut buf = String::with_capacity(512);
  63
  64             if impl_def.is_some() {
  65                 buf.push('\n');
  66             }
  67
  68             let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
  69             format_to!(
  70                 buf,
  71                 "    /// Returns `true` if the {} is [`{}`].
  72     {}fn is_{}(&self) -> bool {{
  73         matches!(self, Self::{})
  74     }}",
  75                 enum_lowercase_name,
  76                 variant_name,
  77                 vis,
  78                 fn_name,
  79                 variant_name
  80             );
  81
  82             let start_offset = impl_def
  83                 .and_then(|impl_def| find_impl_block_end(impl_def, &mut buf))
  84                 .unwrap_or_else(|| {
  85                     buf = generate_impl_text(&ast::Adt::Enum(parent_enum.clone()), &buf);
  86                     parent_enum.syntax().text_range().end()
  87                 });
  88
  89             builder.insert(start_offset, buf);
  90         },
  91     )
  92 }
  93
  94 #[cfg(test)]
  95 mod tests {
  96     use test_utils::mark;
  97
  98     use crate::tests::{check_assist, check_assist_not_applicable};
  99
 100     use super::*;
 101
 102     fn check_not_applicable(ra_fixture: &str) {
 103         check_assist_not_applicable(generate_enum_is_method, ra_fixture)
 104     }
 105
 106     #[test]
 107     fn test_generate_enum_match_from_variant() {
 108         check_assist(
 109             generate_enum_is_method,
 110             r#"
 111 enum Variant {
 112     Undefined,
 113     Minor$0,
 114     Major,
 115 }"#,
 116             r#"enum Variant {
 117     Undefined,
 118     Minor,
 119     Major,
 120 }
 121
 122 impl Variant {
 123     /// Returns `true` if the variant is [`Minor`].
 124     fn is_minor(&self) -> bool {
 125         matches!(self, Self::Minor)
 126     }
 127 }"#,
 128         );
 129     }
 130
 131     #[test]
 132     fn test_generate_enum_match_already_implemented() {
 133         check_not_applicable(
 134             r#"
 135 enum Variant {
 136     Undefined,
 137     Minor$0,
 138     Major,
 139 }
 140
 141 impl Variant {
 142     fn is_minor(&self) -> bool {
 143         matches!(self, Self::Minor)
 144     }
 145 }"#,
 146         );
 147     }
 148
 149     #[test]
 150     fn test_add_from_impl_no_element() {
 151         mark::check!(test_gen_enum_match_on_non_unit_variant_not_implemented);
 152         check_not_applicable(
 153             r#"
 154 enum Variant {
 155     Undefined,
 156     Minor(u32)$0,
 157     Major,
 158 }"#,
 159         );
 160     }
 161
 162     #[test]
 163     fn test_generate_enum_match_from_variant_with_one_variant() {
 164         check_assist(
 165             generate_enum_is_method,
 166             r#"enum Variant { Undefi$0ned }"#,
 167             r#"
 168 enum Variant { Undefined }
 169
 170 impl Variant {
 171     /// Returns `true` if the variant is [`Undefined`].
 172     fn is_undefined(&self) -> bool {
 173         matches!(self, Self::Undefined)
 174     }
 175 }"#,
 176         );
 177     }
 178
 179     #[test]
 180     fn test_generate_enum_match_from_variant_with_visibility_marker() {
 181         check_assist(
 182             generate_enum_is_method,
 183             r#"
 184 pub(crate) enum Variant {
 185     Undefined,
 186     Minor$0,
 187     Major,
 188 }"#,
 189             r#"pub(crate) enum Variant {
 190     Undefined,
 191     Minor,
 192     Major,
 193 }
 194
 195 impl Variant {
 196     /// Returns `true` if the variant is [`Minor`].
 197     pub(crate) fn is_minor(&self) -> bool {
 198         matches!(self, Self::Minor)
 199     }
 200 }"#,
 201         );
 202     }
 203
 204     #[test]
 205     fn test_multiple_generate_enum_match_from_variant() {
 206         check_assist(
 207             generate_enum_is_method,
 208             r#"
 209 enum Variant {
 210     Undefined,
 211     Minor,
 212     Major$0,
 213 }
 214
 215 impl Variant {
 216     /// Returns `true` if the variant is [`Minor`].
 217     fn is_minor(&self) -> bool {
 218         matches!(self, Self::Minor)
 219     }
 220 }"#,
 221             r#"enum Variant {
 222     Undefined,
 223     Minor,
 224     Major,
 225 }
 226
 227 impl Variant {
 228     /// Returns `true` if the variant is [`Minor`].
 229     fn is_minor(&self) -> bool {
 230         matches!(self, Self::Minor)
 231     }
 232
 233     /// Returns `true` if the variant is [`Major`].
 234     fn is_major(&self) -> bool {
 235         matches!(self, Self::Major)
 236     }
 237 }"#,
 238         );
 239     }
 240 }