From da57dfb62f8d36df6be516e5800df6387c5e986e Mon Sep 17 00:00:00 2001
From: Tim Overbeek <158390905+Bleachfuel@users.noreply.github.com>
Date: Thu, 23 Jan 2025 05:06:00 +0100
Subject: [PATCH] DeriveWorld for enums (#17496)

# Objective

Fixes #17457

## Solution

#[derive(FromWorld)] now works with enums by specifying which variant
should be used.

## Showcase

```rust
#[Derive(FromWorld)]
enum Game {
    #[from_world]
    Playing,
    Stopped
}
```

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
Co-authored-by: Benjamin Brienen <benjamin.brienen@outlook.com>
---
 crates/bevy_ecs/macros/src/lib.rs | 60 +++++++++++++++++++++++--------
 crates/bevy_ecs/src/world/mod.rs  | 15 ++++++--
 2 files changed, 59 insertions(+), 16 deletions(-)

diff --git a/crates/bevy_ecs/macros/src/lib.rs b/crates/bevy_ecs/macros/src/lib.rs
index 45a948cc8c..f93dfabab9 100644
--- a/crates/bevy_ecs/macros/src/lib.rs
+++ b/crates/bevy_ecs/macros/src/lib.rs
@@ -612,20 +612,41 @@ pub fn derive_substates(input: TokenStream) -> TokenStream {
     states::derive_substates(input)
 }
 
-#[proc_macro_derive(FromWorld)]
+#[proc_macro_derive(FromWorld, attributes(from_world))]
 pub fn derive_from_world(input: TokenStream) -> TokenStream {
     let bevy_ecs_path = bevy_ecs_path();
     let ast = parse_macro_input!(input as DeriveInput);
-    let struct_name = ast.ident;
+    let name = ast.ident;
     let (impl_generics, ty_generics, where_clauses) = ast.generics.split_for_impl();
 
-    let Data::Struct(DataStruct { fields, .. }) = &ast.data else {
-        return syn::Error::new(
-            Span::call_site(),
-            "#[derive(FromWorld)]` only supports structs",
-        )
-        .into_compile_error()
-        .into();
+    let (fields, variant_ident) = match &ast.data {
+        Data::Struct(data) => (&data.fields, None),
+        Data::Enum(data) => {
+            match data.variants.iter().find(|variant| {
+                variant
+                    .attrs
+                    .iter()
+                    .any(|attr| attr.path().is_ident("from_world"))
+            }) {
+                Some(variant) => (&variant.fields, Some(&variant.ident)),
+                None => {
+                    return syn::Error::new(
+                        Span::call_site(),
+                        "No #[from_world] attribute was found on any of this enum's variants.",
+                    )
+                    .into_compile_error()
+                    .into();
+                }
+            }
+        }
+        Data::Union(_) => {
+            return syn::Error::new(
+                Span::call_site(),
+                "#[derive(FromWorld)]` does not support unions",
+            )
+            .into_compile_error()
+            .into();
+        }
     };
 
     let field_init_expr = quote!(#bevy_ecs_path::world::FromWorld::from_world(world));
@@ -645,12 +666,23 @@ pub fn derive_from_world(input: TokenStream) -> TokenStream {
         syn::Fields::Unit => Punctuated::new(),
     };
 
+    let field_initializers: TokenStream2 = if !field_initializers.is_empty() {
+        quote!({ #field_initializers })
+    } else {
+        quote!(#field_initializers)
+    };
+
+    let field_initializers = match variant_ident {
+        Some(variant_ident) => quote!( Self::#variant_ident #field_initializers),
+        None => quote!( Self #field_initializers),
+    };
+
     TokenStream::from(quote! {
-        impl #impl_generics #bevy_ecs_path::world::FromWorld for #struct_name #ty_generics #where_clauses {
-            fn from_world(world: &mut #bevy_ecs_path::world::World) -> Self {
-                #[allow(clippy::init_numbered_fields)]
-                Self { #field_initializers }
+            impl #impl_generics #bevy_ecs_path::world::FromWorld for #name #ty_generics #where_clauses {
+                fn from_world(world: &mut #bevy_ecs_path::world::World) -> Self {
+                    #[allow(clippy::init_numbered_fields)]
+                    #field_initializers
+                }
             }
-        }
     })
 }
diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs
index 7d2eb08a09..22520ac71a 100644
--- a/crates/bevy_ecs/src/world/mod.rs
+++ b/crates/bevy_ecs/src/world/mod.rs
@@ -3716,9 +3716,13 @@ unsafe impl Sync for World {}
 ///
 /// This can be helpful for complex initialization or context-aware defaults.
 ///
-/// [`FromWorld`] is automatically implemented for any type implementing [`Default`],
-/// and may also be derived for any struct whose fields all implement `FromWorld`:
+/// [`FromWorld`] is automatically implemented for any type implementing [`Default`]
+/// and may also be derived for:
+/// - any struct whose fields all implement `FromWorld`
+/// - any enum where one variant has the attribute `#[from_world]`
+///
 /// ```rs
+///
 /// #[derive(Default)]
 /// struct A;
 ///
@@ -3735,6 +3739,13 @@ unsafe impl Sync for World {}
 ///
 /// #[derive(FromWorld)]
 /// struct D(A, B, C);
+///
+/// #[derive(FromWorld)]
+/// enum E {
+///     #[from_world]
+///     F,
+///     G
+/// }
 /// ```
 pub trait FromWorld {
     /// Creates `Self` using data from the given [`World`].