dd

crates/bevy_ecs/macros/src/states.rs

@@ -1,6 +1,7 @@
 use proc_macro::{Span, TokenStream};
+use proc_macro2::Ident;
 use quote::{format_ident, quote};
-use syn::{parse_macro_input, Data::Enum, DeriveInput};
+use syn::{parse_macro_input, spanned::Spanned, Data::Enum, DeriveInput, Field, Fields};
 
 use crate::bevy_ecs_path;
 
@@ -9,36 +10,90 @@ pub fn derive_states(input: TokenStream) -> TokenStream {
     let error = || {
         syn::Error::new(
             Span::call_site().into(),
-            "derive(States) only supports fieldless enums",
+            "derive(States) only supports enums (additionaly, all of the fields need to implement States too)",
         )
         .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 mut error: Option<syn::Error> = None;
+    let non_unnamed = enumeration
+        .variants
+        .iter()
+        .filter(|v| matches!(v.fields, Fields::Named(_)));
+    for variant in non_unnamed {
+        let err = syn::Error::new(
+            variant.span(),
+            format!(
+                "Expected either unit (e.g. None) or unnamed field (e.g. Some(T)), found {}",
+                match variant.fields {
+                    Fields::Named(_) => "named structs (e.g. Foo { bar: bool }",
+                    _ => unreachable!(),
+                }
+            ),
+        );
+        match &mut error {
+            Some(error) => error.combine(err),
+            None => error = Some(err),
+        }
     }
 
+    if let Some(error) = error {
+        return error.into_compile_error().into();
+    }
+
+    {}
     let generics = ast.generics;
     let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
 
     let mut trait_path = bevy_ecs_path();
     trait_path.segments.push(format_ident!("schedule").into());
     trait_path.segments.push(format_ident!("States").into());
     let struct_name = &ast.ident;
-    let idents = enumeration.variants.iter().map(|v| &v.ident);
-    let len = idents.len();
+    let fieldless_idents = enumeration
+        .variants
+        .iter()
+        .filter(|v| v.fields.is_empty())
+        .map(|v| &v.ident);
+    let fieldful_variants = enumeration.variants.iter().filter(|v| !v.fields.is_empty());
+    let fieldful_idents: Vec<&Ident> = fieldful_variants.clone().map(|v| &v.ident).collect();
+    let fieldful_values: Vec<&Field> = fieldful_variants
+        .flat_map(|v| match &v.fields {
+            syn::Fields::Unnamed(field) => &field.unnamed,
+            _ => unreachable!(),
+        })
+        .collect();
+
+    let len = enumeration.variants.len();
+    let (variants_impl, iter_type) = if !fieldful_idents.is_empty() {
+        (
+            quote! {
 
-    quote! {
+                [vec![#(Self::#fieldless_idents,)*], #(<#fieldful_values as #trait_path>::variants().map(|variant| {
+                            Self::#fieldful_idents(variant)
+                }).collect::<Vec<Self>>(),)*].into_iter()
+                .flatten()
+                .collect::<Vec<Self>>()
+                .into_iter()
+            },
+            quote! {std::vec::IntoIter<Self>},
+        )
+    } else {
+        (
+            quote! {[#(Self::#fieldless_idents,)*].into_iter()  },
+            quote! {std::array::IntoIter<Self, #len>},
+        )
+    };
+    let token = quote! {
         impl #impl_generics #trait_path for #struct_name #ty_generics #where_clause {
-            type Iter = std::array::IntoIter<Self, #len>;
+            type Iter = #iter_type;
 
             fn variants() -> Self::Iter {
-                [#(Self::#idents,)*].into_iter()
+                #variants_impl
             }
         }
-    }
-    .into()
+    };
+    token.into()
 }

crates/bevy_ecs/src/lib.rs

@@ -38,9 +38,12 @@ pub mod prelude {
         query::{Added, AnyOf, Changed, Or, QueryState, With, Without},
         removal_detection::RemovedComponents,
         schedule::{
-            apply_state_transition, apply_system_buffers, common_conditions::*, Condition,
-            IntoSystemConfigs, IntoSystemSet, IntoSystemSetConfig, IntoSystemSetConfigs, NextState,
-            OnEnter, OnExit, OnTransition, OnUpdate, Schedule, Schedules, State, States, SystemSet,
+            apply_state_transition, apply_system_buffers,
+            common_conditions::*,
+            substate::{self, SubstateLabel, SubstateLabelInFn},
+            Condition, IntoSystemConfigs, IntoSystemSet, IntoSystemSetConfig, IntoSystemSetConfigs,
+            NextState, OnEnter, OnExit, OnTransition, OnUpdate, Schedule, Schedules, State, States,
+            SystemSet,
         },
         system::{
             adapter as system_adapter,
@@ -61,6 +64,7 @@ type TypeIdMap<V> = rustc_hash::FxHashMap<TypeId, V>;
 mod tests {
     use crate as bevy_ecs;
     use crate::prelude::Or;
+    use crate::schedule::States;
     use crate::{
         bundle::Bundle,
         change_detection::Ref,
@@ -1707,4 +1711,31 @@ mod tests {
             "new entity was spawned and received C component"
         );
     }
+
+    #[test]
+    fn fieldful_and_fieldless_states() {
+        use bevy_ecs_macros::States;
+        #[derive(Hash, Default, Eq, PartialEq, Clone, Debug, States)]
+        pub enum Foo {
+            #[default]
+            Fieldless,
+
+            Fieldful(Bar),
+        }
+        #[derive(Hash, Eq, Default, PartialEq, Clone, Debug, States)]
+        pub enum Bar {
+            #[default]
+            Alice,
+            Bob,
+        }
+
+        assert_eq!(
+            Foo::variants().collect::<Vec<Foo>>(),
+            vec![
+                Foo::Fieldless,
+                Foo::Fieldful(Bar::Alice),
+                Foo::Fieldful(Bar::Bob)
+            ]
+        );
+    }
 }

crates/bevy_ecs/src/schedule/mod.rs

@@ -6,7 +6,7 @@ mod graph_utils;
 mod schedule;
 mod set;
 mod state;
-
+pub mod substate;
 pub use self::condition::*;
 pub use self::config::*;
 pub use self::executor::*;

crates/bevy_ecs/src/schedule/state.rs

@@ -1,5 +1,6 @@
 use std::fmt::Debug;
 use std::hash::Hash;
+
 use std::mem;
 
 use crate as bevy_ecs;
@@ -54,6 +55,15 @@ pub struct OnEnter<S: States>(pub S);
 #[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct OnExit<S: States>(pub S);
 
+
+/*
+#[derive(Clone, PartialEq, Eq, Hash, Debug, Default, ScheduleLabel)]
+///
+/// Schedule label that matches all substates in the enum tree that got changed too.
+///
+pub struct SubstateFromAny<L: ScheduleLabel>(pub L);
+*/
+
 /// The label of a [`Schedule`](super::Schedule) that **only** runs whenever [`State<S>`]
 /// exits the `from` state, AND enters the `to` state.
 ///
@@ -132,3 +142,4 @@ pub fn apply_state_transition<S: States>(world: &mut World) {
         world.try_run_schedule(OnEnter(entered)).ok();
     }
 }
+

crates/bevy_ecs/src/schedule/substate.rs

@@ -0,0 +1,67 @@
+use crate::schedule::ScheduleLabel;
+
+use crate as bevy_ecs;
+use bevy_utils::{all_tuples};
+use std::marker::PhantomData;
+
+#[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
+///
+///
+/// This should be able to be used on some unnamed fields in enums. However, please beware that fields that have the same substate might pass the type checking, as they really match type-wise. (
+/// So
+/// ```rs
+/// #[derive(States, Clone, Copy, Default, Eq, PartialEq, Hash, Debug)]
+/// pub enum MyState {
+///     Foo(AliceOrBob)
+///     Bar(AliceOrBob)
+/// }
+///
+/// pub AliceOrBob {
+///     Alice,
+///     Bob
+/// }
+/// // these should compile
+/// let mut foo: Fn(AliceOrBob) -> MyState = MyState::Foo;
+/// foo = MyState::Bar;
+///
+/// ```
+pub struct SubstateLabelInFn<Ret, Args, L: SubstateLabel>(PhantomData<(L, Ret, Args)>);
+
+/// A trait which most variadic functions should technically implement?
+/// The issue with Fn trait types is that it doesn't support variadic functions.  This one should
+pub trait VariadicFn<Ret, Args> {}
+
+macro_rules! impl_substate_tuple {
+    ($($name: tt),*)  => {
+        impl<Function, Ret, $( $name, )* >  VariadicFn<Ret, ($($name,)*)> for Function where Function: Fn($($name,)*) -> Ret {}
+    }
+}
+
+pub trait Arguments {}
+macro_rules! impl_args_tuple {
+    ($($name: tt),*) => {
+        impl< $( $name, )* > Arguments for  ($( $name, )* ) {
+
+        }
+    }
+}
+
+all_tuples!(impl_args_tuple, 0, 16, A);
+all_tuples!(impl_substate_tuple, 0, 16, S);
+pub trait SubstateLabel
+where
+    Self: Sized,
+{
+    fn with<F: VariadicFn<Ret, Args>, Ret, Args>(_f: &F) -> SubstateLabelInFn<Ret, Args, Self> {
+        SubstateLabelInFn(PhantomData)
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct OnEnter;
+impl SubstateLabel for OnEnter {}
+
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct OnExit;
+
+impl SubstateLabel for OnExit {}

tests/how_to_test_systems.rs

@@ -1,3 +1,5 @@
+use std::sync::{Arc, Mutex};
+
 use bevy::{ecs::event::Events, prelude::*};
 
 #[derive(Component, Default)]
@@ -168,3 +170,54 @@ fn update_score_on_event() {
     // Check resulting changes
     assert_eq!(app.world.resource::<Score>().0, 3);
 }
+
+#[test]
+fn substate() {
+    #[derive(Clone, Default, Debug, Hash, Eq, PartialEq /* States */)]
+    enum ParentState {
+        #[default]
+        Unit,
+        Child(ChildState),
+    }
+    // Manual implementation first as https://github.com/bevyengine/bevy/pull/8188 is blocked for now.
+    impl States for ParentState {
+        type Iter = std::array::IntoIter<Self, 3>;
+
+        fn variants() -> Self::Iter {
+            [
+                Self::Unit,
+                Self::Child(ChildState::Bar),
+                Self::Child(ChildState::Foo),
+            ]
+            .into_iter()
+        }
+    }
+
+    #[derive(Clone, Default, Debug, Hash, Eq, PartialEq, States)]
+    enum ChildState {
+        Foo,
+        #[default]
+        Bar,
+    }
+
+    // Setup app
+    let mut app = App::new();
+    app.add_state::<ParentState>();
+    fn enter_unit(mut next: ResMut<NextState<ParentState>>) {
+        next.set(ParentState::Child(ChildState::Bar));
+    }
+    app.add_systems(OnEnter(ParentState::Unit), enter_unit);
+    let has_runned = Arc::new(Mutex::new(false));
+    let has_runned2 = has_runned.clone();
+    let enter_bar = move |state: Res<State<ParentState>>| {
+        assert_eq!(state.0, ParentState::Child(ChildState::Bar));
+        *has_runned2.lock().unwrap() = true;
+    };
+    app.add_systems(
+        substate::OnEnter::with(&ParentState::Child),
+        enter_bar,
+    );
+    app.update();
+    app.update();
+    assert_eq!(*has_runned.lock().unwrap(), true);
+}