Merge #563

563: Optionized get_mut_or_default r=torkleyy a=JaniM

Fixes #561
Alternative to #562 . Implements an Option API instead of panicking if the entity is dead

I have no idea how I managed to mess up the branch name.



Co-authored-by: Jani Mustonen <[email protected]>

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)
     }

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());