Clean up methods, use substruct methods for consistency, entry error …

…is now a single struct.

src/error.rs

@@ -129,7 +129,6 @@ impl StdError for WrongGeneration {
     }
 }
 
-
 /// An error type which cannot be instantiated.
 /// Used as a placeholder for associated error types if
 /// something cannot fail.
@@ -147,3 +146,19 @@ impl StdError for NoError {
         match *self {}
     }
 }
+
+/// Storage entry error when the request entity was dead.
+#[derive(Debug)]
+pub struct EntryIsDead(pub Entity);
+
+impl Display for EntryIsDead {
+    fn fmt(&self, f: &mut Formatter) -> FmtResult {
+        write!(f, "Attempted to get an entry for entity {:?} after it died.", self.0)
+    }
+}
+
+impl StdError for EntryIsDead {
+    fn description(&self) -> &str {
+        "Attempted to get a storage entry for an entity that is no longer valid/alive"
+    }
+}

src/storage/mod.rs

@@ -192,7 +192,7 @@ where
 }
 
 /// An entry to a storage which has a component associated to the entity.
-pub struct OccupiedEntry<'a, 'b: 'a, T: 'b, D: 'b> {
+pub struct OccupiedEntry<'a, 'b: 'a, T: 'a, D: 'a> {
     entity: Entity,
     storage: &'a mut Storage<'b, T, D>,
 }
@@ -212,30 +212,33 @@ impl<'a, 'b, T, D> OccupiedEntry<'a, 'b, T, D>
           D: DerefMut<Target = MaskedStorage<T>>,
 {
     /// Get a mutable reference to the component associated with the entity.
-    pub fn get_mut(&mut self) -> Option<&mut T> {
-        self.storage.get_mut(self.entity)
+    pub fn get_mut(self) -> &'a mut T {
+        match self.storage.get_mut(self.entity) {
+            Some(component) => component,
+            None => { panic!("`OccupiedEntry` on an entity without a component."); },
+        }
     }
 
     /// Inserts a value into the storage and returns the old one.
-    ///
-    /// Return value can be `None` if the entity passed in was dead.
-    pub fn insert(&mut self, component: T) -> Option<T> {
+    pub fn insert(&mut self, component: T) -> T {
         match self.storage.insert(self.entity, component) {
-            InsertResult::Updated(old) => Some(old),
-            _ => None,
+            InsertResult::Updated(old) => old,
+            InsertResult::Inserted => { panic!("`OccupiedEntry` on an entity without a component."); },
+            InsertResult::EntityIsDead(_) => { panic!("`OccupiedEntry` on a dead entity: {:?}.", self.entity); },
         }
     }
 
     /// Removes the component from the storage and returns it.
-    ///
-    /// Return value can be `None` if the entity passed in was dead.
-    pub fn remove(&mut self) -> Option<T> {
-        self.storage.remove(self.entity)
+    pub fn remove(self) -> T {
+        match self.storage.remove(self.entity) {
+            Some(old) => old,
+            None => { panic!("`OccupiedEntry` on an entity without a component."); },
+        }
     }
 }
 
 /// An entry to a storage which does not have a component associated to the entity.
-pub struct VacantEntry<'a, 'b: 'a, T: 'b, D: 'b> {
+pub struct VacantEntry<'a, 'b: 'a, T: 'a, D: 'a> {
     entity: Entity,
     storage: &'a mut Storage<'b, T, D>,
 }
@@ -244,30 +247,30 @@ impl<'a, 'b, T, D> VacantEntry<'a, 'b, T, D>
     where T: Component,
           D: DerefMut<Target = MaskedStorage<T>>,
 {
-    /// Inserts a value into the storage and returns the old one.
-    ///
-    /// Return value can be `None` if the entity passed in was dead.
-    pub fn insert(&mut self, component: T) -> Option<T> {
+    /// Inserts a value into the storage.
+    pub fn insert(self, component: T) -> &'a mut T {
         match self.storage.insert(self.entity, component) {
-            InsertResult::Updated(old) => Some(old),
-            _ => None,
+            InsertResult::Updated(_) => { panic!("`VacantEntry` on an entity with a component."); },
+            InsertResult::Inserted => {
+                match self.storage.get_mut(self.entity) {
+                    Some(component) => component,
+                    None => { panic!("Storage insertion into entity {:?} failed on `VacantEntry`", self.entity); },
+                }
+            },
+            InsertResult::EntityIsDead(_) => { panic!("`VacantEntry` on a dead entity: {:?}.", self.entity); },
         }
     }
 }
 
 /// Entry to a storage for convenient filling of components or removal based on whether
 /// the entity has a component.
-pub enum StorageEntry<'a, 'b: 'a, T: 'b, D: 'b> {
+pub enum StorageEntry<'a, 'b: 'a, T: 'a, D: 'a> {
     /// Entry variant that is returned if the entity does has a component.
     Occupied(OccupiedEntry<'a, 'b, T, D>),
     /// Entry variant that is returned if the entity does not have a component.
     Vacant(VacantEntry<'a, 'b, T, D>),
 }
 
-pub enum EntryError {
-    EntityIsDead(Entity),
-}
-
 impl<'a, 'b, T, D> StorageEntry<'a, 'b, T, D>
     where T: Component,
           D: DerefMut<Target = MaskedStorage<T>>
@@ -278,29 +281,14 @@ impl<'a, 'b, T, D> StorageEntry<'a, 'b, T, D>
         self.or_insert_with(|| component)
     }
 
