From e028d89e15fd219854cbd7122d63ebd2fa37d390 Mon Sep 17 00:00:00 2001
From: HampusM <hampus@hampusmat.com>
Date: Wed, 10 Apr 2024 22:45:35 +0200
Subject: feat(ecs): add weak ref single

---
 ecs/src/lib.rs  |  9 ++++--
 ecs/src/sole.rs | 94 +++++++++++++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 90 insertions(+), 13 deletions(-)

diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index db82d91..8a5de64 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -353,12 +353,12 @@ pub struct SoleAlreadyExistsError(pub &'static str);
 #[derive(Debug, Default)]
 struct SoleStorage
 {
-    storage: HashMap<TypeId, Lock<Box<dyn Sole>>>,
+    storage: HashMap<TypeId, Arc<Lock<Box<dyn Sole>>>>,
 }
 
 impl SoleStorage
 {
-    fn get<SoleT: Sole>(&self) -> Option<&Lock<Box<dyn Sole>>>
+    fn get<SoleT: Sole>(&self) -> Option<&Arc<Lock<Box<dyn Sole>>>>
     {
         self.storage.get(&TypeId::of::<SoleT>())
     }
@@ -371,7 +371,10 @@ impl SoleStorage
             return Err(SoleAlreadyExistsError(type_name::<SoleT>()));
         }
 
-        self.storage.insert(sole_type_id, Lock::new(Box::new(sole)));
+        // TODO: Reconsider this maybe?
+        #[allow(clippy::arc_with_non_send_sync)]
+        self.storage
+            .insert(sole_type_id, Arc::new(Lock::new(Box::new(sole))));
 
         Ok(())
     }
diff --git a/ecs/src/sole.rs b/ecs/src/sole.rs
index e5166e9..014ab7d 100644
--- a/ecs/src/sole.rs
+++ b/ecs/src/sole.rs
@@ -2,8 +2,9 @@ use std::any::{type_name, Any, TypeId};
 use std::fmt::Debug;
 use std::marker::PhantomData;
 use std::ops::{Deref, DerefMut};
+use std::sync::{Arc, Weak};
 
-use crate::lock::WriteGuard;
+use crate::lock::{Lock, WriteGuard};
 use crate::system::{NoInitParamFlag, Param as SystemParam, System};
 use crate::tuple::FilterExclude as TupleFilterExclude;
 use crate::type_name::TypeName;
@@ -51,9 +52,41 @@ impl TypeName for Box<dyn Sole>
 pub struct Single<'world, SoleT: Sole>
 {
     sole: WriteGuard<'world, Box<dyn Sole>>,
+    sole_weak: Weak<Lock<Box<dyn Sole>>>,
     _ph: PhantomData<SoleT>,
 }
 
+impl<'world, SoleT> Single<'world, SoleT>
+where
+    SoleT: Sole,
+{
+    /// Returns a struct which holds a weak reference to the [`World`] that `Single`
+    /// references and that can be used to aquire a new `Single` if the referenced
+    /// [`World`] is still alive.
+    #[must_use]
+    pub fn to_weak_ref(&self) -> SingleWeakRef<SoleT>
+    {
+        SingleWeakRef {
+            sole: self.sole_weak.clone(),
+            _ph: PhantomData,
+        }
+    }
+
+    fn new(sole: &'world Arc<Lock<Box<dyn Sole>>>) -> Self
+    {
+        Self {
+            sole: sole.write_nonblock().unwrap_or_else(|_| {
+                panic!(
+                    "Failed to aquire read-write lock to single component {}",
+                    type_name::<SoleT>()
+                )
+            }),
+            sole_weak: Arc::downgrade(sole),
+            _ph: PhantomData,
+        }
+    }
+}
+
 unsafe impl<'world, SoleT> SystemParam<'world> for Single<'world, SoleT>
 where
     SoleT: Sole,
@@ -76,16 +109,9 @@ where
         let sole = world_data
             .sole_storage
             .get::<SoleT>()
-            .expect("Sole was not found in world")
-            .write_nonblock()
-            .unwrap_or_else(|_| {
-                panic!(
-                    "Failed to aquire read-write lock to single component {}",
-                    type_name::<SoleT>()
-                )
-            });
+            .expect("Sole was not found in world");
 
-        Self { sole, _ph: PhantomData }
+        Self::new(sole)
     }
 
     fn is_compatible<Other: SystemParam<'world>>() -> bool
@@ -128,6 +154,54 @@ where
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct SingleWeakRef<SoleT>
+where
+    SoleT: Sole,
+{
+    sole: Weak<Lock<Box<dyn Sole>>>,
+    _ph: PhantomData<SoleT>,
+}
+
+impl<SoleT> SingleWeakRef<SoleT>
+where
+    SoleT: Sole,
+{
+    /// Returns a struct which can be used to retrieve a [`Single`].
+    ///
+    /// Returns [`None`] if the referenced [`World`] has been dropped.
+    #[must_use]
+    pub fn access(&self) -> Option<SingleRef<'_, SoleT>>
+    {
+        Some(SingleRef {
+            sole: self.sole.upgrade()?,
+            _pd: PhantomData,
+        })
+    }
+}
+
+/// Intermediate between [`Single`] and [`SingleWeakRef`]. Contains a strong reference to
+/// a world which is not allowed direct access to.
+#[derive(Debug, Clone)]
+pub struct SingleRef<'weak_ref, SoleT>
+where
+    SoleT: Sole,
+{
+    sole: Arc<Lock<Box<dyn Sole>>>,
+    _pd: PhantomData<&'weak_ref SoleT>,
+}
+
+impl<'weak_ref, SoleT> SingleRef<'weak_ref, SoleT>
+where
+    SoleT: Sole,
+{
+    #[must_use]
+    pub fn to_single(&self) -> Single<'_, SoleT>
+    {
+        Single::new(&self.sole)
+    }
+}
+
 struct Comparable
 {
     sole_type_id: TypeId,
-- 
cgit v1.2.3-18-g5258