From 359b771c525f4647df7b0e37e07d8a7d270007c2 Mon Sep 17 00:00:00 2001 From: A4-Tacks Date: Fri, 14 Nov 2025 16:27:29 +0800 Subject: [PATCH] Support multiple variant for generate_from_impl_for_enum Example --- ```rust enum A { $0Foo(u32), Bar$0(i32) } ``` **Before this PR** ```rust enum A { Foo(u32), Bar(i32) } impl From for A { fn from(v: u32) -> Self { Self::Foo(v) } } ``` **After this PR** ```rust enum A { Foo(u32), Bar(i32) } impl From for A { fn from(v: u32) -> Self { Self::Foo(v) } } impl From for A { fn from(v: i32) -> Self { Self::Bar(v) } } ``` --- .../handlers/generate_from_impl_for_enum.rs | 124 +++++++++++++----- crates/ide-assists/src/utils.rs | 9 ++ 2 files changed, 103 insertions(+), 30 deletions(-) diff --git a/crates/ide-assists/src/handlers/generate_from_impl_for_enum.rs b/crates/ide-assists/src/handlers/generate_from_impl_for_enum.rs index d88b0f34b796..6a868239cb20 100644 --- a/crates/ide-assists/src/handlers/generate_from_impl_for_enum.rs +++ b/crates/ide-assists/src/handlers/generate_from_impl_for_enum.rs @@ -2,7 +2,10 @@ use hir::next_solver::{DbInterner, TypingMode}; use ide_db::{RootDatabase, famous_defs::FamousDefs}; use syntax::ast::{self, AstNode, HasName}; -use crate::{AssistContext, AssistId, Assists, utils::generate_trait_impl_text_intransitive}; +use crate::{ + AssistContext, AssistId, Assists, + utils::{generate_trait_impl_text_intransitive, is_selected}, +}; // Assist: generate_from_impl_for_enum // @@ -26,8 +29,68 @@ pub(crate) fn generate_from_impl_for_enum( ctx: &AssistContext<'_>, ) -> Option<()> { let variant = ctx.find_node_at_offset::()?; - let variant_name = variant.name()?; - let enum_ = ast::Adt::Enum(variant.parent_enum()); + let adt = ast::Adt::Enum(variant.parent_enum()); + let variants = selected_variants(ctx, &variant)?; + + let target = variant.syntax().text_range(); + acc.add( + AssistId::generate("generate_from_impl_for_enum"), + "Generate `From` impl for this enum variant(s)", + target, + |edit| { + let start_offset = variant.parent_enum().syntax().text_range().end(); + let from_impl = variants + .into_iter() + .map(|variant_info| { + let from_trait = format!("From<{}>", variant_info.ty); + let impl_code = generate_impl_code(variant_info); + generate_trait_impl_text_intransitive(&adt, &from_trait, &impl_code) + }) + .collect::(); + edit.insert(start_offset, from_impl); + }, + ) +} + +fn generate_impl_code(VariantInfo { name, field_name, ty }: VariantInfo) -> String { + if let Some(field) = field_name { + format!( + r#" fn from({field}: {ty}) -> Self {{ + Self::{name} {{ {field} }} + }}"# + ) + } else { + format!( + r#" fn from(v: {ty}) -> Self {{ + Self::{name}(v) + }}"# + ) + } +} + +struct VariantInfo { + name: ast::Name, + field_name: Option, + ty: ast::Type, +} + +fn selected_variants(ctx: &AssistContext<'_>, variant: &ast::Variant) -> Option> { + variant + .parent_enum() + .variant_list()? + .variants() + .filter(|it| is_selected(it, ctx.selection_trimmed(), true)) + .map(|variant| { + let (name, ty) = extract_variant_info(&ctx.sema, &variant)?; + Some(VariantInfo { name: variant.name()?, field_name: name, ty }) + }) + .collect() +} + +fn extract_variant_info( + sema: &'_ hir::Semantics<'_, RootDatabase>, + variant: &ast::Variant, +) -> Option<(Option, ast::Type)> { let (field_name, field_type) = match variant.kind() { ast::StructKind::Tuple(field_list) => { if field_list.fields().count() != 1 { @@ -45,36 +108,11 @@ pub(crate) fn generate_from_impl_for_enum( ast::StructKind::Unit => return None, }; - if existing_from_impl(&ctx.sema, &variant).is_some() { + if existing_from_impl(sema, variant).is_some() { cov_mark::hit!(test_add_from_impl_already_exists); return None; } - - let target = variant.syntax().text_range(); - acc.add( - AssistId::generate("generate_from_impl_for_enum"), - "Generate `From` impl for this enum variant", - target, - |edit| { - let start_offset = variant.parent_enum().syntax().text_range().end(); - let from_trait = format!("From<{field_type}>"); - let impl_code = if let Some(name) = field_name { - format!( - r#" fn from({name}: {field_type}) -> Self {{ - Self::{variant_name} {{ {name} }} - }}"# - ) - } else { - format!( - r#" fn from(v: {field_type}) -> Self {{ - Self::{variant_name}(v) - }}"# - ) - }; - let from_impl = generate_trait_impl_text_intransitive(&enum_, &from_trait, &impl_code); - edit.insert(start_offset, from_impl); - }, - ) + Some((field_name, field_type)) } fn existing_from_impl( @@ -123,6 +161,32 @@ impl From for A { ); } + #[test] + fn test_generate_from_impl_for_multiple_enum_variants() { + check_assist( + generate_from_impl_for_enum, + r#" +//- minicore: from +enum A { $0Foo(u32), Bar$0(i32) } +"#, + r#" +enum A { Foo(u32), Bar(i32) } + +impl From for A { + fn from(v: u32) -> Self { + Self::Foo(v) + } +} + +impl From for A { + fn from(v: i32) -> Self { + Self::Bar(v) + } +} +"#, + ); + } + // FIXME(next-solver): it would be nice to not be *required* to resolve the // path in order to properly generate assists #[test] diff --git a/crates/ide-assists/src/utils.rs b/crates/ide-assists/src/utils.rs index e43516f6b963..86acab4615c2 100644 --- a/crates/ide-assists/src/utils.rs +++ b/crates/ide-assists/src/utils.rs @@ -1155,6 +1155,15 @@ pub(crate) fn cover_let_chain(mut expr: ast::Expr, range: TextRange) -> Option bool { + selection.intersect(it.syntax().text_range()).is_some_and(|it| !it.is_empty()) + || allow_empty && it.syntax().text_range().contains_range(selection) +} + pub fn is_body_const(sema: &Semantics<'_, RootDatabase>, expr: &ast::Expr) -> bool { let mut is_const = true; preorder_expr(expr, &mut |ev| {