Adding Reflect data types for States and FreelyMutableState. (#14643)

# Objective

- While developing a debug tool I saw the gap where it was not possible
to get all existing states from a World using reflection.
- This PR allows to iterate over all `States` types that exist in a
world, and modify them in case they implement `FreelyMutableState`.
- Two new methods are available on `App` and `SubApp` as helper to
register the data types:
  -  `register_state_reflect` and `register_mutable_state_reflect`

## Solution

- Two new data types are added:
- `ReflectState`: Allows to extract the current value of a state from
the World.
- `ReflectFreelyMutableState`: Allows to set the next state in a world,
similar to call `NextState::set`.
- There is no distinction between `States`, `SubStates` and
`ComputedStates`:
- `States` can register both `ReflectState` and
`ReflectFreelyMutableState`.
- `SubStates` can register both `ReflectState` and
`ReflectFreelyMutableState`.
  -  `ComputedStates` can register only `ReflectState` .

## Testing

- Added tests inside the `bevy_state` crate.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
Co-authored-by: Gino Valente <49806985+MrGVSV@users.noreply.github.com>
Co-authored-by: Jan Hohenheim <jan@hohenheim.ch>

crates/bevy_ecs/src/reflect/mod.rs

@@ -85,7 +85,7 @@ impl DerefMut for AppFunctionRegistry {
 /// this method will panic.
 ///
 /// If none of the strategies succeed, this method will panic.
-fn from_reflect_with_fallback<T: Reflect>(
+pub fn from_reflect_with_fallback<T: Reflect>(
     reflected: &dyn Reflect,
     world: &mut World,
     registry: &TypeRegistry,

crates/bevy_ecs/src/reflect/resource.rs

@@ -176,7 +176,7 @@ impl ReflectResource {
     }
 }
 
-impl<R: Resource + FromReflect> FromType<R> for ReflectResource {
+impl<R: Resource + Reflect> FromType<R> for ReflectResource {
     fn from_type() -> Self {
         ReflectResource(ReflectResourceFns {
             insert: |world, reflected_resource, registry| {

crates/bevy_state/src/app.rs

@@ -8,6 +8,9 @@ use crate::state::{
 };
 use crate::state_scoped::clear_state_scoped_entities;
 
+#[cfg(feature = "bevy_reflect")]
+use bevy_reflect::{FromReflect, GetTypeRegistration, Typed};
+
 /// State installation methods for [`App`] and [`SubApp`].
 pub trait AppExtStates {
     /// Initializes a [`State`] with standard starting values.
@@ -55,6 +58,25 @@ pub trait AppExtStates {
     ///
     /// For more information refer to [`StateScoped`](crate::state_scoped::StateScoped).
     fn enable_state_scoped_entities<S: States>(&mut self) -> &mut Self;
+
+    #[cfg(feature = "bevy_reflect")]
+    /// Registers the state type `T` using [`App::register_type`],
+    /// and adds [`ReflectState`](crate::reflect::ReflectState) type data to `T` in the type registry.
+    ///
+    /// This enables reflection code to access the state. For detailed information, see the docs on [`crate::reflect::ReflectState`] .
+    fn register_type_state<S>(&mut self) -> &mut Self
+    where
+        S: States + FromReflect + GetTypeRegistration + Typed;
+
+    #[cfg(feature = "bevy_reflect")]
+    /// Registers the state type `T` using [`App::register_type`],
+    /// and adds [`crate::reflect::ReflectState`] and [`crate::reflect::ReflectFreelyMutableState`] type data to `T` in the type registry.
+    ///
+    /// This enables reflection code to access and modify the state.
+    /// For detailed information, see the docs on [`crate::reflect::ReflectState`] and [`crate::reflect::ReflectFreelyMutableState`].
+    fn register_type_mutable_state<S>(&mut self) -> &mut Self
+    where
+        S: FreelyMutableState + FromReflect + GetTypeRegistration + Typed;
 }
 
 /// Separate function to only warn once for all state installation methods.
@@ -183,6 +205,30 @@ impl AppExtStates for SubApp {
             clear_state_scoped_entities::<S>.in_set(StateTransitionSteps::ExitSchedules),
         )
     }
+
+    #[cfg(feature = "bevy_reflect")]
+    fn register_type_state<S>(&mut self) -> &mut Self
+    where
+        S: States + FromReflect + GetTypeRegistration + Typed,
+    {
+        self.register_type::<S>();
+        self.register_type::<State<S>>();
+        self.register_type_data::<S, crate::reflect::ReflectState>();
+        self
+    }
+
+    #[cfg(feature = "bevy_reflect")]
+    fn register_type_mutable_state<S>(&mut self) -> &mut Self
+    where
+        S: FreelyMutableState + FromReflect + GetTypeRegistration + Typed,
+    {
+        self.register_type::<S>();
+        self.register_type::<State<S>>();
+        self.register_type::<NextState<S>>();
+        self.register_type_data::<S, crate::reflect::ReflectState>();
+        self.register_type_data::<S, crate::reflect::ReflectFreelyMutableState>();
+        self
+    }
 }
 
 impl AppExtStates for App {
@@ -210,6 +256,24 @@ impl AppExtStates for App {
         self.main_mut().enable_state_scoped_entities::<S>();
         self
     }
+
+    #[cfg(feature = "bevy_reflect")]
+    fn register_type_state<S>(&mut self) -> &mut Self
+    where
+        S: States + FromReflect + GetTypeRegistration + Typed,
+    {
+        self.main_mut().register_type_state::<S>();
+        self
+    }
+
+    #[cfg(feature = "bevy_reflect")]
+    fn register_type_mutable_state<S>(&mut self) -> &mut Self
+    where
+        S: FreelyMutableState + FromReflect + GetTypeRegistration + Typed,
+    {
+        self.main_mut().register_type_mutable_state::<S>();
+        self
+    }
 }
 
 /// Registers the [`StateTransition`] schedule in the [`MainScheduleOrder`] to enable state processing.

crates/bevy_state/src/lib.rs

@@ -39,6 +39,10 @@ pub mod state;
 /// [`clear_state_scoped_entities`](crate::state_scoped::clear_state_scoped_entities) for managing lifetime of entities.
 pub mod state_scoped;
 
+#[cfg(feature = "bevy_reflect")]
+/// Provides definitions for the basic traits required by the state system
+mod reflect;
+
 /// Most commonly used re-exported types.
 pub mod prelude {
     #[cfg(feature = "bevy_app")]
@@ -46,6 +50,9 @@ pub mod prelude {
     pub use crate::app::AppExtStates;
     #[doc(hidden)]
     pub use crate::condition::*;
+    #[cfg(feature = "bevy_app")]
+    #[doc(hidden)]
+    pub use crate::reflect::{ReflectFreelyMutableState, ReflectState};
     #[doc(hidden)]
     pub use crate::state::{
         last_transition, ComputedStates, EnterSchedules, ExitSchedules, NextState, OnEnter, OnExit,

crates/bevy_state/src/reflect.rs

@@ -0,0 +1,161 @@
+use crate::state::{FreelyMutableState, NextState, State, States};
+
+use bevy_ecs::reflect::from_reflect_with_fallback;
+use bevy_ecs::world::World;
+use bevy_reflect::{FromType, Reflect, TypeRegistry};
+
+/// A struct used to operate on the reflected [`States`] trait of a type.
+///
+/// A [`ReflectState`] for type `T` can be obtained via
+/// [`bevy_reflect::TypeRegistration::data`].
+#[derive(Clone)]
+pub struct ReflectState(ReflectStateFns);
+
+/// The raw function pointers needed to make up a [`ReflectState`].
+#[derive(Clone)]
+pub struct ReflectStateFns {
+    /// Function pointer implementing [`ReflectState::reflect()`].
+    pub reflect: fn(&World) -> Option<&dyn Reflect>,
+}
+
+impl ReflectStateFns {
+    /// Get the default set of [`ReflectStateFns`] for a specific component type using its
+    /// [`FromType`] implementation.
+    ///
+    /// This is useful if you want to start with the default implementation before overriding some
+    /// of the functions to create a custom implementation.
+    pub fn new<T: States + Reflect>() -> Self {
+        <ReflectState as FromType<T>>::from_type().0
+    }
+}
+
+impl ReflectState {
+    /// Gets the value of this [`States`] type from the world as a reflected reference.
+    pub fn reflect<'a>(&self, world: &'a World) -> Option<&'a dyn Reflect> {
+        (self.0.reflect)(world)
+    }
+}
+
+impl<S: States + Reflect> FromType<S> for ReflectState {
+    fn from_type() -> Self {
+        ReflectState(ReflectStateFns {
+            reflect: |world| {
+                world
+                    .get_resource::<State<S>>()
+                    .map(|res| res.get() as &dyn Reflect)
+            },
+        })
+    }
+}
+
+/// A struct used to operate on the reflected [`FreelyMutableState`] trait of a type.
+///
+/// A [`ReflectFreelyMutableState`] for type `T` can be obtained via
+/// [`bevy_reflect::TypeRegistration::data`].
+#[derive(Clone)]
+pub struct ReflectFreelyMutableState(ReflectFreelyMutableStateFns);
+
+/// The raw function pointers needed to make up a [`ReflectFreelyMutableState`].
+#[derive(Clone)]
+pub struct ReflectFreelyMutableStateFns {
+    /// Function pointer implementing [`ReflectFreelyMutableState::set_next_state()`].
+    pub set_next_state: fn(&mut World, &dyn Reflect, &TypeRegistry),
+}
+
+impl ReflectFreelyMutableStateFns {
+    /// Get the default set of [`ReflectFreelyMutableStateFns`] for a specific component type using its
+    /// [`FromType`] implementation.
+    ///
+    /// This is useful if you want to start with the default implementation before overriding some
+    /// of the functions to create a custom implementation.
+    pub fn new<T: FreelyMutableState + Reflect>() -> Self {
+        <ReflectFreelyMutableState as FromType<T>>::from_type().0
+    }
+}
+
+impl ReflectFreelyMutableState {
+    /// Tentatively set a pending state transition to a reflected [`ReflectFreelyMutableState`].
+    pub fn set_next_state(&self, world: &mut World, state: &dyn Reflect, registry: &TypeRegistry) {
+        (self.0.set_next_state)(world, state, registry);
+    }
+}
+
+impl<S: FreelyMutableState + Reflect> FromType<S> for ReflectFreelyMutableState {
+    fn from_type() -> Self {
+        ReflectFreelyMutableState(ReflectFreelyMutableStateFns {
+            set_next_state: |world, reflected_state, registry| {
+                let new_state: S = from_reflect_with_fallback(reflected_state, world, registry);
+                if let Some(mut next_state) = world.get_resource_mut::<NextState<S>>() {
+                    next_state.set(new_state);
+                }
+            },
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate as bevy_state;
+    use crate::app::{AppExtStates, StatesPlugin};
+    use crate::reflect::{ReflectFreelyMutableState, ReflectState};
+    use crate::state::State;
+    use bevy_app::App;
+    use bevy_ecs::prelude::AppTypeRegistry;
+    use bevy_reflect::Reflect;
+    use bevy_state_macros::States;
+    use std::any::TypeId;
+
+    #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, States, Reflect)]
+    enum StateTest {
+        A,
+        B,
+    }
+
+    #[test]
+    fn test_reflect_state_operations() {
+        let mut app = App::new();
+        app.add_plugins(StatesPlugin)
+            .insert_state(StateTest::A)
+            .register_type_mutable_state::<StateTest>();
+
+        let type_registry = app.world_mut().resource::<AppTypeRegistry>().0.clone();
+        let type_registry = type_registry.read();
+
+        let (reflect_state, reflect_mutable_state) = (
+            type_registry
+                .get_type_data::<ReflectState>(TypeId::of::<StateTest>())
+                .unwrap()
+                .clone(),
+            type_registry
+                .get_type_data::<ReflectFreelyMutableState>(TypeId::of::<StateTest>())
+                .unwrap()
+                .clone(),
+        );
+
+        let current_value = reflect_state.reflect(app.world()).unwrap();
+        assert_eq!(
+            current_value.downcast_ref::<StateTest>().unwrap(),
+            &StateTest::A
+        );
+
+        reflect_mutable_state.set_next_state(app.world_mut(), &StateTest::B, &type_registry);
+
+        assert_ne!(
+            app.world().resource::<State<StateTest>>().get(),
+            &StateTest::B
+        );
+
+        app.update();
+
+        assert_eq!(
+            app.world().resource::<State<StateTest>>().get(),
+            &StateTest::B
+        );
+
+        let current_value = reflect_state.reflect(app.world()).unwrap();
+        assert_eq!(
+            current_value.downcast_ref::<StateTest>().unwrap(),
+            &StateTest::B
+        );
+    }
+}

crates/bevy_state/src/state/resources.rs

@@ -11,6 +11,9 @@ use super::{freely_mutable_state::FreelyMutableState, states::States};
 #[cfg(feature = "bevy_reflect")]
 use bevy_ecs::prelude::ReflectResource;
 
+#[cfg(feature = "bevy_reflect")]
+use bevy_reflect::prelude::ReflectDefault;
+
 /// A finite-state machine whose transitions have associated schedules
 /// ([`OnEnter(state)`](crate::state::OnEnter) and [`OnExit(state)`](crate::state::OnExit)).
 ///
@@ -115,7 +118,7 @@ impl<S: States> Deref for State<S> {
 #[cfg_attr(
     feature = "bevy_reflect",
     derive(bevy_reflect::Reflect),
-    reflect(Resource)
+    reflect(Resource, Default)
 )]
 pub enum NextState<S: FreelyMutableState> {
     /// No state transition is pending