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