From 0bccf5c44724b620a64b792c9868b8bb9b2b9550 Mon Sep 17 00:00:00 2001
From: hodasemi <michaelh.95@t-online.de>
Date: Sat, 5 Apr 2025 09:43:07 +0200
Subject: [PATCH] Implement optional component return

---
 ecs/src/entity.rs    | 11 +++++++++++
 ecs/src/resources.rs |  7 +++----
 ecs/src/type_map.rs  | 33 +++++++++++++++++++++++++++++++--
 3 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/ecs/src/entity.rs b/ecs/src/entity.rs
index 8144129..c007ca8 100644
--- a/ecs/src/entity.rs
+++ b/ecs/src/entity.rs
@@ -255,6 +255,17 @@ macro_rules! impl_get_components_mut {
                 crate::get_disjoint_mut::GetDisjointMut::<'a, ($(&'a mut $t,)+)>::get_mut(&mut self.components)
             }
         }
+
+        impl<'a, $($t,)+> EntityComponentDisjointMut<'a, ($(Option<&'a mut $t>,)+)> for EntityObject
+        where
+            $(
+                $t: EntityComponent + ComponentDebug + 'static,
+            )+
+        {
+            fn get_components_mut(&'a mut self) -> std::result::Result<($(Option<&'a mut $t>,)+), ComponentNotFoundError> {
+                crate::get_disjoint_mut::GetDisjointMut::<'a, ($(Option<&'a mut $t>,)+)>::get_mut(&mut self.components)
+            }
+        }
     };
 }
 
diff --git a/ecs/src/resources.rs b/ecs/src/resources.rs
index 79b267c..1dc7b9c 100644
--- a/ecs/src/resources.rs
+++ b/ecs/src/resources.rs
@@ -121,9 +121,7 @@ where
     type Error = Error;
 
     fn get_mut(&'a mut self) -> std::result::Result<&'a mut T, Self::Error> {
-        self.map
-            .get_mut(&TypeId::of::<T>())
-            .map(|any| any.downcast_mut().unwrap())
+        self.get_mut_opt()
             .ok_or_else(|| anyhow!("failed downcasting {}", stringify!(T)))
     }
 }
@@ -146,7 +144,8 @@ macro_rules! impl_get_disjoint_mut {
                     types
                         .pop_front()
                         .flatten()
-                        .map(|any| any.downcast_mut().unwrap())
+                        .map(|any| any.downcast_mut())
+                        .flatten()
                         .ok_or_else(|| anyhow!("failed downcasting {}", stringify!($t)))?,
                 )+))
             }
diff --git a/ecs/src/type_map.rs b/ecs/src/type_map.rs
index 7bb9943..5dc54c4 100644
--- a/ecs/src/type_map.rs
+++ b/ecs/src/type_map.rs
@@ -160,7 +160,8 @@ where
     fn get_mut(&'a mut self) -> std::result::Result<&'a mut T, Self::Error> {
         self.map
             .get_mut(&TypeId::of::<T>())
-            .map(|any| any.downcast_mut().unwrap())
+            .map(|any| any.downcast_mut())
+            .flatten()
             .ok_or_else(ComponentNotFoundError::component::<T>)
     }
 }
@@ -183,11 +184,39 @@ macro_rules! impl_get_disjoint_mut {
                     types
                         .pop_front()
                         .flatten()
-                        .map(|any| any.downcast_mut().unwrap())
+                        .map(|any| any.downcast_mut())
+                        .flatten()
                         .ok_or_else(ComponentNotFoundError::component::<$t>)?,
                 )+))
             }
         }
+
+        impl<'a, $($t,)+> GetDisjointMut<'a, ($(Option<&'a mut $t>,)+)> for $struct
+        where
+            $(
+                $t: EntityComponent + ComponentDebug + 'static,
+            )+
+        {
+            type Error = $error;
+
+            fn get_mut(&'a mut self) -> std::result::Result<($(Option<&'a mut $t>,)+), Self::Error> {
+                let mut types: std::collections::VecDeque<_>
+                    = self.map.get_disjoint_mut([$(&TypeId::of::<$t>(),)+]).into_iter().collect();
+
+                Ok(($(
+                    {
+                        // dummy needed for macro iteration
+                        let _ = stringify!($t);
+
+                        types
+                            .pop_front()
+                            .flatten()
+                            .map(|any| any.downcast_mut())
+                            .flatten()
+                    },
+                )+))
+            }
+        }
     };
 }