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