Automatic registration of ECS types

crates/bevy_ecs/Cargo.toml

@@ -13,6 +13,7 @@ categories = ["game-engines", "data-structures"]
 trace = []
 multi-threaded = ["bevy_tasks/multi-threaded"]
 bevy_debug_stepping = []
+bevy_reflect = ["dep:bevy_reflect", "bevy_ecs_macros/bevy_reflect"]
 default = ["bevy_reflect", "bevy_debug_stepping"]
 
 [dependencies]

crates/bevy_ecs/macros/Cargo.toml

@@ -8,6 +8,9 @@ license = "MIT OR Apache-2.0"
 [lib]
 proc-macro = true
 
+[features]
+bevy_reflect = []
+
 [dependencies]
 bevy_macro_utils = { path = "../../bevy_macro_utils", version = "0.14.0-dev" }
 

crates/bevy_ecs/macros/src/component.rs

@@ -1,7 +1,7 @@
 use proc_macro::TokenStream;
 use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::quote;
-use syn::{parse_macro_input, parse_quote, DeriveInput, Ident, LitStr, Path, Result};
+use syn::{parse_macro_input, parse_quote, DeriveInput, Ident, LitStr, Meta, Path, Result};
 
 pub fn derive_event(input: TokenStream) -> TokenStream {
     let mut ast = parse_macro_input!(input as DeriveInput);
@@ -32,9 +32,11 @@ pub fn derive_resource(input: TokenStream) -> TokenStream {
 
     let struct_name = &ast.ident;
     let (impl_generics, type_generics, where_clause) = &ast.generics.split_for_impl();
+    let register_type = get_type_registration(&ast, &bevy_ecs_path, "Resource");
 
     TokenStream::from(quote! {
         impl #impl_generics #bevy_ecs_path::system::Resource for #struct_name #type_generics #where_clause {
+            #register_type
         }
     })
 }
@@ -57,16 +59,19 @@ pub fn derive_component(input: TokenStream) -> TokenStream {
 
     let struct_name = &ast.ident;
     let (impl_generics, type_generics, where_clause) = &ast.generics.split_for_impl();
+    let register_type = get_type_registration(&ast, &bevy_ecs_path, "Component");
 
     TokenStream::from(quote! {
         impl #impl_generics #bevy_ecs_path::component::Component for #struct_name #type_generics #where_clause {
             const STORAGE_TYPE: #bevy_ecs_path::component::StorageType = #storage;
+            #register_type
         }
     })
 }
 
 pub const COMPONENT: &str = "component";
 pub const STORAGE: &str = "storage";
