2 files changed, 223 insertions, 62 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..ee89d4208
--- /dev/null
+++ b/crates/assists/src/handlers/generate_enum_match_method.rs
@@ -0,0 +1,221 @@
+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 {
+//     /// Returns `true` if the version is [`Minor`].
+//     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 enum_lowercase_name = to_lower_snake_case(&parent_enum.name()?.to_string());
+    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,
+                "    /// Returns `true` if the {} is [`{}`].
+    {}fn is_{}(&self) -> bool {{
+        matches!(self, Self::{})
+    }}",
+                enum_lowercase_name,
+                variant_name,
+                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 {
+    /// Returns `true` if the variant is [`Minor`].
+    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 {
+    /// Returns `true` if the variant is [`Undefined`].
+    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 {
+    /// Returns `true` if the variant is [`Minor`].
+    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};