add `generate-enum-match` assist

3 files changed, 322 insertions, 0 deletions

diff --git a/crates/assists/src/handlers/generate_enum_match_method.rs b/crates/assists/src/handlers/generate_enum_match_method.rs
new file mode 100644
index 000000000..079ed27bd
--- /dev/null
+++ b/crates/assists/src/handlers/generate_enum_match_method.rs
@@ -0,0 +1,293 @@
+use hir::Adt;
+use stdx::format_to;
+use syntax::ast::{self, AstNode, NameOwner};
+use syntax::{ast::VisibilityOwner, T};
+use test_utils::mark;
+
+use crate::{AssistContext, AssistId, AssistKind, Assists};
+
+// Assist: generate_enum_match_method
+//
+// Generate an `is_` method for an enum variant.
+//
+// ```
+// enum Version {
+//  Undefined,
+//  Minor$0,
+//  Major,
+// }
+// ```
+// ->
+// ```
+// enum Version {
+//  Undefined,
+//  Minor,
+//  Major,
+// }
+//
+// impl Version {
+//     fn is_minor(&self) -> bool {
+//         matches!(self, Self::Minor)
+//     }
+// }
+// ```
+pub(crate) fn generate_enum_match_method(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
+    let variant = ctx.find_node_at_offset::<ast::Variant>()?;
+    let variant_name = variant.name()?;
+    let parent_enum = variant.parent_enum();
+    if !matches!(variant.kind(), ast::StructKind::Unit) {
+        mark::hit!(test_gen_enum_match_on_non_unit_variant_not_implemented);
+        return None;
+    }
+
+    let fn_name = to_lower_snake_case(&format!("{}", variant_name));
+
+    // Return early if we've found an existing new fn
+    let impl_def = find_struct_impl(&ctx, &parent_enum, format!("is_{}", fn_name).as_str())?;
+
+    let target = variant.syntax().text_range();
+    acc.add(
+        AssistId("generate_enum_match_method", AssistKind::Generate),
+        "Generate an `is_` method for an enum variant",
+        target,
+        |builder| {
+            let mut buf = String::with_capacity(512);
+
+            if impl_def.is_some() {
+                buf.push('\n');
+            }
+
+            let vis = parent_enum.visibility().map_or(String::new(), |v| format!("{} ", v));
+
+            format_to!(
+                buf,
+                "    {}fn is_{}(&self) -> bool {{
+        matches!(self, Self::{})
+    }}",
+                vis,
+                fn_name,
+                variant_name
+            );
+
+            let start_offset = impl_def
+                .and_then(|impl_def| {
+                    buf.push('\n');
+                    let start = impl_def
+                        .syntax()
+                        .descendants_with_tokens()
+                        .find(|t| t.kind() == T!['{'])?
+                        .text_range()
+                        .end();
+
+                    Some(start)
+                })
+                .unwrap_or_else(|| {
+                    buf = generate_impl_text(&parent_enum, &buf);
+                    parent_enum.syntax().text_range().end()
+                });
+
+            builder.insert(start_offset, buf);
+        },
+    )
+}
+
+// Generates the surrounding `impl Type { <code> }` including type and lifetime
+// parameters
+fn generate_impl_text(strukt: &ast::Enum, code: &str) -> String {
+    let mut buf = String::with_capacity(code.len());
+    buf.push_str("\n\nimpl");
+    buf.push_str(" ");
+    buf.push_str(strukt.name().unwrap().text());
+    format_to!(buf, " {{\n{}\n}}", code);
+    buf
+}
+
+fn to_lower_snake_case(s: &str) -> String {
+    let mut buf = String::with_capacity(s.len());
+    let mut prev = false;
+    for c in s.chars() {
+        if c.is_ascii_uppercase() && prev {
+            buf.push('_')
+        }
+        prev = true;
+
+        buf.push(c.to_ascii_lowercase());
+    }
+    buf
+}
+
+// Uses a syntax-driven approach to find any impl blocks for the struct that
+// exist within the module/file
+//
+// Returns `None` if we've found an existing `new` fn
+//
+// FIXME: change the new fn checking to a more semantic approach when that's more
+// viable (e.g. we process proc macros, etc)
+fn find_struct_impl(
+    ctx: &AssistContext,
+    strukt: &ast::Enum,
+    name: &str,
+) -> Option<Option<ast::Impl>> {
+    let db = ctx.db();
+    let module = strukt.syntax().ancestors().find(|node| {
+        ast::Module::can_cast(node.kind()) || ast::SourceFile::can_cast(node.kind())
+    })?;
+
+    let struct_def = ctx.sema.to_def(strukt)?;
+
+    let block = module.descendants().filter_map(ast::Impl::cast).find_map(|impl_blk| {
+        let blk = ctx.sema.to_def(&impl_blk)?;
+
+        // FIXME: handle e.g. `struct S<T>; impl<U> S<U> {}`
+        // (we currently use the wrong type parameter)
+        // also we wouldn't want to use e.g. `impl S<u32>`
+        let same_ty = match blk.target_ty(db).as_adt() {
+            Some(def) => def == Adt::Enum(struct_def),
+            None => false,
+        };
+        let not_trait_impl = blk.target_trait(db).is_none();
+
+        if !(same_ty && not_trait_impl) {
+            None
+        } else {
+            Some(impl_blk)
+        }
+    });
+
+    if let Some(ref impl_blk) = block {
+        if has_fn(impl_blk, name) {
+            mark::hit!(test_gen_enum_match_impl_already_exists);
+            return None;
+        }
+    }
+
+    Some(block)
+}
+
+fn has_fn(imp: &ast::Impl, rhs_name: &str) -> bool {
+    if let Some(il) = imp.assoc_item_list() {
+        for item in il.assoc_items() {
+            if let ast::AssocItem::Fn(f) = item {
+                if let Some(name) = f.name() {
+                    if name.text().eq_ignore_ascii_case(rhs_name) {
+                        return true;
+                    }
+                }
+            }
+        }
+    }
+
+    false
+}
+
+#[cfg(test)]
+mod tests {
+    use ide_db::helpers::FamousDefs;
+    use test_utils::mark;
+
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    fn check_not_applicable(ra_fixture: &str) {
+        let fixture =
+            format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
+        check_assist_not_applicable(generate_enum_match_method, &fixture)
+    }
+
+    #[test]
+    fn test_generate_enum_match_from_variant() {
+        check_assist(
+            generate_enum_match_method,
+            r#"
+enum Variant {
+    Undefined,
+    Minor$0,
+    Major,
+}"#,
+            r#"enum Variant {
+    Undefined,
+    Minor,
+    Major,
+}
+
+impl Variant {
+    fn is_minor(&self) -> bool {
+        matches!(self, Self::Minor)
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_match_already_implemented() {
+        mark::check!(test_gen_enum_match_impl_already_exists);
+        check_not_applicable(
+            r#"
+enum Variant {
+    Undefined,
+    Minor$0,
+    Major,
+}
+
+impl Variant {
+    fn is_minor(&self) -> bool {
+        matches!(self, Self::Minor)
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_add_from_impl_no_element() {
+        mark::check!(test_gen_enum_match_on_non_unit_variant_not_implemented);
+        check_not_applicable(
+            r#"
+enum Variant {
+    Undefined,
+    Minor(u32)$0,
+    Major,
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_match_from_variant_with_one_variant() {
+        check_assist(
+            generate_enum_match_method,
+            r#"enum Variant { Undefi$0ned }"#,
+            r#"
+enum Variant { Undefined }
+
+impl Variant {
+    fn is_undefined(&self) -> bool {
+        matches!(self, Self::Undefined)
+    }
+}"#,
+        );
+    }
+
+    #[test]
+    fn test_generate_enum_match_from_variant_with_visibility_marker() {
+        check_assist(
+            generate_enum_match_method,
+            r#"
+pub(crate) enum Variant {
+    Undefined,
+    Minor$0,
+    Major,
+}"#,
+            r#"pub(crate) enum Variant {
+    Undefined,
+    Minor,
+    Major,
+}
+
+impl Variant {
+    pub(crate) fn is_minor(&self) -> bool {
+        matches!(self, Self::Minor)
+    }
+}"#,
+        );
+    }
+}
diff --git a/crates/assists/src/lib.rs b/crates/assists/src/lib.rs
index 559b9651e..a18232877 100644
--- a/crates/assists/src/lib.rs
+++ b/crates/assists/src/lib.rs
@@ -126,6 +126,7 @@ mod handlers {
     mod flip_trait_bound;
     mod generate_default_from_enum_variant;
     mod generate_derive;
+    mod generate_enum_match_method;
     mod generate_from_impl_for_enum;
     mod generate_function;
     mod generate_impl;
@@ -183,6 +184,7 @@ mod handlers {
             flip_trait_bound::flip_trait_bound,
             generate_default_from_enum_variant::generate_default_from_enum_variant,
             generate_derive::generate_derive,
+            generate_enum_match_method::generate_enum_match_method,
             generate_from_impl_for_enum::generate_from_impl_for_enum,
             generate_function::generate_function,
             generate_impl::generate_impl,
diff --git a/crates/assists/src/tests/generated.rs b/crates/assists/src/tests/generated.rs
index 9aa807f10..ae7b400e2 100644
--- a/crates/assists/src/tests/generated.rs
+++ b/crates/assists/src/tests/generated.rs
@@ -433,6 +433,33 @@ struct Point {
 }
 
 #[test]
+fn doctest_generate_enum_match_method() {
+    check_doc_test(
+        "generate_enum_match_method",
+        r#####"
+enum Version {
+ Undefined,
+ Minor$0,
+ Major,
+}
+"#####,
+        r#####"
+enum Version {
+ Undefined,
+ Minor,
+ Major,
+}
+
+impl Version {
+    fn is_minor(&self) -> bool {
+        matches!(self, Self::Minor)
+    }
+}
+"#####,
+    )
+}
+
+#[test]
 fn doctest_generate_from_impl_for_enum() {
     check_doc_test(
         "generate_from_impl_for_enum",