Queries as Entities

crates/bevy_ecs/src/archetype.rs

@@ -19,11 +19,13 @@
 //! [`Table`]: crate::storage::Table
 //! [`World::archetypes`]: crate::world::World::archetypes
 
+use crate as bevy_ecs;
 use crate::{
     bundle::BundleId,
     component::{ComponentId, Components, StorageType},
     entity::{Entity, EntityLocation},
     observer::Observers,
+    prelude::Event,
     storage::{ImmutableSparseSet, SparseArray, SparseSet, SparseSetIndex, TableId, TableRow},
 };
 use std::{
@@ -108,6 +110,9 @@ impl ArchetypeId {
     }
 }
 
+#[derive(Event)]
+pub(crate) struct ArchetypeCreated(pub ArchetypeId);
+
 #[derive(Copy, Clone, Eq, PartialEq)]
 pub(crate) enum ComponentStatus {
     Added,

crates/bevy_ecs/src/query/state.rs

@@ -1,15 +1,17 @@
 use crate::{
-    archetype::{Archetype, ArchetypeComponentId, ArchetypeGeneration, ArchetypeId},
+    archetype::{Archetype, ArchetypeComponentId, ArchetypeCreated, ArchetypeGeneration, ArchetypeId},
     batching::BatchingStrategy,
-    component::{ComponentId, Components, Tick},
+    component::{Component, ComponentHooks, ComponentId, Components, StorageType, Tick},
     entity::Entity,
     prelude::FromWorld,
     query::{
         Access, DebugCheckedUnwrap, FilteredAccess, QueryCombinationIter, QueryIter, QueryParIter,
     },
+    observer::{Observer, Trigger},
     storage::{SparseSetIndex, TableId},
     world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId},
 };
+use crate::world::DeferredWorld;
 use bevy_utils::tracing::warn;
 #[cfg(feature = "trace")]
 use bevy_utils::tracing::Span;
@@ -56,6 +58,7 @@ pub(super) union StorageId {
 // SAFETY NOTE:
 // Do not add any new fields that use the `D` or `F` generic parameters as this may
 // make `QueryState::as_transmuted_state` unsound if not done with care.
+#[derive(Clone)]
 pub struct QueryState<D: QueryData, F: QueryFilter = ()> {
     world_id: WorldId,
     pub(crate) archetype_generation: ArchetypeGeneration,
@@ -74,6 +77,27 @@ pub struct QueryState<D: QueryData, F: QueryFilter = ()> {
     par_iter_span: Span,
 }
 
+impl<D: QueryData + 'static, F: QueryFilter + 'static> Component for QueryState<D, F> {
+    const STORAGE_TYPE: StorageType = StorageType::SparseSet;
+
+    fn register_component_hooks(hooks: &mut ComponentHooks) {
+        hooks.on_add(|mut world, entity, _| {
+            let mut observer = Observer::new(|trigger: Trigger<ArchetypeCreated>,
+                                              mut world: DeferredWorld| {
+                let archetype_id: ArchetypeId = trigger.event().0;
+                let archetype_ptr: *const Archetype = world.archetypes().get(archetype_id).unwrap();
+                let mut entity = world.entity_mut(trigger.entity());
+                let mut state = entity.get_mut::<QueryState<D, F>>().unwrap();
+                unsafe {
+                    state.new_archetype_internal(&*archetype_ptr);
+                }
+            });
+            observer.watch_entity(entity);
+            world.commands().spawn(observer);
+        });
+    }
+}
+
 impl<D: QueryData, F: QueryFilter> fmt::Debug for QueryState<D, F> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("QueryState")