feat: Add World::get_reflect() and World::get_reflect_mut() (#14416)

# Objective Sometimes one wants to retrieve a `&dyn Reflect` for an entity's component, which so far required multiple, non-obvious steps and `unsafe`-code. The docs for [`MutUntyped`](https://docs.rs/bevy/latest/bevy/ecs/change_detection/struct.MutUntyped.html#method.map_unchanged) contain an example of the unsafe part. ## Solution This PR adds the two methods: ```rust // immutable variant World::get_reflect(&self, entity: Entity, type_id: TypeId) -> Result<&dyn Reflect, GetComponentReflectError> // mutable variant World::get_reflect_mut(&mut self, entity: Entity, type_id: TypeId) -> Result<Mut<'_, dyn Reflect>, GetComponentReflectError> ``` which take care of the necessary steps, check required invariants etc., and contain the unsafety so the caller doesn't have to deal with it. ## Testing - Did you test these changes? If so, how? - Added tests and a doc test, also (successfully) ran `cargo run -p ci`. - Are there any parts that need more testing? - Could add tests for each individual error variant, but it's not required imo. - How can other people (reviewers) test your changes? Is there anything specific they need to know? - Run `cargo test --doc --package bevy_ecs --all-features -- world::World::get_reflect --show-output` for the doctest - Run `cargo test --package bevy_ecs --lib --all-features -- world::tests::reflect_tests --show-output` for the unittests - If relevant, what platforms did you test these changes on, and are there any important ones you can't test? - Don't think it's relevant, but tested on 64bit linux (only). --- ## Showcase Copy of the doctest example which gives a good overview of what this enables: ```rust use bevy_ecs::prelude::*; use bevy_reflect::Reflect; use std::any::TypeId; // define a `Component` and derive `Reflect` for it #[derive(Component, Reflect)] struct MyComponent; // create a `World` for this example let mut world = World::new(); // Note: This is usually handled by `App::register_type()`, but this example can not use `App`. world.init_resource::<AppTypeRegistry>(); world.get_resource_mut::<AppTypeRegistry>().unwrap().write().register::<MyComponent>(); // spawn an entity with a `MyComponent` let entity = world.spawn(MyComponent).id(); // retrieve a reflected reference to the entity's `MyComponent` let comp_reflected: &dyn Reflect = world.get_reflect(entity, TypeId::of::<MyComponent>()).unwrap(); // make sure we got the expected type assert!(comp_reflected.is::<MyComponent>()); ``` ## Migration Guide No breaking changes, but users can use the new methods if they did it manually before. --------- Co-authored-by: Gino Valente <49806985+MrGVSV@users.noreply.github.com>
2024-07-23 23:57:54 +07:00 · 2024-07-23 23:57:54 +07:00 · abceebebba
commit abceebebba
parent 8a79185880
2 changed files with 340 additions and 0 deletions
--- a/crates/bevy_ecs/src/world/mod.rs
+++ b/crates/bevy_ecs/src/world/mod.rs
@ -9,6 +9,9 @@ mod identifier;
 mod spawn_batch;
 pub mod unsafe_world_cell;

+#[cfg(feature = "bevy_reflect")]
+pub mod reflect;
+
 pub use crate::{
    change_detection::{Mut, Ref, CHECK_TICK_THRESHOLD},
    world::command_queue::CommandQueue,
--- a/crates/bevy_ecs/src/world/reflect.rs
+++ b/crates/bevy_ecs/src/world/reflect.rs
@ -0,0 +1,337 @@
+//! Provides additional functionality for [`World`] when the `bevy_reflect` feature is enabled.
+
+use core::any::TypeId;
+
+use thiserror::Error;
+
+use bevy_reflect::{Reflect, ReflectFromPtr};
+
+use crate::prelude::*;
+use crate::world::ComponentId;
+
+impl World {
+    /// Retrieves a reference to the given `entity`'s [`Component`] of the given `type_id` using
+    /// reflection.
+    ///
+    /// Requires implementing [`Reflect`] for the [`Component`] (e.g., using [`#[derive(Reflect)`](derive@bevy_reflect::Reflect))
+    /// and `app.register_type::<TheComponent>()` to have been called[^note-reflect-impl].
+    ///
+    /// If you want to call this with a [`ComponentId`], see [`World::components`] and [`Components::get_id`] to get
+    /// the corresponding [`TypeId`].
+    ///
+    /// Also see the crate documentation for [`bevy_reflect`] for more information on
+    /// [`Reflect`] and bevy's reflection capabilities.
+    ///
+    /// # Errors
+    ///
+    /// See [`GetComponentReflectError`] for the possible errors and their descriptions.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use bevy_ecs::prelude::*;
+    /// use bevy_reflect::Reflect;
+    /// use std::any::TypeId;
+    ///
+    /// // define a `Component` and derive `Reflect` for it
+    /// #[derive(Component, Reflect)]
+    /// struct MyComponent;
+    ///
+    /// // create a `World` for this example
+    /// let mut world = World::new();
+    ///
+    /// // Note: This is usually handled by `App::register_type()`, but this example cannot use `App`.
+    /// world.init_resource::<AppTypeRegistry>();
+    /// world.get_resource_mut::<AppTypeRegistry>().unwrap().write().register::<MyComponent>();
+    ///
+    /// // spawn an entity with a `MyComponent`
+    /// let entity = world.spawn(MyComponent).id();
+    ///
+    /// // retrieve a reflected reference to the entity's `MyComponent`
+    /// let comp_reflected: &dyn Reflect = world.get_reflect(entity, TypeId::of::<MyComponent>()).unwrap();
+    ///
+    /// // make sure we got the expected type
+    /// assert!(comp_reflected.is::<MyComponent>());
+    /// ```
+    ///
+    /// # Note
+    /// Requires the `bevy_reflect` feature (included in the default features).
+    ///
+    /// [`Components::get_id`]: crate::component::Components::get_id
+    /// [`ReflectFromPtr`]: bevy_reflect::ReflectFromPtr
+    /// [`TypeData`]: bevy_reflect::TypeData
+    /// [`Reflect`]: bevy_reflect::Reflect
+    /// [`App::register_type`]: ../../bevy_app/struct.App.html#method.register_type
+    /// [^note-reflect-impl]: More specifically: Requires [`TypeData`] for [`ReflectFromPtr`] to be registered for the given `type_id`,
+    ///     which is automatically handled when deriving [`Reflect`] and calling [`App::register_type`].
+    #[inline]
+    pub fn get_reflect(
+        &self,
+        entity: Entity,
+        type_id: TypeId,
+    ) -> Result<&dyn Reflect, GetComponentReflectError> {
+        let Some(component_id) = self.components().get_id(type_id) else {
+            return Err(GetComponentReflectError::NoCorrespondingComponentId(
+                type_id,
+            ));
+        };
+
+        let Some(comp_ptr) = self.get_by_id(entity, component_id) else {
+            let component_name = self
+                .components()
+                .get_name(component_id)
+                .map(ToString::to_string);
+
+            return Err(GetComponentReflectError::EntityDoesNotHaveComponent {
+                entity,
+                type_id,
+                component_id,
+                component_name,
+            });
+        };
+
+        let Some(type_registry) = self.get_resource::<AppTypeRegistry>().map(|atr| atr.read())
+        else {
+            return Err(GetComponentReflectError::MissingAppTypeRegistry);
+        };
+
+        let Some(reflect_from_ptr) = type_registry.get_type_data::<ReflectFromPtr>(type_id) else {
+            return Err(GetComponentReflectError::MissingReflectFromPtrTypeData(
+                type_id,
+            ));
+        };
+
+        // SAFETY:
+        // - `comp_ptr` is guaranteed to point to an object of type `type_id`
+        // - `reflect_from_ptr` was constructed for type `type_id`
+        // - Assertion that checks this equality is present
+        unsafe {
+            assert_eq!(
+                reflect_from_ptr.type_id(),
+                type_id,
+                "Mismatch between Ptr's type_id and ReflectFromPtr's type_id",
+            );
+
+            Ok(reflect_from_ptr.as_reflect(comp_ptr))
+        }
+    }
+
+    /// Retrieves a mutable reference to the given `entity`'s [`Component`] of the given `type_id` using
+    /// reflection.
+    ///
+    /// Requires implementing [`Reflect`] for the [`Component`] (e.g., using [`#[derive(Reflect)`](derive@bevy_reflect::Reflect))
+    /// and `app.register_type::<TheComponent>()` to have been called.
+    ///
+    /// This is the mutable version of [`World::get_reflect`], see its docs for more information
+    /// and an example.
+    ///
+    /// Just calling this method does not trigger [change detection](crate::change_detection).
+    ///
+    /// # Errors
+    ///
+    /// See [`GetComponentReflectError`] for the possible errors and their descriptions.
+    ///
+    /// # Example
+    ///
+    /// See the documentation for [`World::get_reflect`].
+    ///
+    /// # Note
+    /// Requires the feature `bevy_reflect` (included in the default features).
+    ///
+    /// [`Reflect`]: bevy_reflect::Reflect
+    #[inline]
+    pub fn get_reflect_mut(
+        &mut self,
+        entity: Entity,
+        type_id: TypeId,
+    ) -> Result<Mut<'_, dyn Reflect>, GetComponentReflectError> {
+        // little clone() + read() dance so we a) don't keep a borrow of `self` and b) don't drop a
+        // temporary (from read()) too  early.
+        let Some(app_type_registry) = self.get_resource::<AppTypeRegistry>().cloned() else {
+            return Err(GetComponentReflectError::MissingAppTypeRegistry);
+        };
+        let type_registry = app_type_registry.read();
+
+        let Some(reflect_from_ptr) = type_registry.get_type_data::<ReflectFromPtr>(type_id) else {
+            return Err(GetComponentReflectError::MissingReflectFromPtrTypeData(
+                type_id,
+            ));
+        };
+
+        let Some(component_id) = self.components().get_id(type_id) else {
+            return Err(GetComponentReflectError::NoCorrespondingComponentId(
+                type_id,
+            ));
+        };
+
+        // HACK: Only required for the `None`-case/`else`-branch, but it borrows `self`, which will
+        // already be mutablyy borrowed by `self.get_mut_by_id()`, and I didn't find a way around it.
+        let component_name = self
+            .components()
+            .get_name(component_id)
+            .map(ToString::to_string);
+
+        let Some(comp_mut_untyped) = self.get_mut_by_id(entity, component_id) else {
+            return Err(GetComponentReflectError::EntityDoesNotHaveComponent {
+                entity,
+                type_id,
+                component_id,
+                component_name,
+            });
+        };
+
+        // SAFETY:
+        // - `comp_mut_untyped` is guaranteed to point to an object of type `type_id`
+        // - `reflect_from_ptr` was constructed for type `type_id`
+        // - Assertion that checks this equality is present
+        let comp_mut_typed = comp_mut_untyped.map_unchanged(|ptr_mut| unsafe {
+            assert_eq!(
+                reflect_from_ptr.type_id(),
+                type_id,
+                "Mismatch between PtrMut's type_id and ReflectFromPtr's type_id",
+            );
+
+            reflect_from_ptr.as_reflect_mut(ptr_mut)
+        });
+
+        Ok(comp_mut_typed)
+    }
+}
+
+/// The error type returned by [`World::get_reflect`] and [`World::get_reflect_mut`].
+#[derive(Error, Debug)]
+pub enum GetComponentReflectError {
+    /// There is no [`ComponentId`] corresponding to the given [`TypeId`].
+    ///
+    /// This is usually handled by calling [`App::register_type`] for the type corresponding to
+    /// the given [`TypeId`].
+    ///
+    /// See the documentation for [`bevy_reflect`] for more information.
+    ///
+    /// [`App::register_type`]: ../../../bevy_app/struct.App.html#method.register_type
+    #[error("No `ComponentId` corresponding to {0:?} found (did you call App::register_type()?)")]
+    NoCorrespondingComponentId(TypeId),
+
+    /// The given [`Entity`] does not have a [`Component`] corresponding to the given [`TypeId`].
+    #[error("The given `Entity` {entity:?} does not have a `{component_name:?}` component ({component_id:?}, which corresponds to {type_id:?})")]
+    EntityDoesNotHaveComponent {
+        /// The given [`Entity`].
+        entity: Entity,
+        /// The given [`TypeId`].
+        type_id: TypeId,
+        /// The [`ComponentId`] corresponding to the given [`TypeId`].
+        component_id: ComponentId,
+        /// The name corresponding to the [`Component`] with the given [`TypeId`], or `None`
+        /// if not available.
+        component_name: Option<String>,
+    },
+
+    /// The [`World`] was missing the [`AppTypeRegistry`] resource.
+    #[error("The `World` was missing the `AppTypeRegistry` resource")]
+    MissingAppTypeRegistry,
+
+    /// The [`World`]'s [`TypeRegistry`] did not contain [`TypeData`] for [`ReflectFromPtr`] for the given [`TypeId`].
+    ///
+    /// This is usually handled by calling [`App::register_type`] for the type corresponding to
+    /// the given [`TypeId`].
+    ///
+    /// See the documentation for [`bevy_reflect`] for more information.
+    ///
+    /// [`TypeData`]: bevy_reflect::TypeData
+    /// [`TypeRegistry`]: bevy_reflect::TypeRegistry
+    /// [`ReflectFromPtr`]: bevy_reflect::ReflectFromPtr
+    /// [`App::register_type`]: ../../../bevy_app/struct.App.html#method.register_type
+    #[error("The `World`'s `TypeRegistry` did not contain `TypeData` for `ReflectFromPtr` for the given {0:?} (did you call `App::register_type()`?)")]
+    MissingReflectFromPtrTypeData(TypeId),
+}
+
+#[cfg(test)]
+mod tests {
+    use std::any::TypeId;
+
+    use bevy_reflect::Reflect;
+
+    use crate::{
+        // For bevy_ecs_macros
+        self as bevy_ecs,
+        prelude::{AppTypeRegistry, Component, DetectChanges, World},
+    };
+
+    #[derive(Component, Reflect)]
+    struct RFoo(i32);
+
+    #[derive(Component)]
+    struct Bar;
+
+    #[test]
+    fn get_component_as_reflect() {
+        let mut world = World::new();
+        world.init_resource::<AppTypeRegistry>();
+
+        let app_type_registry = world.get_resource_mut::<AppTypeRegistry>().unwrap();
+        app_type_registry.write().register::<RFoo>();
+
+        {
+            let entity_with_rfoo = world.spawn(RFoo(42)).id();
+            let comp_reflect = world
+                .get_reflect(entity_with_rfoo, TypeId::of::<RFoo>())
+                .expect("Reflection of RFoo-component failed");
+
+            assert!(comp_reflect.is::<RFoo>());
+        }
+
+        {
+            let entity_without_rfoo = world.spawn_empty().id();
+            let reflect_opt = world.get_reflect(entity_without_rfoo, TypeId::of::<RFoo>());
+
+            assert!(reflect_opt.is_err());
+        }
+
+        {
+            let entity_with_bar = world.spawn(Bar).id();
+            let reflect_opt = world.get_reflect(entity_with_bar, TypeId::of::<Bar>());
+
+            assert!(reflect_opt.is_err());
+        }
+    }
+
+    #[test]
+    fn get_component_as_mut_reflect() {
+        let mut world = World::new();
+        world.init_resource::<AppTypeRegistry>();
+
+        let app_type_registry = world.get_resource_mut::<AppTypeRegistry>().unwrap();
+        app_type_registry.write().register::<RFoo>();
+
+        {
+            let entity_with_rfoo = world.spawn(RFoo(42)).id();
+            let mut comp_reflect = world
+                .get_reflect_mut(entity_with_rfoo, TypeId::of::<RFoo>())
+                .expect("Mutable reflection of RFoo-component failed");
+
+            let comp_rfoo_reflected = comp_reflect
+                .downcast_mut::<RFoo>()
+                .expect("Wrong type reflected (expected RFoo)");
+            assert_eq!(comp_rfoo_reflected.0, 42);
+            comp_rfoo_reflected.0 = 1337;
+
+            let rfoo_ref = world.entity(entity_with_rfoo).get_ref::<RFoo>().unwrap();
+            assert!(rfoo_ref.is_changed());
+            assert_eq!(rfoo_ref.0, 1337);
+        }
+
+        {
+            let entity_without_rfoo = world.spawn_empty().id();
+            let reflect_opt = world.get_reflect_mut(entity_without_rfoo, TypeId::of::<RFoo>());
+
+            assert!(reflect_opt.is_err());
+        }
+
+        {
+            let entity_with_bar = world.spawn(Bar).id();
+            let reflect_opt = world.get_reflect_mut(entity_with_bar, TypeId::of::<Bar>());
+
+            assert!(reflect_opt.is_err());
+        }
+    }
+}