+const REFLECT: &str = "reflect";
 
 struct Attrs {
     storage: StorageTy,
@@ -109,6 +114,37 @@ fn parse_component_attr(ast: &DeriveInput) -> Result<Attrs> {
     Ok(attrs)
 }
 
+fn get_type_registration(
+    ast: &DeriveInput,
+    bevy_ecs_path: &Path,
+    reflect_trait: &'static str,
+) -> Option<proc_macro2::TokenStream> {
+    // Generics require the generic parameters to implement Reflect
+    // This is unsupported for now.
+    if !cfg!(feature = "bevy_reflect") || !ast.generics.params.is_empty() {
+        return None;
+    }
+
+    ast.attrs.iter().any(|attr| {
+        if !attr.path().is_ident(REFLECT) || !matches!(attr.meta, Meta::List(_)) {
+            return false;
+        }
+
+        attr.parse_nested_meta(|meta| {
+            meta.path.is_ident(reflect_trait)
+                .then_some(())
+                .ok_or_else(|| meta.error("missing required reflect attribute"))
+        })
+        .is_ok()
+    })
+    .then(|| quote! {
+        #[doc(hidden)]
+        fn __register_type(registry: &#bevy_ecs_path::private::bevy_reflect::TypeRegistryArc) {
+            #bevy_ecs_path::private::bevy_reflect::register_type_shim::<Self>(registry);
+        }
+    })
+}
+
 fn storage_path(bevy_ecs_path: &Path, ty: StorageTy) -> TokenStream2 {
     let storage_type = match ty {
         StorageTy::Table => Ident::new("Table", Span::call_site()),

crates/bevy_ecs/src/component.rs

@@ -12,7 +12,7 @@ use crate::{
 pub use bevy_ecs_macros::Component;
 use bevy_ptr::{OwningPtr, UnsafeCellDeref};
 #[cfg(feature = "bevy_reflect")]
-use bevy_reflect::Reflect;
+use bevy_reflect::{Reflect, TypeRegistryArc};
 use bevy_utils::TypeIdMap;
 use std::cell::UnsafeCell;
 use std::{
@@ -156,6 +156,13 @@ pub trait Component: Send + Sync + 'static {
 
     /// Called when registering this component, allowing mutable access to it's [`ComponentHooks`].
     fn register_component_hooks(_hooks: &mut ComponentHooks) {}
+
+    /// Shim for automatically registering components. Intentionally hidden as it's not part of the
+    /// public interface. Only available if the `bevy_reflect` feature is enabled.
+    #[doc(hidden)]
+    #[allow(unused_variables)]
+    #[cfg(feature = "bevy_reflect")]
+    fn __register_type(registry: &TypeRegistryArc) {}
 }
 
 /// The storage used for a specific component type.
@@ -525,6 +532,8 @@ pub struct Components {
     components: Vec<ComponentInfo>,
     indices: TypeIdMap<ComponentId>,
     resource_indices: TypeIdMap<ComponentId>,
+    #[cfg(feature = "bevy_reflect")]
+    pub(crate) type_registry: TypeRegistryArc,
 }
 
 impl Components {
@@ -552,6 +561,8 @@ impl Components {
                 ComponentDescriptor::new::<T>(),
             );
             T::register_component_hooks(&mut components[index.index()].hooks);
+            #[cfg(feature = "bevy_reflect")]
+            T::__register_type(&self.type_registry);
             index
         })
     }
@@ -720,9 +731,16 @@ impl Components {
     pub fn init_resource<T: Resource>(&mut self) -> ComponentId {
         // SAFETY: The [`ComponentDescriptor`] matches the [`TypeId`]
         unsafe {
-            self.get_or_insert_resource_with(TypeId::of::<T>(), || {
-                ComponentDescriptor::new_resource::<T>()
-            })
+            Self::get_or_insert_resource_with(
+                &mut self.components,
+                &mut self.resource_indices,
+                TypeId::of::<T>(),
+                || {
+                    #[cfg(feature = "bevy_reflect")]
+                    T::__register_type(&self.type_registry);
+                    ComponentDescriptor::new_resource::<T>()
+                },
+            )
         }
     }
 
@@ -733,9 +751,12 @@ impl Components {
     pub fn init_non_send<T: Any>(&mut self) -> ComponentId {
         // SAFETY: The [`ComponentDescriptor`] matches the [`TypeId`]
         unsafe {
-            self.get_or_insert_resource_with(TypeId::of::<T>(), || {
-                ComponentDescriptor::new_non_send::<T>(StorageType::default())
-            })
+            Self::get_or_insert_resource_with(
+                &mut self.components,
+                &mut self.resource_indices,
+                TypeId::of::<T>(),
+                || ComponentDescriptor::new_non_send::<T>(StorageType::default()),
+            )
         }
     }
 
@@ -744,12 +765,12 @@ impl Components {
     /// The [`ComponentDescriptor`] must match the [`TypeId`]
     #[inline]
     unsafe fn get_or_insert_resource_with(
-        &mut self,
+        components: &mut Vec<ComponentInfo>,
+        resource_indices: &mut TypeIdMap<ComponentId>,
         type_id: TypeId,
         func: impl FnOnce() -> ComponentDescriptor,
     ) -> ComponentId {
-        let components = &mut self.components;
-        *self.resource_indices.entry(type_id).or_insert_with(|| {
+        *resource_indices.entry(type_id).or_insert_with(|| {
             let descriptor = func();
             let component_id = ComponentId(components.len());
             components.push(ComponentInfo::new(component_id, descriptor));
@@ -971,3 +992,40 @@ impl<T: Component> FromWorld for InitComponentId<T> {
         }
     }
 }
+
+#[cfg(test)]
+#[cfg(feature = "bevy_reflect")]
+mod test {
+    use crate as bevy_ecs;
+    use crate::{
+        component::Components, prelude::Component, prelude::Resource, reflect::ReflectComponent,
+        reflect::ReflectResource, storage::Storages,
+    };
+    use bevy_reflect::Reflect;
+    use core::any::TypeId;
+
+    #[test]
+    fn init_component_registers_component() {
+        #[derive(Component, Reflect)]
+        #[reflect(Component)]
+        struct A(usize);
+
+        let mut components = Components::default();
+        let mut storages = Storages::default();
+        components.init_component::<A>(&mut storages);
+
+        assert!(components.type_registry.read().contains(TypeId::of::<A>()));
+    }
+
+    #[test]
+    fn init_resource_registers_resource() {
+        #[derive(Resource, Reflect)]
+        #[reflect(Resource)]
+        struct A(usize);
+
+        let mut components = Components::default();
+        components.init_resource::<A>();
+
+        assert!(components.type_registry.read().contains(TypeId::of::<A>()));
+    }
+}

crates/bevy_ecs/src/lib.rs

@@ -52,6 +52,15 @@ pub mod prelude {
     };
 }
 
+#[cfg(feature = "bevy_reflect")]
+#[doc(hidden)]
+pub mod private {
+    pub mod bevy_reflect {
+        pub use crate::reflect::register_type_shim;
+        pub use bevy_reflect::TypeRegistryArc;
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate as bevy_ecs;

crates/bevy_ecs/src/reflect/entity_commands.rs

@@ -338,13 +338,12 @@ mod tests {
     fn insert_reflected() {
         let mut world = World::new();
 
-        let type_registry = AppTypeRegistry::default();
+        world.init_resource::<AppTypeRegistry>();
         {
-            let mut registry = type_registry.write();
+            let mut registry = world.resource::<AppTypeRegistry>().write();
             registry.register::<ComponentA>();
             registry.register_type_data::<ComponentA, ReflectComponent>();
         }
-        world.insert_resource(type_registry);
 
         let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
         let mut commands = system_state.get_mut(&mut world);
@@ -405,13 +404,7 @@ mod tests {
     fn remove_reflected() {
         let mut world = World::new();
 
-        let type_registry = AppTypeRegistry::default();
-        {
-            let mut registry = type_registry.write();
-            registry.register::<ComponentA>();
-            registry.register_type_data::<ComponentA, ReflectComponent>();
-        }
-        world.insert_resource(type_registry);
+        world.init_resource::<AppTypeRegistry>();
 
         let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
         let mut commands = system_state.get_mut(&mut world);

crates/bevy_ecs/src/reflect/mod.rs

@@ -1,11 +1,58 @@
 //! Types that enable reflection support.
+//!
+//! # Automatic Reflect Type Registration
+//! The [`Component`] and [`Resource`] resource derive macros will automatically register
+//! its types that implement [`Reflect`] into the [`AppTypeRegistry`] resource when first
+//! seen by the ECS [`World`].
+//!
+//! If you find that your component or resource is not registered, they may need to be manually
+//! registered. There are a few exceptions:
+//!
+//!  * Automatic registration is only supported via the derive macros. Manual implementations of
+//!    `Component`, `Resource`, or `Reflect` must be manually registered.
+//!  * The associated ECS trait must be reflected (via `reflect(Component)` or `reflect(Resource)`).
+//!  * Generic types are not supported, and must be manually registered.
+//!  * Types are registered when the World first initializes the type. This may cause registrations
+//!    to be missing due to mistiming. These initialization points include but are not limited to:
+//!    - spawning an entity with the component or inserting the resource
+//!    - inserting the component existing entity
+//!    - attempting to remove the component or resource, even if it's not present.
+//!    - a system that references the component or resource is added to a schedule
+//!
+//! ```rust
+//! use bevy_ecs::prelude::*;
+//! use bevy_reflect::Reflect;
+//!
+//! // This will automatically register upon first use!
+//! #[derive(Component, Reflect)]
+//! #[reflect(Component)]
+//! pub struct MyComponent {
+//!     a: usize,
+//!     b: (u32, u8)
+//! }
+//!
+//! // This won't!
+//! #[derive(Component, Reflect)]
+//! #[reflect(Component)]
+//! pub struct GenericComponent<T>(T);
+//!
+//! // This won't!
+//! #[derive(Component, Reflect)]
+//! pub struct NoReflectComponent;
+//! ```
+//!
+//! [`Component`]: crate::prelude::Component
+//! [`Resource`]: crate::prelude::Resource
 
 use std::any::TypeId;
 use std::ops::{Deref, DerefMut};
 
 use crate as bevy_ecs;
-use crate::{system::Resource, world::World};
-use bevy_reflect::{FromReflect, Reflect, TypeRegistry, TypeRegistryArc};
+use crate::{
+    system::Resource,
+    world::{FromWorld, World},
+};
+use bevy_reflect::{FromReflect, GetTypeRegistration, Reflect, TypeRegistry, TypeRegistryArc};
 
 mod bundle;
 mod component;
@@ -21,10 +68,27 @@ pub use from_world::{ReflectFromWorld, ReflectFromWorldFns};
 pub use map_entities::ReflectMapEntities;
 pub use resource::{ReflectResource, ReflectResourceFns};
 
+#[doc(hidden)]
+pub fn register_type_shim<T: GetTypeRegistration>(registry: &TypeRegistryArc) {
+    if let Ok(mut registry) = registry.internal.try_write() {
+        registry.register::<T>();
+        return;
+    }
+    if let Ok(registry) = registry.internal.try_read() {
+        if registry.contains(::core::any::TypeId::of::<T>()) {
+            return;
+        }
+    }
+    panic!(
+        "Deadlock while registering <{}>.",
+        ::std::any::type_name::<T>()
+    );
+}
+
 /// A [`Resource`] storing [`TypeRegistry`] for
 /// type registrations relevant to a whole app.
-#[derive(Resource, Clone, Default)]
-pub struct AppTypeRegistry(pub TypeRegistryArc);
+#[derive(Resource, Clone, Debug)]
+pub struct AppTypeRegistry(TypeRegistryArc);
 
 impl Deref for AppTypeRegistry {
     type Target = TypeRegistryArc;
@@ -42,6 +106,12 @@ impl DerefMut for AppTypeRegistry {
     }
 }
 
+impl FromWorld for AppTypeRegistry {
+    fn from_world(world: &mut World) -> Self {
+        Self(world.type_registry().clone())
+    }
+}
+
 /// Creates a `T` from a `&dyn Reflect`.
 ///
 /// The first approach uses `T`'s implementation of `FromReflect`.