Hybrid Resources

crates/bevy_ecs/macros/src/lib.rs

@@ -588,6 +588,13 @@ pub fn derive_resource(input: TokenStream) -> TokenStream {
     TokenStream::from(resource::derive_resource(&mut ast))
 }
 
+/// TODO: Documentation
+#[proc_macro_derive(HybridResource, attributes(component, require))]
+pub fn derive_hybrid_resource(input: TokenStream) -> TokenStream {
+    let mut ast = parse_macro_input!(input as DeriveInput);
+    TokenStream::from(resource::derive_hybrid_resource(&mut ast))
+}
+
 /// Cheat sheet for derive syntax,
 ///
 /// ## Group Override

crates/bevy_ecs/macros/src/resource.rs

@@ -21,6 +21,44 @@ pub fn derive_resource(ast: &mut DeriveInput) -> TokenStream {
         } else {
             required_components.components_registrator().register_component::<#struct_name #type_generics>()
         };
+        required_components.register_required::<#bevy_ecs::resource::IsUnique>(move || #bevy_ecs::resource::IsUnique::new(resource_component_id));
+        required_components.register_required::<#bevy_ecs::resource::IsResource>(move || #bevy_ecs::resource::IsResource::new(resource_component_id));
+    });
+
+    let component_impl = match derive_component.impl_component(ast, &bevy_ecs, StorageTy::SparseSet)
+    {
+        Ok(value) => value,
+        Err(err) => return err.into_compile_error(),
+    };
+
+    let struct_name = &ast.ident;
+    let (impl_generics, type_generics, where_clause) = &ast.generics.split_for_impl();
+
+    quote! {
+        #component_impl
+        impl #impl_generics #bevy_ecs::resource::Resource for #struct_name #type_generics #where_clause {
+        }
+    }
+}
+
+pub fn derive_hybrid_resource(ast: &mut DeriveInput) -> TokenStream {
+    let bevy_ecs: Path = crate::bevy_ecs_path();
+    let mut derive_component = match DeriveComponent::parse(ast, StorageAttribute::Disallowed) {
+        Ok(value) => value,
+        Err(e) => return e.into_compile_error(),
+    };
+
+    let struct_name = &ast.ident;
+    let (_, type_generics, _) = &ast.generics.split_for_impl();
+
+    // We add the component_id existence check here to avoid recursive init during required components initialization.
+    derive_component.additional_requires.push(quote! {
+        let resource_component_id = if let #FQOption::Some(id) = required_components.components_registrator().component_id::<#struct_name #type_generics>() {
+            id
+        } else {
+            required_components.components_registrator().register_component::<#struct_name #type_generics>()
+        };
+        required_components.register_required::<#bevy_ecs::resource::IsHybridResource>(move || #bevy_ecs::resource::IsHybridResource::new(resource_component_id));
         required_components.register_required::<#bevy_ecs::resource::IsResource>(move || #bevy_ecs::resource::IsResource::new(resource_component_id));
     });
 

crates/bevy_ecs/src/component/constants.rs

@@ -12,3 +12,7 @@ pub const REMOVE: usize = 3;
 pub const DESPAWN: usize = 4;
 /// `usize` of the [`IsResource`](crate::resource::IsResource) component used to mark entities with resources.
 pub const IS_RESOURCE: usize = 5;
+/// TODO
+pub const IS_UNIQUE: usize = 6;
+/// TODO
+pub const IS_HYBRID_RESOURCE: usize = 7;

crates/bevy_ecs/src/resource.rs

