Merge 72efbc5 into 4558f85

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/lib.rs

@@ -209,8 +209,8 @@ pub use shred::AsyncDispatcher;
 
 pub use storage::{BTreeStorage, DenseVecStorage, DistinctStorage, Entry, FlaggedStorage,
                   HashMapStorage, InsertResult, MaskedStorage, NormalRestriction, NullStorage,
-                  ParallelRestriction, ReadStorage, RestrictedStorage, Storage,
-                  UnprotectedStorage, VecStorage, WriteStorage};
+                  OccupiedEntry, ParallelRestriction, ReadStorage, RestrictedStorage, Storage,
+                  StorageEntry, UnprotectedStorage, VacantEntry, VecStorage, WriteStorage};
 pub use world::{Component, CreateIter, CreateIterAtomic, EntitiesRes, Entity, EntityBuilder,
                 Generation, LazyUpdate, World};
 

src/storage/mod.rs

@@ -191,6 +191,108 @@ where
     }
 }
 
+/// An entry to a storage which has a component associated to the entity.
+pub struct OccupiedEntry<'a, 'b: 'a, T: 'a, D: 'a> {
+    entity: Entity,
+    storage: &'a mut Storage<'b, T, D>,
+}
+
+impl<'a, 'b, T, D> OccupiedEntry<'a, 'b, T, D>
+    where T: Component,
+          D: Deref<Target = MaskedStorage<T>>,
+{
+    /// Get a reference to the component associated with the entity.
+    pub fn get(&self) -> Option<&T> {
+        self.storage.get(self.entity)
+    }
+}
+
+impl<'a, 'b, T, D> OccupiedEntry<'a, 'b, T, D>
+    where T: Component,
+          D: DerefMut<Target = MaskedStorage<T>>,
+{
+    /// Get a mutable reference to the component associated with the 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.
+    pub fn insert(&mut self, component: T) -> T {
+        match self.storage.insert(self.entity, component) {
+            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.
+    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: 'a, D: 'a> {
+    entity: Entity,
+    storage: &'a mut Storage<'b, T, D>,
+}
+
+impl<'a, 'b, T, D> VacantEntry<'a, 'b, T, D>
+    where T: Component,
+          D: DerefMut<Target = MaskedStorage<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(_) => { 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: '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>),
+}
+
+impl<'a, 'b, T, D> StorageEntry<'a, 'b, T, D>
+    where T: Component,
+          D: DerefMut<Target = MaskedStorage<T>>
+{
+
+    /// Inserts a component if the entity does not contain a component.
+    pub fn or_insert(self, component: T) -> &'a mut T {
+        self.or_insert_with(|| component)
+    }
+
+    /// 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.get_mut(),
+            StorageEntry::Vacant(vacant) => vacant.insert(component()),
+        }
+    }
+}
+
 impl<'e, T, D> Storage<'e, T, D>
 where
     T: Component,
@@ -205,6 +307,46 @@ where
         }
     }
 
+    /// Returns an entry to the component associated to the entity.
+    ///
+    /// 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) {
+            if self.data.mask.contains(e.id()) && self.entities.is_alive(e) {
+                Ok(StorageEntry::Occupied(OccupiedEntry { entity: e, storage: self }))
+            }
+            else {
+                Ok(StorageEntry::Vacant(VacantEntry { entity: e, storage: self }))
+            }
+        }
+        else {
+            Err(::error::EntryIsDead(e))
+        }
+    }
+
     /// Inserts new data for a given `Entity`.
     /// Returns the result of the operation as a `InsertResult<T>`
     pub fn insert(&mut self, e: Entity, mut v: T) -> InsertResult<T> {

src/storage/tests.rs

@@ -558,6 +558,84 @@ mod test {
         );
     }
 
+    #[test]
+    fn storage_entry() {
+        let mut w = World::new();
+        w.register::<Cvec>();
+
+        let e1 = w.create_entity().build();
+        let e2 = w.create_entity().with(Cvec(10)).build();
+
+        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));
+        }
+
+        if let Ok(entry) = s1.entry(e2) {
+            entry.or_insert(Cvec(5));
+        }
+
+        // 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]
     #[should_panic]
     fn wrong_storage() {