States derive macro

Author: SpecificProtagonist

crates/bevy_ecs/macros/src/lib.rs

@@ -3,6 +3,7 @@ extern crate proc_macro;
 mod component;
 mod fetch;
 mod set;
+mod states;
 
 use crate::{fetch::derive_world_query_impl, set::derive_set};
 use bevy_macro_utils::{derive_boxed_label, get_named_struct_fields, BevyManifest};
@@ -558,3 +559,8 @@ pub fn derive_resource(input: TokenStream) -> TokenStream {
 pub fn derive_component(input: TokenStream) -> TokenStream {
     component::derive_component(input)
 }
+
+#[proc_macro_derive(States)]
+pub fn derive_states(input: TokenStream) -> TokenStream {
+    states::derive_states(input)
+}

crates/bevy_ecs/macros/src/states.rs

@@ -0,0 +1,38 @@
+use proc_macro::{Span, TokenStream};
+use quote::quote;
+use syn::{parse_macro_input, Data::Enum, DeriveInput};
+
+pub fn derive_states(input: TokenStream) -> TokenStream {
+    let ast = parse_macro_input!(input as DeriveInput);
+    let error = || {
+        syn::Error::new(
+            Span::call_site().into(),
+            "derive(States) only supports fieldless enums",
+        )
+        .into_compile_error()
+        .into()
+    };
+    let Enum(enumeration) = ast.data else {
+        return error();
+    };
+    if enumeration.variants.iter().any(|v| !v.fields.is_empty()) {
+        return error();
+    }
+
+    let generics = ast.generics;
+    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
+
+    let struct_name = &ast.ident;
+    let idents = enumeration.variants.iter().map(|v| &v.ident);
+    let len = idents.len();
+
+    quote! {
+        impl #impl_generics bevy::ecs::schedule_v3::States for #struct_name #ty_generics #where_clause {
+            type Iter = std::array::IntoIter<Self, #len>;
+
+            fn variants() -> Self::Iter {
+                [#(Self::#idents,)*].into_iter()
+            }
+        }
+    }.into()
+}

crates/bevy_ecs/src/lib.rs

@@ -48,6 +48,7 @@ pub mod prelude {
         },
         world::{FromWorld, World},
     };
+    pub use bevy_ecs_macros::States;
 }
 
 pub use bevy_ecs_macros::all_tuples;

crates/bevy_ecs/src/schedule/state.rs

@@ -25,22 +25,14 @@ use crate::world::World;
 /// ```rust
 /// use bevy_ecs::prelude::States;
 ///
-/// #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Default)]
+/// #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Default, States)]
 /// enum GameState {
 ///  #[default]
 ///   MainMenu,
 ///   SettingsMenu,
 ///   InGame,
 /// }
 ///
-/// impl States for GameState {
-///   type Iter = std::array::IntoIter<GameState, 3>;
-///
-///   fn variants() -> Self::Iter {
-///     [GameState::MainMenu, GameState::SettingsMenu, GameState::InGame].into_iter()
-///   }
-/// }
-///
 /// ```
 pub trait States: 'static + Send + Sync + Clone + PartialEq + Eq + Hash + Debug + Default {
     type Iter: Iterator<Item = Self>;

examples/ecs/state.rs

@@ -25,21 +25,13 @@ fn main() {
         .run();
 }
 
-#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Hash)]
+#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Hash, States)]
 enum AppState {
     #[default]
     Menu,
     InGame,
 }
 
-impl States for AppState {
-    type Iter = std::array::IntoIter<AppState, 2>;
-
-    fn variants() -> Self::Iter {
-        [AppState::Menu, AppState::InGame].into_iter()
-    }
-}
-
 #[derive(Resource)]
 struct MenuData {
     button_entity: Entity,

examples/games/alien_cake_addict.rs

@@ -5,21 +5,13 @@ use std::f32::consts::PI;
 use bevy::prelude::*;
 use rand::Rng;
 
-#[derive(Clone, Eq, PartialEq, Debug, Hash, Default)]
+#[derive(Clone, Eq, PartialEq, Debug, Hash, Default, States)]
 enum GameState {
     #[default]
     Playing,
     GameOver,
 }
 
-impl States for GameState {
-    type Iter = std::array::IntoIter<GameState, 2>;
-
-    fn variants() -> Self::Iter {
-        [GameState::Playing, GameState::GameOver].into_iter()
-    }
-}
-
 #[derive(Resource)]
 struct BonusSpawnTimer(Timer);
 

examples/games/game_menu.rs

@@ -7,22 +7,14 @@ use bevy::prelude::*;
 const TEXT_COLOR: Color = Color::rgb(0.9, 0.9, 0.9);
 
 // Enum that will be used as a global state for the game
-#[derive(Clone, Copy, Default, Eq, PartialEq, Debug, Hash)]
+#[derive(Clone, Copy, Default, Eq, PartialEq, Debug, Hash, States)]
 enum GameState {
     #[default]
     Splash,
     Menu,
     Game,
 }
 
-impl States for GameState {
-    type Iter = std::array::IntoIter<GameState, 3>;
-
-    fn variants() -> Self::Iter {
-        [GameState::Splash, GameState::Menu, GameState::Game].into_iter()
-    }
-}
-
 // One of the two settings that can be set through the menu. It will be a resource in the app
 #[derive(Resource, Debug, Component, PartialEq, Eq, Clone, Copy)]
 enum DisplayQuality {
@@ -312,7 +304,7 @@ mod menu {
     }
 
     // State used for the current menu screen
-    #[derive(Clone, Copy, Default, Eq, PartialEq, Debug, Hash)]
+    #[derive(Clone, Copy, Default, Eq, PartialEq, Debug, Hash, States)]
     enum MenuState {
         Main,
         Settings,
@@ -322,21 +314,6 @@ mod menu {
         Disabled,
     }
 
-    impl States for MenuState {
-        type Iter = std::array::IntoIter<MenuState, 5>;
-
-        fn variants() -> Self::Iter {
-            [
-                MenuState::Main,
-                MenuState::Settings,
-                MenuState::SettingsDisplay,
-                MenuState::SettingsSound,
-                MenuState::Disabled,
-            ]
-            .into_iter()
-        }
-    }
-
     // Tag component used to tag entities added on the main menu screen
     #[derive(Component)]
     struct OnMainMenuScreen;