2 files changed, 65 insertions, 4 deletions

diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index 8a5de64..285f3b3 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -350,17 +350,26 @@ impl Drop for ComponentStorage
 #[error("Sole {0} already exists")]
 pub struct SoleAlreadyExistsError(pub &'static str);
 
+#[derive(Debug)]
+struct StoredSole
+{
+    sole: Arc<Lock<Box<dyn Sole>>>,
+    drop_last: bool,
+}
+
 #[derive(Debug, Default)]
 struct SoleStorage
 {
-    storage: HashMap<TypeId, Arc<Lock<Box<dyn Sole>>>>,
+    storage: HashMap<TypeId, ManuallyDrop<StoredSole>>,
 }
 
 impl SoleStorage
 {
     fn get<SoleT: Sole>(&self) -> Option<&Arc<Lock<Box<dyn Sole>>>>
     {
-        self.storage.get(&TypeId::of::<SoleT>())
+        self.storage
+            .get(&TypeId::of::<SoleT>())
+            .map(|sole| &sole.sole)
     }
 
     fn insert<SoleT: Sole>(&mut self, sole: SoleT) -> Result<(), SoleAlreadyExistsError>
@@ -371,11 +380,61 @@ impl SoleStorage
             return Err(SoleAlreadyExistsError(type_name::<SoleT>()));
         }
 
+        let drop_last = sole.drop_last();
+
         // TODO: Reconsider this maybe?
         #[allow(clippy::arc_with_non_send_sync)]
-        self.storage
-            .insert(sole_type_id, Arc::new(Lock::new(Box::new(sole))));
+        self.storage.insert(
+            sole_type_id,
+            ManuallyDrop::new(StoredSole {
+                sole: Arc::new(Lock::new(Box::new(sole))),
+                drop_last,
+            }),
+        );
 
         Ok(())
     }
 }
+
+impl Drop for SoleStorage
+{
+    fn drop(&mut self)
+    {
+        let mut soles_to_drop_last = Vec::new();
+
+        for sole in self.storage.values_mut() {
+            if sole.drop_last {
+                #[cfg(feature = "debug")]
+                tracing::debug!(
+                    "Sole {} pushed to dropping last queue",
+                    sole.sole.read_nonblock().unwrap().type_name()
+                );
+
+                soles_to_drop_last.push(sole);
+                continue;
+            }
+
+            #[cfg(feature = "debug")]
+            tracing::debug!(
+                "Dropping sole {}",
+                sole.sole.read_nonblock().unwrap().type_name()
+            );
+
+            unsafe {
+                ManuallyDrop::drop(sole);
+            }
+        }
+
+        for sole in &mut soles_to_drop_last {
+            #[cfg(feature = "debug")]
+            tracing::debug!(
+                "Dropping sole {} last",
+                sole.sole.read_nonblock().unwrap().type_name()
+            );
+
+            unsafe {
+                ManuallyDrop::drop(sole);
+            }
+        }
+    }
+}
diff --git a/ecs/src/sole.rs b/ecs/src/sole.rs
index 014ab7d..fcff8de 100644
--- a/ecs/src/sole.rs
+++ b/ecs/src/sole.rs
@@ -13,6 +13,8 @@ use crate::WorldData;
 /// A type which has a single instance and is shared globally.
 pub trait Sole: Any + TypeName
 {
+    fn drop_last(&self) -> bool;
+
     fn as_any_mut(&mut self) -> &mut dyn Any;
 
     fn as_any(&self) -> &dyn Any;