5 files changed, 317 insertions, 64 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..270b438b7
--- /dev/null
+++ b/crates/assists/src/handlers/generate_enum_match_method.rs
@@ -0,0 +1,213 @@
+use stdx::{format_to, to_lower_snake_case};
+use syntax::ast::{self, AstNode, NameOwner};
+use syntax::{ast::VisibilityOwner, T};
+use test_utils::mark;
+
+use crate::{utils::find_struct_impl, 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(&variant_name.to_string());
+
+    // Return early if we've found an existing new fn
+    let impl_def = find_struct_impl(
+        &ctx,
+        &ast::AdtDef::Enum(parent_enum.clone()),
+        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(strukt.name().unwrap().text());
+    format_to!(buf, " {{\n{}\n}}", code);
+    buf
+}
+
+#[cfg(test)]
+mod tests {
+    use test_utils::mark;
+
+    use crate::tests::{check_assist, check_assist_not_applicable};
+
+    use super::*;
+
+    fn check_not_applicable(ra_fixture: &str) {
+        check_assist_not_applicable(generate_enum_match_method, ra_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() {
+        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/handlers/generate_new.rs b/crates/assists/src/handlers/generate_new.rs
index b7390855a..84832273f 100644
--- a/crates/assists/src/handlers/generate_new.rs
+++ b/crates/assists/src/handlers/generate_new.rs
@@ -1,4 +1,3 @@
-use hir::Adt;
 use itertools::Itertools;
 use stdx::format_to;
 use syntax::{
@@ -6,7 +5,7 @@ use syntax::{
     SmolStr, T,
 };
 
-use crate::{AssistContext, AssistId, AssistKind, Assists};
+use crate::{utils::find_struct_impl, AssistContext, AssistId, AssistKind, Assists};
 
 // Assist: generate_new
 //
@@ -38,7 +37,7 @@ pub(crate) fn generate_new(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
     };
 
     // Return early if we've found an existing new fn
-    let impl_def = find_struct_impl(&ctx, &strukt)?;
+    let impl_def = find_struct_impl(&ctx, &ast::AdtDef::Struct(strukt.clone()), "new")?;
 
     let target = strukt.syntax().text_range();
     acc.add(AssistId("generate_new", AssistKind::Generate), "Generate `new`", target, |builder| {
@@ -111,65 +110,6 @@ fn generate_impl_text(strukt: &ast::Struct, code: &str) -> String {
     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::Struct) -> 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::Struct(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_new_fn(impl_blk) {
-            return None;
-        }
-    }
-
-    Some(block)
-}
-
-fn has_new_fn(imp: &ast::Impl) -> 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("new") {
-                        return true;
-                    }
-                }
-            }
-        }
-    }
-
-    false
-}
-
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};
diff --git a/crates/assists/src/lib.rs b/crates/assists/src/lib.rs
index 062a902ab..83fbf6986 100644
--- a/crates/assists/src/lib.rs
+++ b/crates/assists/src/lib.rs
@@ -127,6 +127,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;
@@ -185,6 +186,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 e84f208a3..960815bd9 100644
--- a/crates/assists/src/tests/generated.rs
+++ b/crates/assists/src/tests/generated.rs
@@ -460,6 +460,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",
diff --git a/crates/assists/src/utils.rs b/crates/assists/src/utils.rs
index 4e762e18b..3842558d8 100644
--- a/crates/assists/src/utils.rs
+++ b/crates/assists/src/utils.rs
@@ -2,7 +2,7 @@
 
 use std::ops;
 
-use hir::HasSource;
+use hir::{Adt, HasSource};
 use ide_db::{helpers::SnippetCap, RootDatabase};
 use itertools::Itertools;
 use syntax::{
@@ -15,7 +15,10 @@ use syntax::{
     SyntaxNode, TextSize, T,
 };
 
-use crate::ast_transform::{self, AstTransform, QualifyPaths, SubstituteTypeParams};
+use crate::{
+    assist_context::AssistContext,
+    ast_transform::{self, AstTransform, QualifyPaths, SubstituteTypeParams},
+};
 
 pub(crate) fn unwrap_trivial_block(block: ast::BlockExpr) -> ast::Expr {
     extract_trivial_expression(&block)
@@ -267,3 +270,71 @@ pub(crate) fn does_pat_match_variant(pat: &ast::Pat, var: &ast::Pat) -> bool {
 
     pat_head == var_head
 }
+
+// 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)
+pub(crate) fn find_struct_impl(
+    ctx: &AssistContext,
+    strukt: &ast::AdtDef,
+    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 = match strukt {
+        ast::AdtDef::Enum(e) => Adt::Enum(ctx.sema.to_def(e)?),
+        ast::AdtDef::Struct(s) => Adt::Struct(ctx.sema.to_def(s)?),
+        ast::AdtDef::Union(u) => Adt::Union(ctx.sema.to_def(u)?),
+    };
+
+    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 == 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) {
+            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
+}