crates/assists/src/handlers/generate_enum_match_method.rs

   1 use hir::Adt;
   2 use stdx::format_to;
   3 use syntax::ast::{self, AstNode, NameOwner};
   4 use syntax::{ast::VisibilityOwner, T};
   5 use test_utils::mark;
   6
   7 use crate::{AssistContext, AssistId, AssistKind, Assists};
   8
   9 // Assist: generate_enum_match_method
  10 //
  11 // Generate an `is_` method for an enum variant.
  12 //
  13 // ```
  14 // enum Version {
  15 //  Undefined,
  16 //  Minor$0,
  17 //  Major,
  18 // }
  19 // ```
  20 // ->
  21 // ```
  22 // enum Version {
  23 //  Undefined,
  24 //  Minor,
  25 //  Major,
  26 // }
  27 //
  28 // impl Version {
  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 fn_name = to_lower_snake_case(&format!("{}", variant_name));
  44
  45     // Return early if we've found an existing new fn
  46     let impl_def = find_struct_impl(&ctx, &parent_enum, format!("is_{}", fn_name).as_str())?;
  47
  48     let target = variant.syntax().text_range();
  49     acc.add(
  50         AssistId("generate_enum_match_method", AssistKind::Generate),
  51         "Generate an `is_` method for an enum variant",
  52         target,
  53         |builder| {
  54             let mut buf = String::with_capacity(512);
  55
  56             if impl_def.is_some() {
  57                 buf.push('\n');
  58             }
  59
  60             let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
  61
  62             format_to!(
  63                 buf,
  64                 "    {}fn is_{}(&self) -> bool {{
  65         matches!(self, Self::{})
  66     }}",
  67                 vis,
  68                 fn_name,
  69                 variant_name
  70             );
  71
  72             let start_offset = impl_def
  73                 .and_then(|impl_def| {
  74                     buf.push('\n');
  75                     let start = impl_def
  76                         .syntax()
  77                         .descendants_with_tokens()
  78                         .find(|t| t.kind() == T!['{'])?
  79                         .text_range()
  80                         .end();
  81
  82                     Some(start)
  83                 })
  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(" ");
 100     buf.push_str(strukt.name().unwrap().text());
 101     format_to!(buf, " {{\n{}\n}}", code);
 102     buf
 103 }
 104
 105 fn to_lower_snake_case(s: &str) -> String {
 106     let mut buf = String::with_capacity(s.len());
 107     let mut prev = false;
 108     for c in s.chars() {
 109         if c.is_ascii_uppercase() && prev {
 110             buf.push('_')
 111         }
 112         prev = true;
 113
 114         buf.push(c.to_ascii_lowercase());
 115     }
 116     buf
 117 }
 118
 119 // Uses a syntax-driven approach to find any impl blocks for the struct that
 120 // exist within the module/file
 121 //
 122 // Returns `None` if we've found an existing `new` fn
 123 //
 124 // FIXME: change the new fn checking to a more semantic approach when that's more
 125 // viable (e.g. we process proc macros, etc)
 126 fn find_struct_impl(
 127     ctx: &AssistContext,
 128     strukt: &ast::Enum,
 129     name: &str,
 130 ) -> Option<Option<ast::Impl>> {
 131     let db = ctx.db();
 132     let module = strukt.syntax().ancestors().find(|node| {
 133         ast::Module::can_cast(node.kind()) || ast::SourceFile::can_cast(node.kind())
 134     })?;
 135
 136     let struct_def = ctx.sema.to_def(strukt)?;
 137
 138     let block = module.descendants().filter_map(ast::Impl::cast).find_map(|impl_blk| {
 139         let blk = ctx.sema.to_def(&impl_blk)?;
 140
 141         // FIXME: handle e.g. `struct S<T>; impl<U> S<U> {}`
 142         // (we currently use the wrong type parameter)
 143         // also we wouldn't want to use e.g. `impl S<u32>`
 144         let same_ty = match blk.target_ty(db).as_adt() {
 145             Some(def) => def == Adt::Enum(struct_def),
 146             None => false,
 147         };
 148         let not_trait_impl = blk.target_trait(db).is_none();
 149
 150         if !(same_ty && not_trait_impl) {
 151             None
 152         } else {
 153             Some(impl_blk)
 154         }
 155     });
 156
 157     if let Some(ref impl_blk) = block {
 158         if has_fn(impl_blk, name) {
 159             mark::hit!(test_gen_enum_match_impl_already_exists);
 160             return None;
 161         }
 162     }
 163
 164     Some(block)
 165 }
 166
 167 fn has_fn(imp: &ast::Impl, rhs_name: &str) -> bool {
 168     if let Some(il) = imp.assoc_item_list() {
 169         for item in il.assoc_items() {
 170             if let ast::AssocItem::Fn(f) = item {
 171                 if let Some(name) = f.name() {
 172                     if name.text().eq_ignore_ascii_case(rhs_name) {
 173                         return true;
 174                     }
 175                 }
 176             }
 177         }
 178     }
 179
 180     false
 181 }
 182
 183 #[cfg(test)]
 184 mod tests {
 185     use ide_db::helpers::FamousDefs;
 186     use test_utils::mark;
 187
 188     use crate::tests::{check_assist, check_assist_not_applicable};
 189
 190     use super::*;
 191
 192     fn check_not_applicable(ra_fixture: &str) {
 193         let fixture =
 194             format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
 195         check_assist_not_applicable(generate_enum_match_method, &fixture)
 196     }
 197
 198     #[test]
 199     fn test_generate_enum_match_from_variant() {
 200         check_assist(
 201             generate_enum_match_method,
 202             r#"
 203 enum Variant {
 204     Undefined,
 205     Minor$0,
 206     Major,
 207 }"#,
 208             r#"enum Variant {
 209     Undefined,
 210     Minor,
 211     Major,
 212 }
 213
 214 impl Variant {
 215     fn is_minor(&self) -> bool {
 216         matches!(self, Self::Minor)
 217     }
 218 }"#,
 219         );
 220     }
 221
 222     #[test]
 223     fn test_generate_enum_match_already_implemented() {
 224         mark::check!(test_gen_enum_match_impl_already_exists);
 225         check_not_applicable(
 226             r#"
 227 enum Variant {
 228     Undefined,
 229     Minor$0,
 230     Major,
 231 }
 232
 233 impl Variant {
 234     fn is_minor(&self) -> bool {
 235         matches!(self, Self::Minor)
 236     }
 237 }"#,
 238         );
 239     }
 240
 241     #[test]
 242     fn test_add_from_impl_no_element() {
 243         mark::check!(test_gen_enum_match_on_non_unit_variant_not_implemented);
 244         check_not_applicable(
 245             r#"
 246 enum Variant {
 247     Undefined,
 248     Minor(u32)$0,
 249     Major,
 250 }"#,
 251         );
 252     }
 253
 254     #[test]
 255     fn test_generate_enum_match_from_variant_with_one_variant() {
 256         check_assist(
 257             generate_enum_match_method,
 258             r#"enum Variant { Undefi$0ned }"#,
 259             r#"
 260 enum Variant { Undefined }
 261
 262 impl Variant {
 263     fn is_undefined(&self) -> bool {
 264         matches!(self, Self::Undefined)
 265     }
 266 }"#,
 267         );
 268     }
 269
 270     #[test]
 271     fn test_generate_enum_match_from_variant_with_visibility_marker() {
 272         check_assist(
 273             generate_enum_match_method,
 274             r#"
 275 pub(crate) enum Variant {
 276     Undefined,
 277     Minor$0,
 278     Major,
 279 }"#,
 280             r#"pub(crate) enum Variant {
 281     Undefined,
 282     Minor,
 283     Major,
 284 }
 285
 286 impl Variant {
 287     pub(crate) fn is_minor(&self) -> bool {
 288         matches!(self, Self::Minor)
 289     }
 290 }"#,
 291         );
 292     }
 293 }