1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
use hibitset::BitSetAll;

use super::*;
use crate::join::Join;

impl<'e, T, D> Storage<'e, T, D>
where
    T: Component,
    D: DerefMut<Target = MaskedStorage<T>>,
{
    /// 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;
    /// # use specs::prelude::*;
    /// # struct Comp {
    /// #    field: u32
    /// # }
    /// # impl Component for Comp {
    /// #    type Storage = DenseVecStorage<Self>;
    /// # }
    /// # fn main() {
    /// # let mut world = World::new();
    /// # world.register::<Comp>();
    /// # let entity = world.create_entity().build();
    /// # let mut storage = world.write_storage::<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>, WrongGeneration>
    where
        'e: 'a,
    {
        if self.entities.is_alive(e) {
            unsafe {
                let entries = self.entries();
                // SAFETY: This is safe since we're not swapping out the mask or the values.
                let (_, mut value): (BitSetAll, _) = entries.open();
                // SAFETY: We did check the mask, because the mask is `BitSetAll` and every
                // index is part of it.
                Ok(Entries::get(&mut value, e.id()))
            }
        } else {
            let gen = self
                .entities
                .alloc
                .generation(e.id())
                .unwrap_or_else(Generation::one);
            Err(WrongGeneration {
                action: "attempting to get an entry to a storage",
                actual_gen: gen,
                entity: e,
            })
        }
    }

    /// Returns a `Join`-able structure that yields all indices, returning
    /// `Entry` for all elements
    ///
    /// WARNING: Do not have a join of only `Entries`s. Otherwise the join will
    /// iterate over every single index of the bitset. If you want a join with
    /// all `Entries`s, add an `EntitiesRes` to the join as well to bound the
    /// join to all entities that are alive.
    ///
    /// ## Example
    ///
    /// ```rust
    /// # extern crate specs;
    /// # use specs::prelude::*;
    /// #
    /// # #[derive(Default)]
    /// # struct Counter(u32);
    /// #
    /// # impl Counter {
    /// #     fn increase(&mut self) {
    /// #         self.0 += 1
    /// #     }
    /// #     fn reached_limit(&self) -> bool {
    /// #         return self.0 >= 100;
    /// #     }
    /// #     fn reset(&mut self) {
    /// #         return self.0 = 0;
    /// #     }
    /// # }
    /// #
    /// # impl Component for Counter {
    /// #     type Storage = VecStorage<Self>;
    /// # }
    /// #
    /// # #[derive(Default)]
    /// # struct AllowCounter;
    /// #
    /// # impl Component for AllowCounter {
    /// #     type Storage = NullStorage<Self>;
    /// # }
    /// #
    /// # let mut world = World::new();
    /// # world.register::<Counter>();
    /// # for _ in 0..15 {
    /// #     world.create_entity().build();
    /// # }
    /// #
    /// # world.exec(|(mut counters, marker): (WriteStorage<Counter>, ReadStorage<AllowCounter>)| {
    /// for (mut counter, _) in (counters.entries(), &marker).join() {
    ///     let counter = counter.or_insert_with(Default::default);
    ///     counter.increase();
    ///
    ///     if counter.reached_limit() {
    ///         counter.reset();
    ///         // Do something
    ///     }
    /// }
    /// # });
    /// ```
    pub fn entries<'a>(&'a mut self) -> Entries<'a, 'e, T, D> {
        Entries(self)
    }
}

/// `Join`-able structure that yields all indices, returning `Entry` for all
/// elements
pub struct Entries<'a, 'b: 'a, T: 'a, D: 'a>(&'a mut Storage<'b, T, D>);

impl<'a, 'b: 'a, T: 'a, D: 'a> Join for Entries<'a, 'b, T, D>
where
    T: Component,
    D: Deref<Target = MaskedStorage<T>>,
{
    type Mask = BitSetAll;
    type Type = StorageEntry<'a, 'b, T, D>;
    type Value = &'a mut Storage<'b, T, D>;

    // SAFETY: No invariants to meet and no unsafe code.
    unsafe fn open(self) -> (Self::Mask, Self::Value) {
        (BitSetAll, self.0)
    }

    // SAFETY: We are lengthening the lifetime of `value` to `'a`;
    // TODO: how to prove this is safe?
    unsafe fn get(value: &mut Self::Value, id: Index) -> Self::Type {
        // This is HACK. See implementation of Join for &'a mut Storage<'e, T, D> for
        // details why it is necessary.
        let storage: *mut Storage<'b, T, D> = *value as *mut Storage<'b, T, D>;
        if (*storage).data.mask.contains(id) {
            StorageEntry::Occupied(OccupiedEntry {
                id,
                storage: &mut *storage,
            })
        } else {
            StorageEntry::Vacant(VacantEntry {
                id,
                storage: &mut *storage,
            })
        }
    }

    #[inline]
    fn is_unconstrained() -> bool {
        true
    }
}

/// An entry to a storage which has a component associated to the entity.
pub struct OccupiedEntry<'a, 'b: 'a, T: 'a, D: 'a> {
    id: Index,
    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) -> &T {
        // SAFETY: This is safe since `OccupiedEntry` is only constructed
        // after checking the mask.
        unsafe { self.storage.data.inner.get(self.id) }
    }
}

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(&mut self) -> &mut T {
        // SAFETY: This is safe since `OccupiedEntry` is only constructed
        // after checking the mask.
        unsafe { self.storage.data.inner.get_mut(self.id) }
    }

    /// Converts the `OccupiedEntry` into a mutable reference bounded by
    /// the storage's lifetime.
    pub fn into_mut(self) -> &'a mut T {
        // SAFETY: This is safe since `OccupiedEntry` is only constructed
        // after checking the mask.
        unsafe { self.storage.data.inner.get_mut(self.id) }
    }

    /// Inserts a value into the storage and returns the old one.
    pub fn insert(&mut self, mut component: T) -> T {
        std::mem::swap(&mut component, self.get_mut());
        component
    }

    /// Removes the component from the storage and returns it.
    pub fn remove(self) -> T {
        self.storage.data.remove(self.id).unwrap()
    }
}

/// 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> {
    id: Index,
    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 {
        self.storage.data.mask.add(self.id);
        // SAFETY: This is safe since we added `self.id` to the mask.
        unsafe {
            self.storage.data.inner.insert(self.id, component);
            self.storage.data.inner.get_mut(self.id)
        }
    }
}

/// 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 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 and returns the old value in case this entry was
    /// already occupied.
    pub fn replace(self, component: T) -> Option<T> {
        match self {
            StorageEntry::Occupied(mut occupied) => Some(occupied.insert(component)),
            StorageEntry::Vacant(vacant) => {
                vacant.insert(component);
                None
            }
        }
    }

    /// Inserts a component if the entity does not have it already.
    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. Ensures this entry has a value and if not,
    /// inserts one using the result of the passed closure. Returns a reference
    /// to the value afterwards.
    pub fn or_insert_with<F>(self, default: F) -> &'a mut T
    where
        F: FnOnce() -> T,
    {
        match self {
            StorageEntry::Occupied(occupied) => occupied.into_mut(),
            StorageEntry::Vacant(vacant) => vacant.insert(default()),
        }
    }
}