@@ -612,20 +612,41 @@ pub fn derive_substates(input: TokenStream) -> TokenStream {
612
612
states:: derive_substates ( input)
613
613
}
614
614
615
- #[ proc_macro_derive( FromWorld ) ]
615
+ #[ proc_macro_derive( FromWorld , attributes ( from_world ) ) ]
616
616
pub fn derive_from_world ( input : TokenStream ) -> TokenStream {
617
617
let bevy_ecs_path = bevy_ecs_path ( ) ;
618
618
let ast = parse_macro_input ! ( input as DeriveInput ) ;
619
- let struct_name = ast. ident ;
619
+ let name = ast. ident ;
620
620
let ( impl_generics, ty_generics, where_clauses) = ast. generics . split_for_impl ( ) ;
621
621
622
- let Data :: Struct ( DataStruct { fields, .. } ) = & ast. data else {
623
- return syn:: Error :: new (
624
- Span :: call_site ( ) ,
625
- "#[derive(FromWorld)]` only supports structs" ,
626
- )
627
- . into_compile_error ( )
628
- . into ( ) ;
622
+ let ( fields, variant_ident) = match & ast. data {
623
+ Data :: Struct ( data) => ( & data. fields , None ) ,
624
+ Data :: Enum ( data) => {
625
+ match data. variants . iter ( ) . find ( |variant| {
626
+ variant
627
+ . attrs
628
+ . iter ( )
629
+ . any ( |attr| attr. path ( ) . is_ident ( "from_world" ) )
630
+ } ) {
631
+ Some ( variant) => ( & variant. fields , Some ( & variant. ident ) ) ,
632
+ None => {
633
+ return syn:: Error :: new (
634
+ Span :: call_site ( ) ,
635
+ "No #[from_world] attribute was found on any of this enum's variants." ,
636
+ )
637
+ . into_compile_error ( )
638
+ . into ( ) ;
639
+ }
640
+ }
641
+ }
642
+ Data :: Union ( _) => {
643
+ return syn:: Error :: new (
644
+ Span :: call_site ( ) ,
645
+ "#[derive(FromWorld)]` does not support unions" ,
646
+ )
647
+ . into_compile_error ( )
648
+ . into ( ) ;
649
+ }
629
650
} ;
630
651
631
652
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 {
645
666
syn:: Fields :: Unit => Punctuated :: new ( ) ,
646
667
} ;
647
668
669
+ let field_initializers: TokenStream2 = if !field_initializers. is_empty ( ) {
670
+ quote ! ( { #field_initializers } )
671
+ } else {
672
+ quote ! ( #field_initializers)
673
+ } ;
674
+
675
+ let field_initializers = match variant_ident {
676
+ Some ( variant_ident) => quote ! ( Self :: #variant_ident #field_initializers) ,
677
+ None => quote ! ( Self #field_initializers) ,
678
+ } ;
679
+
648
680
TokenStream :: from ( quote ! {
649
- impl #impl_generics #bevy_ecs_path:: world:: FromWorld for #struct_name #ty_generics #where_clauses {
650
- fn from_world( world: & mut #bevy_ecs_path:: world:: World ) -> Self {
651
- #[ allow( clippy:: init_numbered_fields) ]
652
- Self { #field_initializers }
681
+ impl #impl_generics #bevy_ecs_path:: world:: FromWorld for #name #ty_generics #where_clauses {
682
+ fn from_world( world: & mut #bevy_ecs_path:: world:: World ) -> Self {
683
+ #[ allow( clippy:: init_numbered_fields) ]
684
+ #field_initializers
685
+ }
653
686
}
654
- }
655
687
} )
656
688
}
0 commit comments