crates/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, find_struct_impl},
   8     AssistContext, AssistId, AssistKind, Assists,
   9 };
  10
  11 // Assist: generate_enum_match_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_match_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_match_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(impl_def, &mut buf))
  84                 .unwrap_or_else(|| {
  85                     buf = generate_impl_text(&parent_enum, &buf);
  86                     parent_enum.syntax().text_range().end()
  87                 });
  88
  89             builder.insert(start_offset, buf);
  90         },
  91     )
  92 }
  93
  94 // Generates the surrounding `impl Type { <code> }` including type and lifetime
  95 // parameters
  96 fn generate_impl_text(strukt: &ast::Enum, code: &str) -> String {
  97     let mut buf = String::with_capacity(code.len());
  98     buf.push_str("\n\nimpl ");
  99     buf.push_str(strukt.name().unwrap().text());
 100     format_to!(buf, " {{\n{}\n}}", code);
 101     buf
 102 }
 103
 104 #[cfg(test)]
 105 mod tests {
 106     use test_utils::mark;
 107
 108     use crate::tests::{check_assist, check_assist_not_applicable};
 109
 110     use super::*;
 111
 112     fn check_not_applicable(ra_fixture: &str) {
 113         check_assist_not_applicable(generate_enum_match_method, ra_fixture)
 114     }
 115
 116     #[test]
 117     fn test_generate_enum_match_from_variant() {
 118         check_assist(
 119             generate_enum_match_method,
 120             r#"
 121 enum Variant {
 122     Undefined,
 123     Minor$0,
 124     Major,
 125 }"#,
 126             r#"enum Variant {
 127     Undefined,
 128     Minor,
 129     Major,
 130 }
 131
 132 impl Variant {
 133     /// Returns `true` if the variant is [`Minor`].
 134     fn is_minor(&self) -> bool {
 135         matches!(self, Self::Minor)
 136     }
 137 }"#,
 138         );
 139     }
 140
 141     #[test]
 142     fn test_generate_enum_match_already_implemented() {
 143         check_not_applicable(
 144             r#"
 145 enum Variant {
 146     Undefined,
 147     Minor$0,
 148     Major,
 149 }
 150
 151 impl Variant {
 152     fn is_minor(&self) -> bool {
 153         matches!(self, Self::Minor)
 154     }
 155 }"#,
 156         );
 157     }
 158
 159     #[test]
 160     fn test_add_from_impl_no_element() {
 161         mark::check!(test_gen_enum_match_on_non_unit_variant_not_implemented);
 162         check_not_applicable(
 163             r#"
 164 enum Variant {
 165     Undefined,
 166     Minor(u32)$0,
 167     Major,
 168 }"#,
 169         );
 170     }
 171
 172     #[test]
 173     fn test_generate_enum_match_from_variant_with_one_variant() {
 174         check_assist(
 175             generate_enum_match_method,
 176             r#"enum Variant { Undefi$0ned }"#,
 177             r#"
 178 enum Variant { Undefined }
 179
 180 impl Variant {
 181     /// Returns `true` if the variant is [`Undefined`].
 182     fn is_undefined(&self) -> bool {
 183         matches!(self, Self::Undefined)
 184     }
 185 }"#,
 186         );
 187     }
 188
 189     #[test]
 190     fn test_generate_enum_match_from_variant_with_visibility_marker() {
 191         check_assist(
 192             generate_enum_match_method,
 193             r#"
 194 pub(crate) enum Variant {
 195     Undefined,
 196     Minor$0,
 197     Major,
 198 }"#,
 199             r#"pub(crate) enum Variant {
 200     Undefined,
 201     Minor,
 202     Major,
 203 }
 204
 205 impl Variant {
 206     /// Returns `true` if the variant is [`Minor`].
 207     pub(crate) fn is_minor(&self) -> bool {
 208         matches!(self, Self::Minor)
 209     }
 210 }"#,
 211         );
 212     }
 213 }