From 9d5b04141cf3d02c8d099db5f3b50dfb1fdab0ad Mon Sep 17 00:00:00 2001
From: Jani Mustonen <janijohannes@kapsi.fi>
Date: Fri, 8 Mar 2019 11:54:18 +0200
Subject: [PATCH] option version of get_mut_or_default

---
 src/storage/generic.rs | 34 +++++++++++++++++++++++++++++++
 src/storage/tests.rs   | 46 +++++++++++++++++++++++++++++++++++++++---
 2 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/src/storage/generic.rs b/src/storage/generic.rs
index 7c38cee61..30172539a 100644
--- a/src/storage/generic.rs
+++ b/src/storage/generic.rs
@@ -83,6 +83,14 @@ pub trait GenericWriteStorage {
     /// Get mutable access to an `Entity`s component
     fn get_mut(&mut self, entity: Entity) -> Option<&mut Self::Component>;
 
+    /// Get mutable access to an `Entity`s component. If the component does not exist, it
+    /// is automatically created using `Default::default()`.
+    ///
+    /// Returns None if the entity is dead.
+    fn get_mut_or_default(&mut self, entity: Entity) -> Option<&mut Self::Component>
+    where
+        Self::Component: Default;
+
     /// Insert a component for an `Entity`
     fn insert(&mut self, entity: Entity, comp: Self::Component) -> InsertResult<Self::Component>;
 
@@ -103,6 +111,19 @@ where
         WriteStorage::get_mut(self, entity)
     }
 
+    fn get_mut_or_default(&mut self, entity: Entity) -> Option<&mut Self::Component>
+    where
+        Self::Component: Default,
+    {
+        if !self.contains(entity) {
+            self.insert(entity, Default::default())
+                .ok()
+                .and_then(move |_| self.get_mut(entity))
+        } else {
+            self.get_mut(entity)
+        }
+    }
+
     fn insert(&mut self, entity: Entity, comp: Self::Component) -> InsertResult<Self::Component> {
         WriteStorage::insert(self, entity, comp)
     }
@@ -126,6 +147,19 @@ where
         WriteStorage::get_mut(*self, entity)
     }
 
+    fn get_mut_or_default(&mut self, entity: Entity) -> Option<&mut Self::Component>
+    where
+        Self::Component: Default,
+    {
+        if !self.contains(entity) {
+            self.insert(entity, Default::default())
+                .ok()
+                .and_then(move |_| self.get_mut(entity))
+        } else {
+            self.get_mut(entity)
+        }
+    }
+
     fn insert(&mut self, entity: Entity, comp: Self::Component) -> InsertResult<Self::Component> {
         WriteStorage::insert(*self, entity, comp)
     }
diff --git a/src/storage/tests.rs b/src/storage/tests.rs
index 6df8db152..11411ea8c 100644
--- a/src/storage/tests.rs
+++ b/src/storage/tests.rs
@@ -136,7 +136,7 @@ mod test {
         type Storage = NullStorage<Self>;
     }
 
-    #[derive(PartialEq, Eq, Debug)]
+    #[derive(PartialEq, Eq, Debug, Default)]
     struct Cvec(u32);
     impl From<u32> for Cvec {
         fn from(v: u32) -> Cvec {
@@ -295,6 +295,43 @@ mod test {
         }
     }
 
+    fn test_get_mut_or_default<T: Component + Default + From<u32> + AsMut<u32> + Debug + Eq>()
+    where
+        T::Storage: Default,
+    {
+        let mut w = World::new();
+        let mut s: Storage<T, _> = create(&mut w);
+
+        // Insert the first 500 components manually, leaving indices 500..1000 unoccupied.
+        for i in 0..500 {
+            if let Err(err) = s.insert(Entity::new(i, Generation::new(1)), (i).into()) {
+                panic!("Failed to insert component into entity! {:?}", err);
+            }
+        }
+
+        for i in 0..1_000 {
+            *s.get_mut_or_default(Entity::new(i, Generation::new(1)))
+                .unwrap()
+                .as_mut() += i;
+        }
+
+        // The first 500 were initialized, and should be i*2.
+        for i in 0..500 {
+            assert_eq!(
+                s.get(Entity::new(i, Generation::new(1))).unwrap(),
+                &(i + i).into()
+            );
+        }
+
+        // The rest were Default-initialized, and should equal i.
+        for i in 500..1_000 {
+            assert_eq!(
+                s.get(Entity::new(i, Generation::new(1))).unwrap(),
+                &(i).into()
+            );
+        }
+    }
+
     fn test_add_gen<T: Component + From<u32> + Debug + Eq>()
     where
         T::Storage: Default,
@@ -391,6 +428,10 @@ mod test {
         test_get_mut::<Cvec>();
     }
     #[test]
+    fn vec_test_get_mut_or_default() {
+        test_get_mut_or_default::<Cvec>();
+    }
+    #[test]
     fn vec_test_add_gen() {
         test_add_gen::<Cvec>();
     }
@@ -754,7 +795,6 @@ mod test {
             comp.0 += 1;
         }
 
-
         {
             inserted.clear();
             modified.clear();
@@ -828,7 +868,7 @@ mod test {
         }
 
         let mut sum = 0;
-        
+
         w.exec(|(e, mut s): (Entities, WriteStorage<CEntries>)| {
             sum = (&e, s.entries()).join().fold(0, |acc, (_, value)| {
                 let v = value.or_insert(2.into());