bones_ecs/components/
typed.rs

1use std::{marker::PhantomData, rc::Rc};
2
3use crate::prelude::*;
4
5use super::untyped::UntypedComponentStore;
6
7/// A typed wrapper around [`UntypedComponentStore`].
8#[derive(Clone)]
9#[repr(transparent)]
10pub struct ComponentStore<T: HasSchema> {
11    untyped: UntypedComponentStore,
12    _phantom: PhantomData<T>,
13}
14
15impl<T: HasSchema> Default for ComponentStore<T> {
16    fn default() -> Self {
17        Self {
18            untyped: UntypedComponentStore::for_type::<T>(),
19            _phantom: PhantomData,
20        }
21    }
22}
23
24impl<T: HasSchema> TryFrom<UntypedComponentStore> for ComponentStore<T> {
25    type Error = SchemaMismatchError;
26
27    fn try_from(untyped: UntypedComponentStore) -> Result<Self, Self::Error> {
28        if untyped.schema == T::schema() {
29            Ok(Self {
30                untyped,
31                _phantom: PhantomData,
32            })
33        } else {
34            Err(SchemaMismatchError)
35        }
36    }
37}
38
39impl<T: HasSchema> ComponentStore<T> {
40    /// Converts to the internal, untyped [`ComponentStore`].
41    #[inline]
42    pub fn into_untyped(self) -> UntypedComponentStore {
43        self.untyped
44    }
45
46    /// Creates a [`ComponentStore`] from an [`UntypedComponentStore`].
47    /// # Panics
48    /// Panics if the schema doesn't match `T`.
49    #[track_caller]
50    pub fn from_untyped(untyped: UntypedComponentStore) -> Self {
51        untyped.try_into().unwrap()
52    }
53
54    // TODO: Replace ComponentStore functions with non-validating ones.
55    // Right now functions like `insert`, `get`, and `get_mut` use the checked and panicing versions
56    // of the `untyped` functions. These functions do an extra check to see that the schema matches,
57    // but we've already validated that in the construction of the `ComponentStore`, so we should
58    // bypass the extra schema check for performance.
59
60    /// Inserts a component for the given `Entity` index.
61    /// Returns the previous component, if any.
62    #[inline]
63    pub fn insert(&mut self, entity: Entity, component: T) -> Option<T> {
64        self.untyped.insert(entity, component)
65    }
66
67    /// Gets an immutable reference to the component of `Entity`.
68    #[inline]
69    pub fn get(&self, entity: Entity) -> Option<&T> {
70        self.untyped.get(entity)
71    }
72
73    /// Gets a mutable reference to the component of `Entity`.
74    #[inline]
75    pub fn get_mut(&mut self, entity: Entity) -> Option<&mut T> {
76        self.untyped.get_mut(entity)
77    }
78
79    /// Get a mutable reference to component if it exists.
80    /// Otherwise inserts `T` generated by calling parameter: `f`.
81    #[inline]
82    pub fn get_mut_or_insert(&mut self, entity: Entity, f: impl FnOnce() -> T) -> &mut T {
83        self.untyped.get_mut_or_insert(entity, f)
84    }
85
86    /// Get mutable references to the component data for multiple entities at the same time.
87    ///
88    /// # Panics
89    ///
90    /// This will panic if the same entity is specified multiple times. This is invalid because it
91    /// would mean you would have two mutable references to the same component data at the same
92    /// time.
93    #[track_caller]
94    pub fn get_many_mut<const N: usize>(&mut self, entities: [Entity; N]) -> [Option<&mut T>; N] {
95        let mut result = self.untyped.get_many_ref_mut(entities);
96
97        std::array::from_fn(move |i| {
98            // SOUND: we know that the schema matches.
99            result[i]
100                .take()
101                .map(|x| unsafe { x.cast_into_mut_unchecked() })
102        })
103    }
104
105    /// Removes the component of `Entity`.
106    /// Returns `Some(T)` if the entity did have the component.
107    /// Returns `None` if the entity did not have the component.
108    #[inline]
109    pub fn remove(&mut self, entity: Entity) -> Option<T> {
110        self.untyped.remove(entity)
111    }
112
113    /// Gets an immutable reference to the component if there is exactly one instance of it.
114    #[inline]
115    pub fn get_single_with_bitset(&self, bitset: Rc<BitSetVec>) -> Result<&T, QuerySingleError> {
116        // SOUND: we know the schema matches.
117        self.untyped
118            .get_single_with_bitset(bitset)
119            .map(|x| unsafe { x.cast_into_unchecked() })
120    }
121
122    /// Gets a mutable reference to the component if there is exactly one instance of it.
123    #[inline]
124    pub fn get_single_with_bitset_mut(
125        &mut self,
126        bitset: Rc<BitSetVec>,
127    ) -> Result<&mut T, QuerySingleError> {
128        // SOUND: we know the schema matches.
129        self.untyped
130            .get_single_with_bitset_mut(bitset)
131            .map(|x| unsafe { x.cast_into_mut_unchecked() })
132    }
133
134    /// Iterates immutably over all components of this type.
135    /// Very fast but doesn't allow joining with other component types.
136    #[inline]
137    pub fn iter(&self) -> impl Iterator<Item = &T> {
138        // SOUND: we know the schema matches.
139        self.untyped
140            .iter()
141            .map(|x| unsafe { x.cast_into_unchecked() })
142    }
143
144    /// Iterates mutably over all components of this type.
145    /// Very fast but doesn't allow joining with other component types.
146    #[inline]
147    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
148        // SOUND: we know the schema matches.
149        self.untyped
150            .iter_mut()
151            .map(|x| unsafe { x.cast_into_mut_unchecked() })
152    }
153}
154
155/// This trait factors out functions for iterating with bitset over component store.
156/// Separated from `impl ComponentStore` for usage in generic trait types that must
157/// be able to create [`ComponentBitsetIterator`] and related types.
158///
159/// Automatically implemented for [`ComponentStore`].
160pub trait ComponentIterBitset<'a, T: HasSchema> {
161    /// Gets an immutable reference to the component if there is exactly one instance of it.
162    fn get_single_with_bitset(&self, bitset: Rc<BitSetVec>) -> Result<&T, QuerySingleError>;
163
164    /// Gets a mutable reference to the component if there is exactly one instance of it.
165    fn get_single_mut_with_bitset(
166        &mut self,
167        bitset: Rc<BitSetVec>,
168    ) -> Result<&mut T, QuerySingleError>;
169
170    /// Iterates immutably over the components of this type where `bitset`
171    /// indicates the indices of entities.
172    /// Slower than `iter()` but allows joining between multiple component types.
173    fn iter_with_bitset(&self, bitset: Rc<BitSetVec>) -> ComponentBitsetIterator<'_, T>;
174
175    /// Iterates immutably over the components of this type where `bitset`
176    /// indicates the indices of entities.
177    /// Slower than `iter()` but allows joining between multiple component types.
178    fn iter_with_bitset_optional(
179        &self,
180        bitset: Rc<BitSetVec>,
181    ) -> ComponentBitsetOptionalIterator<'_, T>;
182
183    /// Iterates mutable over the components of this type where `bitset`
184    /// indicates the indices of entities.
185    /// Slower than `iter()` but allows joining between multiple component types.
186    fn iter_mut_with_bitset(&mut self, bitset: Rc<BitSetVec>) -> ComponentBitsetIteratorMut<'_, T>;
187
188    /// Iterates mutably over the components of this type where `bitset`
189    /// indicates the indices of entities.
190    /// Slower than `iter()` but allows joining between multiple component types.
191    fn iter_mut_with_bitset_optional(
192        &mut self,
193        bitset: Rc<BitSetVec>,
194    ) -> ComponentBitsetOptionalIteratorMut<'_, T>;
195
196    /// Get bitset of [`ComponentStore`] / implementor.
197    fn bitset(&self) -> &BitSetVec;
198
199    /// Check whether or not this component store has data for the given entity.
200    fn contains(&self, entity: Entity) -> bool;
201
202    /// Get [`ComponentStore`] for usage with generic types implementing [`ComponentIterBitset`].
203    fn component_store(&self) -> &ComponentStore<T>;
204}
205
206impl<'a, T: HasSchema> ComponentIterBitset<'a, T> for ComponentStore<T> {
207    /// Gets an immutable reference to the component if there is exactly one instance of it.
208    fn get_single_with_bitset(&self, bitset: Rc<BitSetVec>) -> Result<&T, QuerySingleError> {
209        // SOUND: we know the schema matches.
210        fn map<T>(r: SchemaRef<'_>) -> &T {
211            unsafe { r.cast_into_unchecked() }
212        }
213        self.untyped.get_single_with_bitset(bitset).map(map)
214    }
215
216    /// Gets a mutable reference to the component if there is exactly one instance of it.
217    fn get_single_mut_with_bitset(
218        &mut self,
219        bitset: Rc<BitSetVec>,
220    ) -> Result<&mut T, QuerySingleError> {
221        // SOUND: we know the schema matches.
222        fn map<T>(r: SchemaRefMut<'_>) -> &mut T {
223            unsafe { r.cast_into_mut_unchecked() }
224        }
225        self.untyped.get_single_with_bitset_mut(bitset).map(map)
226    }
227
228    /// Iterates immutably over the components of this type where `bitset`
229    /// indicates the indices of entities.
230    /// Slower than `iter()` but allows joining between multiple component types.
231    #[inline]
232    fn iter_with_bitset(&self, bitset: Rc<BitSetVec>) -> ComponentBitsetIterator<'_, T> {
233        // SOUND: we know the schema matches.
234        fn map<T>(r: SchemaRef<'_>) -> &T {
235            unsafe { r.cast_into_unchecked() }
236        }
237        self.untyped.iter_with_bitset(bitset).map(map)
238    }
239
240    /// Iterates immutably over the components of this type where `bitset`
241    /// indicates the indices of entities where iterator returns an Option.
242    /// None is returned for entities in bitset when Component is not in [`ComponentStore`]
243    #[inline]
244    fn iter_with_bitset_optional(
245        &self,
246        bitset: Rc<BitSetVec>,
247    ) -> ComponentBitsetOptionalIterator<'_, T> {
248        // SOUND: we know the schema matches.
249        fn map<T>(r: Option<SchemaRef<'_>>) -> Option<&T> {
250            r.map(|r| unsafe { r.cast_into_unchecked() })
251        }
252        self.untyped.iter_with_bitset_optional(bitset).map(map)
253    }
254
255    /// Iterates mutable over the components of this type where `bitset`
256    /// indicates the indices of entities.
257    /// Slower than `iter()` but allows joining between multiple component types.
258    #[inline]
259    fn iter_mut_with_bitset(&mut self, bitset: Rc<BitSetVec>) -> ComponentBitsetIteratorMut<'_, T> {
260        // SOUND: we know the schema matches.
261        fn map<T>(r: SchemaRefMut<'_>) -> &mut T {
262            unsafe { r.cast_into_mut_unchecked() }
263        }
264
265        self.untyped.iter_mut_with_bitset(bitset).map(map)
266    }
267
268    /// Iterates mutably over the components of this type where `bitset`
269    /// indicates the indices of entities where iterator returns an Option.
270    /// None is returned for entities in bitset when Component is not in [`ComponentStore`
271    #[inline]
272    fn iter_mut_with_bitset_optional(
273        &mut self,
274        bitset: Rc<BitSetVec>,
275    ) -> ComponentBitsetOptionalIteratorMut<'_, T> {
276        // SOUND: we know the schema matches.
277        fn map<T>(r: Option<SchemaRefMut<'_>>) -> Option<&mut T> {
278            r.map(|r| unsafe { r.cast_into_mut_unchecked() })
279        }
280        self.untyped.iter_mut_with_bitset_optional(bitset).map(map)
281    }
282
283    /// Read the bitset containing the list of entites with this component type on it.
284    #[inline]
285    fn bitset(&self) -> &BitSetVec {
286        self.untyped.bitset()
287    }
288
289    /// Check whether or not this component store has data for the given entity.
290    #[inline]
291    fn contains(&self, entity: Entity) -> bool {
292        self.bitset().contains(entity)
293    }
294
295    //// Get [`ComponentStore`] for usage with generic types implementing [`ComponentIterBitset`].
296    #[inline]
297    fn component_store(&self) -> &ComponentStore<T> {
298        self
299    }
300}
301
302#[cfg(test)]
303mod tests {
304    #![allow(non_snake_case)]
305
306    use std::{iter, rc::Rc};
307
308    use crate::prelude::*;
309
310    #[derive(Debug, Clone, Copy, PartialEq, Eq, HasSchema, Default)]
311    #[repr(C)]
312    struct A(u32);
313
314    fn entity(index: u32) -> Entity {
315        Entity::new(index, 0)
316    }
317
318    fn store(entities: &[u32]) -> ComponentStore<A> {
319        let mut store = ComponentStore::default();
320        for &i in entities {
321            store.insert(entity(i), A(i));
322        }
323        store
324    }
325
326    fn bitset(enabled: &[usize]) -> Rc<BitSetVec> {
327        let mut bitset = BitSetVec::default();
328        for &i in enabled {
329            bitset.bit_set(i);
330        }
331        Rc::new(bitset)
332    }
333
334    #[test]
335    fn create_remove_components() {
336        let mut entities = Entities::default();
337        let e1 = entities.create();
338        let e2 = entities.create();
339
340        let mut storage = ComponentStore::<A>::default();
341        storage.insert(e1, A(1));
342        storage.insert(e2, A(2));
343        assert!(storage.get(e1).is_some());
344        storage.remove(e1);
345        assert!(storage.get(e1).is_none());
346        assert!(storage.iter().eq(iter::once(&A(2))));
347    }
348
349    #[test]
350    fn get_mut_or_insert() {
351        let mut entities = Entities::default();
352        let e1 = entities.create();
353
354        let mut storage = ComponentStore::<A>::default();
355        {
356            // Test that inserted component is correct
357            let comp = storage.get_mut_or_insert(e1, || A(1));
358            assert_eq!(comp.0, 1);
359
360            // Mutate component
361            comp.0 = 2;
362        }
363
364        // Should not insert the unexpected value but retrieve original mutated component.
365        let comp = storage.get_mut_or_insert(e1, || A(u32::MAX));
366
367        // Test that existing component is retrieved
368        assert_eq!(comp.0, 2);
369    }
370
371    #[test]
372    fn iter() {
373        let mut store = store(&[10, 20, 30]);
374        let expected = vec![10, 20, 30];
375
376        let actual = store.iter().map(|c| c.0).collect::<Vec<_>>();
377        assert_eq!(actual, expected);
378
379        let actual = store.iter_mut().map(|c| c.0).collect::<Vec<_>>();
380        assert_eq!(actual, expected);
381    }
382
383    #[test]
384    fn get_single_with_bitset() {
385        {
386            let (mut store, bitset) = (store(&[]), bitset(&[]));
387            let result = store.get_single_with_bitset(bitset.clone());
388            assert_eq!(result, Err(QuerySingleError::NoEntities));
389            let result = store.get_single_with_bitset_mut(bitset);
390            assert_eq!(result, Err(QuerySingleError::NoEntities));
391        }
392
393        {
394            let (mut store, bitset) = (store(&[1]), bitset(&[]));
395            let result = store.get_single_with_bitset(bitset.clone());
396            assert_eq!(result, Err(QuerySingleError::NoEntities));
397            let result = store.get_single_with_bitset_mut(bitset);
398            assert_eq!(result, Err(QuerySingleError::NoEntities));
399        }
400
401        {
402            let (mut store, bitset) = (store(&[]), bitset(&[1]));
403            let result = store.get_single_with_bitset(bitset.clone());
404            assert_eq!(result, Err(QuerySingleError::NoEntities));
405            let result = store.get_single_with_bitset_mut(bitset);
406            assert_eq!(result, Err(QuerySingleError::NoEntities));
407        }
408
409        {
410            let (mut store, bitset) = (store(&[3]), bitset(&[3]));
411            let result = store.get_single_with_bitset(bitset.clone());
412            assert_eq!(result, Ok(&A(3)));
413            let result = store.get_single_with_bitset_mut(bitset);
414            assert_eq!(result, Ok(&mut A(3)));
415        }
416
417        {
418            let (mut store, bitset) = (store(&[5, 6]), bitset(&[5, 6]));
419            let result = store.get_single_with_bitset(bitset.clone());
420            assert_eq!(result, Err(QuerySingleError::MultipleEntities));
421            let result = store.get_single_with_bitset_mut(bitset);
422            assert_eq!(result, Err(QuerySingleError::MultipleEntities));
423        }
424
425        {
426            let (mut store, bitset) = (store(&[11, 21, 31]), bitset(&[12, 22, 32]));
427            let result = store.get_single_with_bitset(bitset.clone());
428            assert_eq!(result, Err(QuerySingleError::NoEntities));
429            let result = store.get_single_with_bitset_mut(bitset);
430            assert_eq!(result, Err(QuerySingleError::NoEntities));
431        }
432
433        {
434            let (mut store, bitset) = (store(&[11, 20, 31, 41]), bitset(&[12, 20, 32, 42]));
435            let result = store.get_single_with_bitset(bitset.clone());
436            assert_eq!(result, Ok(&A(20)));
437            let result = store.get_single_with_bitset_mut(bitset);
438            assert_eq!(result, Ok(&mut A(20)));
439        }
440
441        {
442            let (mut store, bitset) = (store(&[11, 21, 30, 41, 50]), bitset(&[12, 22, 30, 42, 50]));
443            let result = store.get_single_with_bitset(bitset.clone());
444            assert_eq!(result, Err(QuerySingleError::MultipleEntities));
445            let result = store.get_single_with_bitset_mut(bitset);
446            assert_eq!(result, Err(QuerySingleError::MultipleEntities));
447        }
448    }
449
450    #[test]
451    fn iter_with_bitset() {
452        {
453            let (mut store, bitset) = (store(&[]), bitset(&[]));
454
455            let mut iter = store.iter_with_bitset(bitset.clone());
456            assert_eq!(iter.next(), None);
457
458            let mut iter = store.iter_mut_with_bitset(bitset);
459            assert_eq!(iter.next(), None);
460        }
461
462        {
463            let (mut store, bitset) = (store(&[]), bitset(&[10, 20]));
464
465            let mut iter = store.iter_with_bitset(bitset.clone());
466            assert_eq!(iter.next(), None);
467
468            let mut iter = store.iter_mut_with_bitset(bitset);
469            assert_eq!(iter.next(), None);
470        }
471
472        {
473            let (mut store, bitset) = (store(&[10, 20]), bitset(&[]));
474
475            let mut iter = store.iter_with_bitset(bitset.clone());
476            assert_eq!(iter.next(), None);
477
478            let mut iter = store.iter_mut_with_bitset(bitset);
479            assert_eq!(iter.next(), None);
480        }
481
482        {
483            let (mut store, bitset) = (store(&[11, 21, 31]), bitset(&[12, 22, 32]));
484
485            let mut iter = store.iter_with_bitset(bitset.clone());
486            assert_eq!(iter.next(), None);
487
488            let mut iter = store.iter_mut_with_bitset(bitset.clone());
489            assert_eq!(iter.next(), None);
490        }
491
492        {
493            let (mut store, bitset) = (store(&[5]), bitset(&[5]));
494
495            let mut iter = store.iter_with_bitset(bitset.clone());
496            assert_eq!(iter.next(), Some(&A(5)));
497            assert_eq!(iter.next(), None);
498
499            let mut iter = store.iter_mut_with_bitset(bitset);
500            assert_eq!(iter.next(), Some(&mut A(5)));
501            assert_eq!(iter.next(), None);
502        }
503
504        {
505            let (mut store, bitset) = (
506                store(&[11, 20, 31, 41, 50, 61]),
507                bitset(&[12, 20, 32, 42, 50, 62]),
508            );
509
510            let mut iter = store.iter_with_bitset(bitset.clone());
511            assert_eq!(iter.next(), Some(&A(20)));
512            assert_eq!(iter.next(), Some(&A(50)));
513            assert_eq!(iter.next(), None);
514
515            let mut iter = store.iter_mut_with_bitset(bitset);
516            assert_eq!(iter.next(), Some(&mut A(20)));
517            assert_eq!(iter.next(), Some(&mut A(50)));
518            assert_eq!(iter.next(), None);
519        }
520    }
521
522    #[test]
523    fn iter_with_bitset_optional() {
524        {
525            let (mut store, bitset) = (store(&[]), bitset(&[]));
526
527            let mut iter = store.iter_with_bitset_optional(bitset.clone());
528            assert_eq!(iter.next(), None);
529
530            let mut iter = store.iter_mut_with_bitset_optional(bitset);
531            assert_eq!(iter.next(), None);
532        }
533
534        {
535            let (mut store, bitset) = (store(&[]), bitset(&[5, 6]));
536
537            let mut iter = store.iter_with_bitset_optional(bitset.clone());
538            assert_eq!(iter.next(), Some(None));
539            assert_eq!(iter.next(), Some(None));
540            assert_eq!(iter.next(), None);
541
542            let mut iter = store.iter_mut_with_bitset_optional(bitset);
543            assert_eq!(iter.next(), Some(None));
544            assert_eq!(iter.next(), Some(None));
545            assert_eq!(iter.next(), None);
546        }
547
548        {
549            let (mut store, bitset) = (store(&[11, 21, 31]), bitset(&[12, 22, 32]));
550
551            let mut iter = store.iter_with_bitset_optional(bitset.clone());
552            assert_eq!(iter.next(), Some(None));
553            assert_eq!(iter.next(), Some(None));
554            assert_eq!(iter.next(), Some(None));
555            assert_eq!(iter.next(), None);
556
557            let mut iter = store.iter_mut_with_bitset_optional(bitset);
558            assert_eq!(iter.next(), Some(None));
559            assert_eq!(iter.next(), Some(None));
560            assert_eq!(iter.next(), Some(None));
561            assert_eq!(iter.next(), None);
562        }
563
564        {
565            let (mut store, bitset) = (store(&[11, 20, 31]), bitset(&[12, 20, 32]));
566
567            let mut iter = store.iter_with_bitset_optional(bitset.clone());
568            assert_eq!(iter.next(), Some(None));
569            assert_eq!(iter.next(), Some(Some(&A(20))));
570            assert_eq!(iter.next(), Some(None));
571            assert_eq!(iter.next(), None);
572
573            let mut iter = store.iter_mut_with_bitset_optional(bitset);
574            assert_eq!(iter.next(), Some(None));
575            assert_eq!(iter.next(), Some(Some(&mut A(20))));
576            assert_eq!(iter.next(), Some(None));
577            assert_eq!(iter.next(), None);
578        }
579    }
580}
bones_ecs/components/typed.rs

bones_ecs/components/
typed.rs