-    /// Inserts a component using a default function.
+    /// Inserts a component using a lazily called function that is only called
+    /// when inserting the component.
     pub fn or_insert_with<F>(self, component: F) -> &'a mut T
         where F: FnOnce() -> T,
     {
         match self {
-            StorageEntry::Occupied(occupied) => occupied.storage.get_mut(occupied.entity).unwrap(),
-            StorageEntry::Vacant(vacant) => {
-                let insert_result = vacant.storage.insert(vacant.entity, component());
-                if let InsertResult::EntityIsDead(_) = insert_result {
-                    panic!("Entry with an invalid entity");
-                }
-                else {
-                    vacant.storage.get_mut(vacant.entity).unwrap()
-                }
-            }
-        }
-    }
-
-    /// Returns the entity of the current entry.
-    pub fn entity(&self) -> &Entity {
-        match self {
-            &StorageEntry::Occupied(ref occupied) => &occupied.entity,
-            &StorageEntry::Vacant(ref vacant) => &vacant.entity,
+            StorageEntry::Occupied(occupied) => occupied.get_mut(),
+            StorageEntry::Vacant(vacant) => vacant.insert(component()),
         }
     }
 }
@@ -321,8 +309,29 @@ where
 
     /// Returns an entry to the component associated to the entity.
     ///
-    /// Note: This does not immediately error if the entity is dead.
-    pub fn entry<'a>(&'a mut self, e: Entity) -> Result<StorageEntry<'a, 'e, T, D>, EntryError>
+    /// Behaves somewhat similarly to `std::collections::HashMap`'s entry api.
+    ///
+    ///## Example
+    /// 
+    /// ```rust
+    ///# extern crate specs;
+    ///# struct Comp {
+    ///#    field: u32
+    ///# }
+    ///# impl specs::Component for Comp {
+    ///#    type Storage = specs::DenseVecStorage<Self>;
+    ///# }
+    ///# fn main() {
+    ///# let mut world = specs::World::new();
+    ///# world.register::<Comp>();
+    ///# let entity = world.create_entity().build();
+    ///# let mut storage = world.write::<Comp>();
+    /// if let Ok(entry) = storage.entry(entity) {
+    ///     entry.or_insert(Comp { field: 55 });
+    /// }
+    ///# }
+    /// ```
+    pub fn entry<'a>(&'a mut self, e: Entity) -> Result<StorageEntry<'a, 'e, T, D>, ::error::EntryIsDead>
         where 'e: 'a,
     {
         if self.entities.is_alive(e) {
@@ -334,7 +343,7 @@ where
             }
         }
         else {
-            Err(EntryError::EntityIsDead(e))
+            Err(::error::EntryIsDead(e))
         }
     }
 

src/storage/tests.rs

@@ -569,7 +569,12 @@ mod test {
         let e3 = w.create_entity().build();
         let e4 = w.create_entity().with(Cvec(10)).build();
 
+        let e5 = w.create_entity().build();
+        let e6 = w.create_entity().with(Cvec(10)).build();
+
         let mut s1 = w.write::<Cvec>();
+
+        // Basic entry usage.
         if let Ok(entry) = s1.entry(e1) {
             entry.or_insert(Cvec(5));
         }
@@ -578,30 +583,57 @@ mod test {
             entry.or_insert(Cvec(5));
         }
 
-        let mut increment = 0;
-        if let Ok(entry) = s1.entry(e3) {
-            entry.or_insert_with(|| {
-                increment += 1;
-                Cvec(5)
-            });
-
-            assert_eq!(increment, 1);
-        }
-
-        if let Ok(entry) = s1.entry(e4) {
-            entry.or_insert_with(|| {
-                increment += 1;
-                Cvec(5)
-            });
-
-            // Should not have been incremented.
-            assert_eq!(increment, 1);
-        }
-
-        assert_eq!(*s1.get(e1).unwrap(), Cvec(5));
-        assert_eq!(*s1.get(e2).unwrap(), Cvec(10));
-        assert_eq!(*s1.get(e3).unwrap(), Cvec(5));
-        assert_eq!(*s1.get(e4).unwrap(), Cvec(10));
+        // Verify that lazy closures are called only when inserted.
+        {
+            let mut increment = 0;
+            let mut lazy_increment = |entity: Entity, valid: u32| {
+                if let Ok(entry) = s1.entry(entity) {
+                    entry.or_insert_with(|| {
+                        increment += 1;
+                        Cvec(5)
+                    });
+
+                    assert_eq!(increment, valid);
+                }
+            };
+
+            lazy_increment(e3, 1);
+            lazy_increment(e4, 1);
+        }
+
+        // Sanity checks that the entry is occupied after insertions.
+        {
+            let mut occupied = |entity, value| {
+                assert_eq!(*s1.get(entity).unwrap(), value);
+
+                match s1.entry(entity) {
+                    Ok(StorageEntry::Occupied(occupied)) => assert_eq!(*occupied.get_mut(), value),
+                    _ => panic!("Entity not occupied {:?}", entity),
+                }
+            };
+
+            occupied(e1, Cvec(5));
+            occupied(e2, Cvec(10));
+            occupied(e3, Cvec(5));
+            occupied(e4, Cvec(10));
+        }
+
+        // Swap between occupied and vacant depending on the type of entry.
+        {
+            let mut toggle = |entity: Entity| {
+                match s1.entry(entity) {
+                    Ok(StorageEntry::Occupied(occupied)) => { occupied.remove(); },
+                    Ok(StorageEntry::Vacant(vacant)) => { vacant.insert(Cvec(15)); },
+                    Err(_) => { },
+                }
+            };
+
+            toggle(e5);
+            toggle(e6);
+        }
+
+        assert_eq!(s1.get(e5), Some(&Cvec(15)));
+        assert_eq!(s1.get(e6), None);
     }
 
     #[test]