@@ -120,7 +120,6 @@ impl ResourceEntities {
 /// A marker component for entities that have a Resource component.
 #[cfg_attr(feature = "bevy_reflect", derive(Reflect), reflect(Component, Debug))]
 #[derive(Component, Debug)]
-#[component(on_insert, on_discard, on_despawn)]
 pub struct IsResource(ComponentId);
 
 impl IsResource {
@@ -133,13 +132,31 @@ impl IsResource {
     pub fn resource_component_id(&self) -> ComponentId {
         self.0
     }
+}
+
+/// A marker component for entities that have a unique component.
+#[cfg_attr(feature = "bevy_reflect", derive(Reflect), reflect(Component, Debug))]
+#[derive(Component, Debug)]
+#[component(on_insert, on_discard, on_despawn)]
+pub struct IsUnique(ComponentId);
+
+impl IsUnique {
+    /// Creates a new instance with the given `component_id`
+    pub fn new(component_id: ComponentId) -> Self {
+        Self(component_id)
+    }
+
+    /// The [`ComponentId`] of the resource component (the _actual_ resource value component, not the [`IsResource`] component).
+    pub fn component_id(&self) -> ComponentId {
+        self.0
+    }
 
     pub(crate) fn on_insert(mut world: DeferredWorld, context: HookContext) {
         let resource_component_id = world
             .entity(context.entity)
             .get::<Self>()
             .unwrap()
-            .resource_component_id();
+            .component_id();
 
         if let Some(original_entity) = world.resource_entities.get(resource_component_id) {
             if !world.entities().contains(original_entity) {
@@ -186,7 +203,81 @@ impl IsResource {
             .entity(context.entity)
             .get::<Self>()
             .unwrap()
-            .resource_component_id();
+            .component_id();
+
+        if let Some(resource_entity) = world.resource_entities.get(resource_component_id)
+            && resource_entity == context.entity
+        {
+            // SAFETY: We have exclusive world access (as long as we don't make structural changes).
+            let cache = unsafe { world.as_unsafe_world_cell().resource_entities() };
+            // SAFETY: There are no shared references to the map.
+            // We only expose `&ResourceCache` to code with access to a resource (such as `&World`),
+            // and that would conflict with the `DeferredWorld` passed to the resource hook.
+            unsafe { &mut *cache.0.get() }.remove(resource_component_id);
+
+            world
+                .commands()
+                .entity(context.entity)
+                .remove_by_id(resource_component_id);
+        }
+    }
+
+    pub(crate) fn on_despawn(_world: DeferredWorld, _context: HookContext) {
+        warn!("Resource entities are not supposed to be despawned.");
+    }
+}
+
+/// A marker component for entities that have a Resource component.
+#[cfg_attr(feature = "bevy_reflect", derive(Reflect), reflect(Component, Debug))]
+#[derive(Component, Debug)]
+#[component(on_insert, on_discard, on_despawn)]
+pub struct IsHybridResource(ComponentId);
+
+impl IsHybridResource {
+    /// Creates a new instance with the given `component_id`
+    pub fn new(component_id: ComponentId) -> Self {
+        Self(component_id)
+    }
+
+    /// The [`ComponentId`] of the resource component (the _actual_ resource value component, not the [`IsResource`] component).
+    pub fn component_id(&self) -> ComponentId {
+        self.0
+    }
+
+    pub(crate) fn on_insert(mut world: DeferredWorld, context: HookContext) {
+        let resource_component_id = world
+            .entity(context.entity)
+            .get::<Self>()
+            .unwrap()
+            .component_id();
+
+        if let Some(original_entity) = world.resource_entities.get(resource_component_id) {
+            if !world.entities().contains(original_entity) {
+                let name = world
+                    .components()
+                    .get_name(resource_component_id)
+                    .expect("resource is registered");
+                panic!(
+                    "Resource entity {} of {} has been despawned, when it's not supposed to be.",
+                    original_entity, name
+                );
+            }
+        } else {
+            // SAFETY: We have exclusive world access (as long as we don't make structural changes).
+            let cache = unsafe { world.as_unsafe_world_cell().resource_entities() };
+            // SAFETY: There are no shared references to the map.
+            // We only expose `&ResourceCache` to code with access to a resource (such as `&World`),
+            // and that would conflict with the `DeferredWorld` passed to the resource hook.
+            unsafe { &mut *cache.0.get() }.insert(resource_component_id, context.entity);
+        }
+    }
+
+    pub(crate) fn on_discard(mut world: DeferredWorld, context: HookContext) {
+        let resource_component_id = world
+            .entity(context.entity)
+            .get::<Self>()
+            .unwrap()
+            .component_id();
 
         if let Some(resource_entity) = world.resource_entities.get(resource_component_id)
             && resource_entity == context.entity
@@ -213,6 +304,15 @@ impl IsResource {
 /// [`ComponentId`] of the [`IsResource`] component.
 pub const IS_RESOURCE: ComponentId = ComponentId::new(crate::component::IS_RESOURCE);
 
+/// TODO
+pub const IS_UNIQUE: ComponentId = ComponentId::new(crate::component::IS_UNIQUE);
+
+/// TODO
+pub const IS_HYBRID_RESOURCE: ComponentId = ComponentId::new(crate::component::IS_HYBRID_RESOURCE);
+
+/// TODO: Documentation
+pub trait HybridResource: Resource {}
+
 #[cfg(test)]
 mod tests {
     use core::sync::atomic::{AtomicBool, Ordering::Relaxed};
@@ -222,11 +322,12 @@ mod tests {
         entity::Entity,
         lifecycle::HookContext,
         ptr::OwningPtr,
-        resource::{IsResource, Resource},
+        resource::{IsResource, IsUnique, Resource},
         world::{DeferredWorld, World},
     };
     use alloc::vec::Vec;
     use bevy_ecs_macros::Component;
+    use bevy_ecs_macros::HybridResource;
     use bevy_platform::prelude::String;
 
     #[test]
@@ -302,7 +403,7 @@ mod tests {
 
         // Removing IsResource should invalidate the current TestResource entity
         // This uses commands because IsResource's despawn-on-removal invalidates the EntityWorldMut and panics
-        world.entity_mut(first_entity).remove::<IsResource>();
+        world.entity_mut(first_entity).remove::<IsUnique>();
         assert!(world.get_resource::<TestResource>().is_none());
 
         assert!(
@@ -327,9 +428,9 @@ mod tests {
         let id = world.spawn(TestResource).id();
         // This spawned resource conflicts with the canonical resource, so it was cleaned up.
         assert!(world.entity(id).get::<TestResource>().is_none());
-        assert!(world.entity(id).get::<IsResource>().is_none());
+        assert!(world.entity(id).get::<IsUnique>().is_none());
         assert!(world.entity(second_entity).get::<TestResource>().is_some());
-        assert!(world.entity(second_entity).get::<IsResource>().is_some());
+        assert!(world.entity(second_entity).get::<IsUnique>().is_some());
     }
 
     #[test]
@@ -372,4 +473,32 @@ mod tests {
             1
         );
     }
+
+    #[test]
+    fn hybrid_resource() {
+        #[derive(HybridResource)]
+        struct Foo(i32);
+
+        #[derive(Component)]
+        struct Bar;
+
+        let mut world = World::default();
+
+        // global default
+        world.insert_resource(Foo(20));
+
+        // local value
+        world.spawn((Foo(9), Bar));
+
+        // there are two of them
+        assert_eq!(world.query::<&Foo>().iter(&world).count(), 2);
+
+        // the main one
+        assert_eq!(world.resource::<Foo>().0, 20);
+
+        // the local one
+        let mut query = world.query::<(&Foo, &Bar)>();
+        let (foo, _) = query.iter(&world).next().unwrap();
+        assert_eq!(foo.0, 9);
+    }
 }

crates/bevy_ecs/src/world/mod.rs

@@ -57,7 +57,10 @@ use crate::{
     prelude::{Add, Despawn, Discard, Insert, Remove},
     query::{DebugCheckedUnwrap, QueryData, QueryFilter, QueryState},
     relationship::RelationshipHookMode,
-    resource::{IsResource, Resource, ResourceEntities, IS_RESOURCE},
+    resource::{
+        IsHybridResource, IsResource, IsUnique, Resource, ResourceEntities, IS_HYBRID_RESOURCE,
+        IS_RESOURCE, IS_UNIQUE,
+    },
     schedule::{Schedule, ScheduleLabel, Schedules},
     storage::{NonSendData, Storages},
     system::Commands,
@@ -178,6 +181,12 @@ impl World {
         let is_resource = self.register_component::<IsResource>();
         assert_eq!(IS_RESOURCE, is_resource);
 
+        let is_unique = self.register_component::<IsUnique>();
+        assert_eq!(IS_UNIQUE, is_unique);
+
+        let is_hybrid_resource = self.register_component::<IsHybridResource>();
+        assert_eq!(IS_HYBRID_RESOURCE, is_hybrid_resource);
+
         // This sets up `Disabled` as a disabling component, via the FromWorld impl
         self.init_resource::<DefaultQueryFilters>();
     }