feat(ecs): add weak reference query

1 files changed, 78 insertions, 6 deletions

diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index df25f46..7a56f37 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -6,6 +6,7 @@ use std::fmt::Debug;
 use std::marker::PhantomData;
 use std::ops::RangeBounds;
 use std::slice::Iter as SliceIter;
+use std::sync::{Arc, Weak};
 use std::vec::Drain;
 
 use crate::actions::Action;
@@ -120,7 +121,7 @@ impl World
     where
         Comps: ComponentSequence,
     {
-        Query::new(&self.data)
+        Query::new(&self.data.component_storage)
     }
 
     /// Peforms the actions that have been queued up using [`Actions`].
@@ -159,7 +160,7 @@ impl World
 pub struct WorldData
 {
     events: HashMap<EventId, Vec<usize>>,
-    component_storage: Lock<ComponentStorage>,
+    component_storage: Arc<Lock<ComponentStorage>>,
     action_queue: Lock<ActionQueue>,
 }
 
@@ -190,12 +191,14 @@ impl TypeName for ActionQueue
     }
 }
 
+/// A entity query.
 #[derive(Debug)]
 pub struct Query<'world, Comps>
 where
     Comps: ComponentSequence,
 {
     component_storage: ReadGuard<'world, ComponentStorage>,
+    component_storage_lock: Weak<Lock<ComponentStorage>>,
     comps_pd: PhantomData<Comps>,
 }
 
@@ -214,13 +217,22 @@ where
         }
     }
 
-    fn new(world_data: &'world WorldData) -> Self
+    /// Returns a weak reference query to the same components.
+    pub fn to_weak_ref(&self) -> WeakRefQuery<Comps>
+    {
+        WeakRefQuery {
+            component_storage: self.component_storage_lock.clone(),
+            comps_pd: PhantomData,
+        }
+    }
+
+    fn new(component_storage: &'world Arc<Lock<ComponentStorage>>) -> Self
     {
         Self {
-            component_storage: world_data
-                .component_storage
+            component_storage: component_storage
                 .read_nonblock()
                 .expect("Failed to acquire read-only component storage lock"),
+            component_storage_lock: Arc::downgrade(component_storage),
             comps_pd: PhantomData,
         }
     }
@@ -258,7 +270,7 @@ where
         world_data: &'world WorldData,
     ) -> Self
     {
-        Self::new(world_data)
+        Self::new(&world_data.component_storage)
     }
 
     fn is_compatible<Other: SystemParam<'world>>() -> bool
@@ -282,6 +294,66 @@ where
     }
 }
 
+/// A entity query containing a weak reference to the world.
+#[derive(Debug)]
+pub struct WeakRefQuery<Comps>
+where
+    Comps: ComponentSequence,
+{
+    component_storage: Weak<Lock<ComponentStorage>>,
+    comps_pd: PhantomData<Comps>,
+}
+
+impl<Comps> WeakRefQuery<Comps>
+where
+    Comps: ComponentSequence,
+{
+    /// Returns a struct which can be used to retrieve a [`Query`].
+    ///
+    /// Returns [`None`] if the [`World`] has been dropped.
+    pub fn access(&self) -> Option<RefQuery<'_, Comps>>
+    {
+        Some(RefQuery {
+            component_storage: self.component_storage.upgrade()?,
+            _pd: PhantomData,
+        })
+    }
+}
+
+impl<Comps> Clone for WeakRefQuery<Comps>
+where
+    Comps: ComponentSequence,
+{
+    fn clone(&self) -> Self
+    {
+        Self {
+            component_storage: self.component_storage.clone(),
+            comps_pd: PhantomData,
+        }
+    }
+}
+
+/// Intermediate between [`Query`] and [`WeakRefQuery`]. Contains a strong reference to
+/// the world which is not allowed direct access to.
+#[derive(Debug, Clone)]
+pub struct RefQuery<'weak_ref, Comps>
+where
+    Comps: ComponentSequence,
+{
+    component_storage: Arc<Lock<ComponentStorage>>,
+    _pd: PhantomData<&'weak_ref Comps>,
+}
+
+impl<'weak_ref, Comps> RefQuery<'weak_ref, Comps>
+where
+    Comps: ComponentSequence,
+{
+    pub fn to_query(&self) -> Query<'_, Comps>
+    {
+        Query::new(&self.component_storage)
+    }
+}
+
 #[derive(Debug)]
 struct QueryComponentIds
 {