crates/assists/src/handlers/generate_enum_match_method.rs

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