bevy_ecs/storage/
sparse_set.rs

1use crate::{
2    change_detection::{CheckChangeTicks, ComponentTickCells, ComponentTicks, MaybeLocation, Tick},
3    component::{ComponentId, ComponentInfo},
4    entity::{Entity, EntityIndex},
5    query::DebugCheckedUnwrap,
6    storage::{AbortOnPanic, Column, TableRow, VecExtensions},
7};
8use alloc::{boxed::Box, vec::Vec};
9use bevy_ptr::{OwningPtr, Ptr};
10use core::{cell::UnsafeCell, hash::Hash, marker::PhantomData, num::NonZero, panic::Location};
11use nonmax::{NonMaxU32, NonMaxUsize};
12
13#[derive(Debug)]
14pub(crate) struct SparseArray<I, V = I> {
15    values: Vec<Option<V>>,
16    marker: PhantomData<I>,
17}
18
19/// A space-optimized version of [`SparseArray`] that cannot be changed
20/// after construction.
21#[derive(Debug)]
22pub(crate) struct ImmutableSparseArray<I, V = I> {
23    values: Box<[Option<V>]>,
24    marker: PhantomData<I>,
25}
26
27impl<I: SparseSetIndex, V> Default for SparseArray<I, V> {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33impl<I, V> SparseArray<I, V> {
34    #[inline]
35    pub const fn new() -> Self {
36        Self {
37            values: Vec::new(),
38            marker: PhantomData,
39        }
40    }
41}
42
43macro_rules! impl_sparse_array {
44    ($ty:ident) => {
45        impl<I: SparseSetIndex, V> $ty<I, V> {
46            /// Returns `true` if the collection contains a value for the specified `index`.
47            #[inline]
48            pub fn contains(&self, index: I) -> bool {
49                let index = index.sparse_set_index();
50                self.values.get(index).is_some_and(Option::is_some)
51            }
52
53            /// Returns a reference to the value at `index`.
54            ///
55            /// Returns `None` if `index` does not have a value or if `index` is out of bounds.
56            #[inline]
57            pub fn get(&self, index: I) -> Option<&V> {
58                let index = index.sparse_set_index();
59                self.values.get(index).and_then(Option::as_ref)
60            }
61        }
62    };
63}
64
65impl_sparse_array!(SparseArray);
66impl_sparse_array!(ImmutableSparseArray);
67
68impl<I: SparseSetIndex, V> SparseArray<I, V> {
69    /// Inserts `value` at `index` in the array.
70    ///
71    /// # Panics
72    /// - Panics if the insertion forces a reallocation, and any of the new capacity overflows `isize::MAX` bytes.
73    /// - Panics if the insertion forces a reallocation, and any of the new the reallocations causes an out-of-memory error.
74    ///
75    /// If `index` is out-of-bounds, this will enlarge the buffer to accommodate it.
76    #[inline]
77    pub fn insert(&mut self, index: I, value: V) {
78        let index = index.sparse_set_index();
79        if index >= self.values.len() {
80            self.values.resize_with(index + 1, || None);
81        }
82        self.values[index] = Some(value);
83    }
84
85    /// Returns a mutable reference to the value at `index`.
86    ///
87    /// Returns `None` if `index` does not have a value or if `index` is out of bounds.
88    #[inline]
89    pub fn get_mut(&mut self, index: I) -> Option<&mut V> {
90        let index = index.sparse_set_index();
91        self.values.get_mut(index).and_then(Option::as_mut)
92    }
93
94    /// Removes and returns the value stored at `index`.
95    ///
96    /// Returns `None` if `index` did not have a value or if `index` is out of bounds.
97    #[inline]
98    pub fn remove(&mut self, index: I) -> Option<V> {
99        let index = index.sparse_set_index();
100        self.values.get_mut(index).and_then(Option::take)
101    }
102
103    /// Removes all of the values stored within.
104    pub fn clear(&mut self) {
105        self.values.clear();
106    }
107
108    /// Converts the [`SparseArray`] into an immutable variant.
109    pub(crate) fn into_immutable(self) -> ImmutableSparseArray<I, V> {
110        ImmutableSparseArray {
111            values: self.values.into_boxed_slice(),
112            marker: PhantomData,
113        }
114    }
115}
116
117/// A sparse data structure of [`Component`](crate::component::Component)s.
118///
119/// Designed for relatively fast insertions and deletions.
120#[derive(Debug)]
121pub struct ComponentSparseSet {
122    /// Capacity and length match those of `entities`.
123    dense: Column,
124    // Internally this only relies on the Entity index to keep track of where the component data is
125    // stored for entities that are alive. The generation is not required, but is stored
126    // in debug builds to validate that access is correct.
127    #[cfg(not(debug_assertions))]
128    entities: Vec<EntityIndex>,
129    #[cfg(debug_assertions)]
130    entities: Vec<Entity>,
131    sparse: SparseArray<EntityIndex, TableRow>,
132}
133
134impl ComponentSparseSet {
135    /// Creates a new [`ComponentSparseSet`] with a given component type layout and
136    /// initial `capacity`.
137    pub(crate) fn new(component_info: &ComponentInfo, capacity: usize) -> Self {
138        let entities = Vec::with_capacity(capacity);
139        Self {
140            dense: Column::with_capacity(component_info, entities.capacity()),
141            entities,
142            sparse: Default::default(),
143        }
144    }
145
146    /// Removes all of the values stored within.
147    pub(crate) fn clear(&mut self) {
148        // SAFETY: This is using the size of the ComponentSparseSet.
149        unsafe { self.dense.clear(self.len()) };
150        self.entities.clear();
151        self.sparse.clear();
152    }
153
154    /// Returns the number of component values in the sparse set.
155    #[inline]
156    pub fn len(&self) -> usize {
157        self.entities.len()
158    }
159
160    /// Returns `true` if the sparse set contains no component values.
161    #[inline]
162    pub fn is_empty(&self) -> bool {
163        self.entities.is_empty()
164    }
165
166    /// Inserts the `entity` key and component `value` pair into this sparse
167    /// set.
168    ///
169    /// # Aborts
170    /// - Aborts the process if the insertion forces a reallocation, and any of the new capacity overflows `isize::MAX` bytes.
171    /// - Aborts the process if the insertion forces a reallocation, and any of the new the reallocations causes an out-of-memory error.
172    ///
173    /// # Safety
174    /// The `value` pointer must point to a valid address that matches the [`Layout`](std::alloc::Layout)
175    /// inside the [`ComponentInfo`] given when constructing this sparse set.
176    pub(crate) unsafe fn insert(
177        &mut self,
178        entity: Entity,
179        value: OwningPtr<'_>,
180        change_tick: Tick,
181        caller: MaybeLocation,
182    ) {
183        if let Some(&dense_index) = self.sparse.get(entity.index()) {
184            #[cfg(debug_assertions)]
185            assert_eq!(entity, self.entities[dense_index.index()]);
186            self.dense.replace(dense_index, value, change_tick, caller);
187        } else {
188            let dense_index = self.entities.len();
189            let capacity = self.entities.capacity();
190
191            #[cfg(not(debug_assertions))]
192            self.entities.push(entity.index());
193            #[cfg(debug_assertions)]
194            self.entities.push(entity);
195
196            // If any of the following operations panic due to an allocation error, the state
197            // of the `ComponentSparseSet` will be left in an invalid state and potentially cause UB.
198            // We create an AbortOnPanic guard to force panics to terminate the process if this occurs.
199            let _guard = AbortOnPanic;
200            if capacity != self.entities.capacity() {
201                // SAFETY: An entity was just pushed onto `entities`, its capacity cannot be zero.
202                let new_capacity = unsafe { NonZero::new_unchecked(self.entities.capacity()) };
203                if let Some(capacity) = NonZero::new(capacity) {
204                    // SAFETY: This is using the capacity of the previous allocation.
205                    unsafe { self.dense.realloc(capacity, new_capacity) };
206                } else {
207                    self.dense.alloc(new_capacity);
208                }
209            }
210
211            // SAFETY: This entity index does not exist here yet, so there are no duplicates,
212            // and the entity index is `NonMaxU32` so the length must not be max either.
213            let table_row = unsafe { TableRow::new(NonMaxU32::new_unchecked(dense_index as u32)) };
214            self.dense.initialize(table_row, value, change_tick, caller);
215            self.sparse.insert(entity.index(), table_row);
216
217            core::mem::forget(_guard);
218        }
219    }
220
221    /// Returns `true` if the sparse set has a component value for the provided `entity`.
222    #[inline]
223    pub fn contains(&self, entity: Entity) -> bool {
224        #[cfg(debug_assertions)]
225        {
226            if let Some(&dense_index) = self.sparse.get(entity.index()) {
227                #[cfg(debug_assertions)]
228                assert_eq!(entity, self.entities[dense_index.index()]);
229                true
230            } else {
231                false
232            }
233        }
234        #[cfg(not(debug_assertions))]
235        self.sparse.contains(entity.index())
236    }
237
238    /// Returns a reference to the entity's component value.
239    ///
240    /// Returns `None` if `entity` does not have a component in the sparse set.
241    #[inline]
242    pub fn get(&self, entity: Entity) -> Option<Ptr<'_>> {
243        self.sparse.get(entity.index()).map(|&dense_index| {
244            #[cfg(debug_assertions)]
245            assert_eq!(entity, self.entities[dense_index.index()]);
246            // SAFETY: if the sparse index points to something in the dense vec, it exists
247            unsafe { self.dense.get_data_unchecked(dense_index) }
248        })
249    }
250
251    /// Returns references to the entity's component value and its added and changed ticks.
252    ///
253    /// Returns `None` if `entity` does not have a component in the sparse set.
254    #[inline]
255    pub fn get_with_ticks(&self, entity: Entity) -> Option<(Ptr<'_>, ComponentTickCells<'_>)> {
256        let dense_index = *self.sparse.get(entity.index())?;
257        #[cfg(debug_assertions)]
258        assert_eq!(entity, self.entities[dense_index.index()]);
259        // SAFETY: if the sparse index points to something in the dense vec, it exists
260        unsafe {
261            Some((
262                self.dense.get_data_unchecked(dense_index),
263                ComponentTickCells {
264                    added: self.dense.get_added_tick_unchecked(dense_index),
265                    changed: self.dense.get_changed_tick_unchecked(dense_index),
266                    changed_by: self.dense.get_changed_by_unchecked(dense_index),
267                },
268            ))
269        }
270    }
271
272    /// Returns a reference to the "added" tick of the entity's component value.
273    ///
274    /// Returns `None` if `entity` does not have a component in the sparse set.
275    #[inline]
276    pub fn get_added_tick(&self, entity: Entity) -> Option<&UnsafeCell<Tick>> {
277        let dense_index = *self.sparse.get(entity.index())?;
278        #[cfg(debug_assertions)]
279        assert_eq!(entity, self.entities[dense_index.index()]);
280        // SAFETY: if the sparse index points to something in the dense vec, it exists
281        unsafe { Some(self.dense.get_added_tick_unchecked(dense_index)) }
282    }
283
284    /// Returns a reference to the "changed" tick of the entity's component value.
285    ///
286    /// Returns `None` if `entity` does not have a component in the sparse set.
287    #[inline]
288    pub fn get_changed_tick(&self, entity: Entity) -> Option<&UnsafeCell<Tick>> {
289        let dense_index = *self.sparse.get(entity.index())?;
290        #[cfg(debug_assertions)]
291        assert_eq!(entity, self.entities[dense_index.index()]);
292        // SAFETY: if the sparse index points to something in the dense vec, it exists
293        unsafe { Some(self.dense.get_changed_tick_unchecked(dense_index)) }
294    }
295
296    /// Returns a reference to the "added" and "changed" ticks of the entity's component value.
297    ///
298    /// Returns `None` if `entity` does not have a component in the sparse set.
299    #[inline]
300    pub fn get_ticks(&self, entity: Entity) -> Option<ComponentTicks> {
301        let dense_index = *self.sparse.get(entity.index())?;
302        #[cfg(debug_assertions)]
303        assert_eq!(entity, self.entities[dense_index.index()]);
304        // SAFETY: if the sparse index points to something in the dense vec, it exists
305        unsafe { Some(self.dense.get_ticks_unchecked(dense_index)) }
306    }
307
308    /// Returns a reference to the calling location that last changed the entity's component value.
309    ///
310    /// Returns `None` if `entity` does not have a component in the sparse set.
311    #[inline]
312    pub fn get_changed_by(
313        &self,
314        entity: Entity,
315    ) -> MaybeLocation<Option<&UnsafeCell<&'static Location<'static>>>> {
316        MaybeLocation::new_with_flattened(|| {
317            let dense_index = *self.sparse.get(entity.index())?;
318            #[cfg(debug_assertions)]
319            assert_eq!(entity, self.entities[dense_index.index()]);
320            // SAFETY: if the sparse index points to something in the dense vec, it exists
321            unsafe { Some(self.dense.get_changed_by_unchecked(dense_index)) }
322        })
323    }
324
325    /// Returns the drop function for the component type stored in the sparse set,
326    /// or `None` if it doesn't need to be dropped.
327    #[inline]
328    pub fn get_drop(&self) -> Option<unsafe fn(OwningPtr<'_>)> {
329        self.dense.get_drop()
330    }
331
332    /// Removes the `entity` from this sparse set and returns a pointer to the associated value (if
333    /// it exists).
334    #[must_use = "The returned pointer must be used to drop the removed component."]
335    pub(crate) fn remove_and_forget(&mut self, entity: Entity) -> Option<OwningPtr<'_>> {
336        self.sparse.remove(entity.index()).map(|dense_index| {
337            #[cfg(debug_assertions)]
338            assert_eq!(entity, self.entities[dense_index.index()]);
339            let last = self.entities.len() - 1;
340            if dense_index.index() >= last {
341                #[cfg(debug_assertions)]
342                assert_eq!(dense_index.index(), last);
343                // SAFETY: This is strictly decreasing the length, so it cannot outgrow
344                // it also cannot underflow as an item was just removed from the sparse array.
345                unsafe { self.entities.set_len(last) };
346                // SAFETY: `last` is guaranteed to be the last element in `dense` as the length is synced with
347                // the `entities` store.
348                unsafe {
349                    self.dense
350                        .get_data_unchecked(dense_index)
351                        .assert_unique()
352                        .promote()
353                }
354            } else {
355                // SAFETY: The above check ensures that `dense_index` and the last element are not
356                // overlapping, and thus also within bounds.
357                unsafe {
358                    self.entities
359                        .swap_remove_nonoverlapping_unchecked(dense_index.index());
360                };
361                // SAFETY: The above check ensures that `dense_index` is in bounds.
362                let swapped_entity = unsafe { self.entities.get_unchecked(dense_index.index()) };
363                #[cfg(not(debug_assertions))]
364                let index = *swapped_entity;
365                #[cfg(debug_assertions)]
366                let index = swapped_entity.index();
367                // SAFETY: The swapped entity was just fetched from the entity Vec, it must have already
368                // been inserted and in bounds.
369                unsafe {
370                    *self.sparse.get_mut(index).debug_checked_unwrap() = dense_index;
371                }
372                // SAFETY: The above check ensures that `dense_index` and the last element are not
373                // overlapping, and thus also within bounds.
374                unsafe {
375                    self.dense
376                        .swap_remove_and_forget_unchecked_nonoverlapping(last, dense_index)
377                }
378            }
379        })
380    }
381
382    /// Removes (and drops) the entity's component value from the sparse set.
383    ///
384    /// Returns `true` if `entity` had a component value in the sparse set.
385    pub(crate) fn remove(&mut self, entity: Entity) -> bool {
386        self.sparse
387            .remove(entity.index())
388            .map(|dense_index| {
389                #[cfg(debug_assertions)]
390                assert_eq!(entity, self.entities[dense_index.index()]);
391                let last = self.entities.len() - 1;
392                if dense_index.index() >= last {
393                    #[cfg(debug_assertions)]
394                    assert_eq!(dense_index.index(), last);
395                    // SAFETY: This is strictly decreasing the length, so it cannot outgrow
396                    // it also cannot underflow as an item was just removed from the sparse array.
397                    unsafe { self.entities.set_len(last) };
398                    // SAFETY: `last` is guaranteed to be the last element in `dense` as the length is synced with
399                    // the `entities` store.
400                    unsafe { self.dense.drop_last_component(last) };
401                } else {
402                    // SAFETY: The above check ensures that `dense_index` and the last element are not
403                    // overlapping, and thus also within bounds.
404                    unsafe {
405                        self.entities
406                            .swap_remove_nonoverlapping_unchecked(dense_index.index());
407                    };
408                    let swapped_entity =
409                        // SAFETY: The above check ensures that `dense_index` is in bounds.
410                        unsafe { self.entities.get_unchecked(dense_index.index()) };
411                    #[cfg(not(debug_assertions))]
412                    let index = *swapped_entity;
413                    #[cfg(debug_assertions)]
414                    let index = swapped_entity.index();
415                    // SAFETY: The swapped entity was just fetched from the entity Vec, it must have already
416                    // been inserted and in bounds.
417                    unsafe {
418                        *self.sparse.get_mut(index).debug_checked_unwrap() = dense_index;
419                    }
420                    // SAFETY: The above check ensures that `dense_index` and the last element are not
421                    // overlapping, and thus also within bounds.
422                    unsafe {
423                        self.dense
424                            .swap_remove_and_drop_unchecked_nonoverlapping(last, dense_index);
425                    }
426                }
427            })
428            .is_some()
429    }
430
431    pub(crate) fn check_change_ticks(&mut self, check: CheckChangeTicks) {
432        // SAFETY: This is using the valid size of the column.
433        unsafe { self.dense.check_change_ticks(self.len(), check) };
434    }
435}
436
437impl Drop for ComponentSparseSet {
438    fn drop(&mut self) {
439        let len = self.entities.len();
440        self.entities.clear();
441        // SAFETY: `cap` and `len` are correct. `dense` is never accessed again after this call.
442        unsafe {
443            self.dense.drop(self.entities.capacity(), len);
444        }
445    }
446}
447
448/// A data structure that blends dense and sparse storage
449///
450/// `I` is the type of the indices, while `V` is the type of data stored in the dense storage.
451#[derive(Debug)]
452pub struct SparseSet<I, V: 'static> {
453    dense: Vec<V>,
454    indices: Vec<I>,
455    sparse: SparseArray<I, NonMaxUsize>,
456}
457
458/// A space-optimized version of [`SparseSet`] that cannot be changed
459/// after construction.
460#[derive(Debug)]
461pub(crate) struct ImmutableSparseSet<I, V: 'static> {
462    dense: Box<[V]>,
463    indices: Box<[I]>,
464    sparse: ImmutableSparseArray<I, NonMaxUsize>,
465}
466
467macro_rules! impl_sparse_set {
468    ($ty:ident) => {
469        impl<I: SparseSetIndex, V> $ty<I, V> {
470            /// Returns the number of elements in the sparse set.
471            #[inline]
472            pub fn len(&self) -> usize {
473                self.dense.len()
474            }
475
476            /// Returns `true` if the sparse set contains a value for `index`.
477            #[inline]
478            pub fn contains(&self, index: I) -> bool {
479                self.sparse.contains(index)
480            }
481
482            /// Returns a reference to the value for `index`.
483            ///
484            /// Returns `None` if `index` does not have a value in the sparse set.
485            pub fn get(&self, index: I) -> Option<&V> {
486                self.sparse.get(index).map(|dense_index| {
487                    // SAFETY: if the sparse index points to something in the dense vec, it exists
488                    unsafe { self.dense.get_unchecked(dense_index.get()) }
489                })
490            }
491
492            /// Returns a mutable reference to the value for `index`.
493            ///
494            /// Returns `None` if `index` does not have a value in the sparse set.
495            pub fn get_mut(&mut self, index: I) -> Option<&mut V> {
496                let dense = &mut self.dense;
497                self.sparse.get(index).map(move |dense_index| {
498                    // SAFETY: if the sparse index points to something in the dense vec, it exists
499                    unsafe { dense.get_unchecked_mut(dense_index.get()) }
500                })
501            }
502
503            /// Returns an iterator visiting all keys (indices) in arbitrary order.
504            pub fn indices(&self) -> &[I] {
505                &self.indices
506            }
507
508            /// Returns an iterator visiting all values in arbitrary order.
509            pub fn values(&self) -> impl Iterator<Item = &V> {
510                self.dense.iter()
511            }
512
513            /// Returns an iterator visiting all values mutably in arbitrary order.
514            pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
515                self.dense.iter_mut()
516            }
517
518            /// Returns an iterator visiting all key-value pairs in arbitrary order, with references to the values.
519            pub fn iter(&self) -> impl Iterator<Item = (&I, &V)> {
520                self.indices.iter().zip(self.dense.iter())
521            }
522
523            /// Returns an iterator visiting all key-value pairs in arbitrary order, with mutable references to the values.
524            pub fn iter_mut(&mut self) -> impl Iterator<Item = (&I, &mut V)> {
525                self.indices.iter().zip(self.dense.iter_mut())
526            }
527        }
528    };
529}
530
531impl_sparse_set!(SparseSet);
532impl_sparse_set!(ImmutableSparseSet);
533
534impl<I: SparseSetIndex, V> Default for SparseSet<I, V> {
535    fn default() -> Self {
536        Self::new()
537    }
538}
539
540impl<I, V> SparseSet<I, V> {
541    /// Creates a new [`SparseSet`].
542    pub const fn new() -> Self {
543        Self {
544            dense: Vec::new(),
545            indices: Vec::new(),
546            sparse: SparseArray::new(),
547        }
548    }
549}
550
551impl<I: SparseSetIndex, V> SparseSet<I, V> {
552    /// Creates a new [`SparseSet`] with a specified initial capacity.
553    ///
554    /// # Panics
555    /// - Panics if the new capacity of the allocation overflows `isize::MAX` bytes.
556    /// - Panics if the new allocation causes an out-of-memory error.
557    pub fn with_capacity(capacity: usize) -> Self {
558        Self {
559            dense: Vec::with_capacity(capacity),
560            indices: Vec::with_capacity(capacity),
561            sparse: Default::default(),
562        }
563    }
564
565    /// Returns the total number of elements the [`SparseSet`] can hold without needing to reallocate.
566    #[inline]
567    pub fn capacity(&self) -> usize {
568        self.dense.capacity()
569    }
570
571    /// Inserts `value` at `index`.
572    ///
573    /// If a value was already present at `index`, it will be overwritten.
574    ///
575    /// # Panics
576    /// - Panics if the insertion forces an reallocation and the new capacity overflows `isize::MAX` bytes.
577    /// - Panics if the insertion forces an reallocation and causes an out-of-memory error.
578    pub fn insert(&mut self, index: I, value: V) {
579        if let Some(dense_index) = self.sparse.get(index.clone()).cloned() {
580            // SAFETY: dense indices stored in self.sparse always exist
581            unsafe {
582                *self.dense.get_unchecked_mut(dense_index.get()) = value;
583            }
584        } else {
585            self.sparse
586                .insert(index.clone(), NonMaxUsize::new(self.dense.len()).unwrap());
587            self.indices.push(index);
588            self.dense.push(value);
589        }
590    }
591
592    /// Returns a reference to the value for `index`, inserting one computed from `func`
593    /// if not already present.
594    ///
595    /// # Panics
596    /// - Panics if the insertion forces an reallocation and the new capacity overflows `isize::MAX` bytes.
597    /// - Panics if the insertion forces an reallocation and causes an out-of-memory error.
598    pub fn get_or_insert_with(&mut self, index: I, func: impl FnOnce() -> V) -> &mut V {
599        if let Some(dense_index) = self.sparse.get(index.clone()).cloned() {
600            // SAFETY: dense indices stored in self.sparse always exist
601            unsafe { self.dense.get_unchecked_mut(dense_index.get()) }
602        } else {
603            let value = func();
604            let dense_index = self.dense.len();
605            self.sparse
606                .insert(index.clone(), NonMaxUsize::new(dense_index).unwrap());
607            self.indices.push(index);
608            self.dense.push(value);
609            // SAFETY: dense index was just populated above
610            unsafe { self.dense.get_unchecked_mut(dense_index) }
611        }
612    }
613
614    /// Returns `true` if the sparse set contains no elements.
615    #[inline]
616    pub fn is_empty(&self) -> bool {
617        self.dense.len() == 0
618    }
619
620    /// Removes and returns the value for `index`.
621    ///
622    /// Returns `None` if `index` does not have a value in the sparse set.
623    pub fn remove(&mut self, index: I) -> Option<V> {
624        self.sparse.remove(index).map(|dense_index| {
625            let index = dense_index.get();
626            let is_last = index == self.dense.len() - 1;
627            let value = self.dense.swap_remove(index);
628            self.indices.swap_remove(index);
629            if !is_last {
630                let swapped_index = self.indices[index].clone();
631                *self.sparse.get_mut(swapped_index).unwrap() = dense_index;
632            }
633            value
634        })
635    }
636
637    /// Clears all of the elements from the sparse set.
638    ///
639    /// # Panics
640    /// - Panics if any of the keys or values implements [`Drop`] and any of those panic.
641    pub fn clear(&mut self) {
642        self.dense.clear();
643        self.indices.clear();
644        self.sparse.clear();
645    }
646
647    /// Converts the sparse set into its immutable variant.
648    pub(crate) fn into_immutable(self) -> ImmutableSparseSet<I, V> {
649        ImmutableSparseSet {
650            dense: self.dense.into_boxed_slice(),
651            indices: self.indices.into_boxed_slice(),
652            sparse: self.sparse.into_immutable(),
653        }
654    }
655}
656
657/// Represents something that can be stored in a [`SparseSet`] as an integer.
658///
659/// Ideally, the `usize` values should be very small (ie: incremented starting from
660/// zero), as the number of bits needed to represent a `SparseSetIndex` in a `FixedBitSet`
661/// is proportional to the **value** of those `usize`.
662pub trait SparseSetIndex: Clone + PartialEq + Eq + Hash {
663    /// Gets the sparse set index corresponding to this instance.
664    fn sparse_set_index(&self) -> usize;
665    /// Creates a new instance of this type with the specified index.
666    fn get_sparse_set_index(value: usize) -> Self;
667}
668
669macro_rules! impl_sparse_set_index {
670    ($($ty:ty),+) => {
671        $(impl SparseSetIndex for $ty {
672            #[inline]
673            fn sparse_set_index(&self) -> usize {
674                *self as usize
675            }
676
677            #[inline]
678            fn get_sparse_set_index(value: usize) -> Self {
679                value as $ty
680            }
681        })*
682    };
683}
684
685impl_sparse_set_index!(u8, u16, u32, u64, usize);
686
687/// A collection of [`ComponentSparseSet`] storages, indexed by [`ComponentId`]
688///
689/// Can be accessed via [`Storages`](crate::storage::Storages)
690#[derive(Default)]
691pub struct SparseSets {
692    sets: SparseSet<ComponentId, ComponentSparseSet>,
693}
694
695impl SparseSets {
696    /// Returns the number of [`ComponentSparseSet`]s this collection contains.
697    #[inline]
698    pub fn len(&self) -> usize {
699        self.sets.len()
700    }
701
702    /// Returns true if this collection contains no [`ComponentSparseSet`]s.
703    #[inline]
704    pub fn is_empty(&self) -> bool {
705        self.sets.is_empty()
706    }
707
708    /// An Iterator visiting all ([`ComponentId`], [`ComponentSparseSet`]) pairs.
709    /// NOTE: Order is not guaranteed.
710    pub fn iter(&self) -> impl Iterator<Item = (ComponentId, &ComponentSparseSet)> {
711        self.sets.iter().map(|(id, data)| (*id, data))
712    }
713
714    /// Gets a reference to the [`ComponentSparseSet`] of a [`ComponentId`]. This may be `None` if the component has never been spawned.
715    #[inline]
716    pub fn get(&self, component_id: ComponentId) -> Option<&ComponentSparseSet> {
717        self.sets.get(component_id)
718    }
719
720    /// Gets a mutable reference of [`ComponentSparseSet`] of a [`ComponentInfo`].
721    /// Create a new [`ComponentSparseSet`] if not exists.
722    ///
723    /// # Panics
724    /// - Panics if the insertion forces an reallocation and the new capacity overflows `isize::MAX` bytes.
725    /// - Panics if the insertion forces an reallocation and causes an out-of-memory error.
726    pub(crate) fn get_or_insert(
727        &mut self,
728        component_info: &ComponentInfo,
729    ) -> &mut ComponentSparseSet {
730        if !self.sets.contains(component_info.id()) {
731            self.sets.insert(
732                component_info.id(),
733                ComponentSparseSet::new(component_info, 64),
734            );
735        }
736
737        self.sets.get_mut(component_info.id()).unwrap()
738    }
739
740    /// Gets a mutable reference to the [`ComponentSparseSet`] of a [`ComponentId`]. This may be `None` if the component has never been spawned.
741    pub(crate) fn get_mut(&mut self, component_id: ComponentId) -> Option<&mut ComponentSparseSet> {
742        self.sets.get_mut(component_id)
743    }
744
745    /// Clear entities stored in each [`ComponentSparseSet`]
746    ///
747    /// # Panics
748    /// - Panics if any of the components stored within implement [`Drop`] and any of them panic.
749    pub(crate) fn clear_entities(&mut self) {
750        for set in self.sets.values_mut() {
751            set.clear();
752        }
753    }
754
755    pub(crate) fn check_change_ticks(&mut self, check: CheckChangeTicks) {
756        for set in self.sets.values_mut() {
757            set.check_change_ticks(check);
758        }
759    }
760}
761
762#[cfg(test)]
763mod tests {
764    use super::SparseSets;
765    use crate::{
766        component::{Component, ComponentDescriptor, ComponentId, ComponentInfo},
767        entity::{Entity, EntityIndex},
768        storage::SparseSet,
769    };
770    use alloc::{vec, vec::Vec};
771
772    #[derive(Debug, Eq, PartialEq)]
773    struct Foo(usize);
774
775    #[test]
776    fn sparse_set() {
777        let mut set = SparseSet::<Entity, Foo>::default();
778        let e0 = Entity::from_index(EntityIndex::from_raw_u32(0).unwrap());
779        let e1 = Entity::from_index(EntityIndex::from_raw_u32(1).unwrap());
780        let e2 = Entity::from_index(EntityIndex::from_raw_u32(2).unwrap());
781        let e3 = Entity::from_index(EntityIndex::from_raw_u32(3).unwrap());
782        let e4 = Entity::from_index(EntityIndex::from_raw_u32(4).unwrap());
783
784        set.insert(e1, Foo(1));
785        set.insert(e2, Foo(2));
786        set.insert(e3, Foo(3));
787
788        assert_eq!(set.get(e0), None);
789        assert_eq!(set.get(e1), Some(&Foo(1)));
790        assert_eq!(set.get(e2), Some(&Foo(2)));
791        assert_eq!(set.get(e3), Some(&Foo(3)));
792        assert_eq!(set.get(e4), None);
793
794        {
795            let iter_results = set.values().collect::<Vec<_>>();
796            assert_eq!(iter_results, vec![&Foo(1), &Foo(2), &Foo(3)]);
797        }
798
799        assert_eq!(set.remove(e2), Some(Foo(2)));
800        assert_eq!(set.remove(e2), None);
801
802        assert_eq!(set.get(e0), None);
803        assert_eq!(set.get(e1), Some(&Foo(1)));
804        assert_eq!(set.get(e2), None);
805        assert_eq!(set.get(e3), Some(&Foo(3)));
806        assert_eq!(set.get(e4), None);
807
808        assert_eq!(set.remove(e1), Some(Foo(1)));
809
810        assert_eq!(set.get(e0), None);
811        assert_eq!(set.get(e1), None);
812        assert_eq!(set.get(e2), None);
813        assert_eq!(set.get(e3), Some(&Foo(3)));
814        assert_eq!(set.get(e4), None);
815
816        set.insert(e1, Foo(10));
817
818        assert_eq!(set.get(e1), Some(&Foo(10)));
819
820        *set.get_mut(e1).unwrap() = Foo(11);
821        assert_eq!(set.get(e1), Some(&Foo(11)));
822    }
823
824    #[test]
825    fn sparse_sets() {
826        let mut sets = SparseSets::default();
827
828        #[derive(Component, Default, Debug)]
829        struct TestComponent1;
830
831        #[derive(Component, Default, Debug)]
832        struct TestComponent2;
833
834        assert_eq!(sets.len(), 0);
835        assert!(sets.is_empty());
836
837        register_component::<TestComponent1>(&mut sets, 1);
838        assert_eq!(sets.len(), 1);
839
840        register_component::<TestComponent2>(&mut sets, 2);
841        assert_eq!(sets.len(), 2);
842
843        // check its shape by iter
844        let mut collected_sets = sets
845            .iter()
846            .map(|(id, set)| (id, set.len()))
847            .collect::<Vec<_>>();
848        collected_sets.sort();
849        assert_eq!(
850            collected_sets,
851            vec![(ComponentId::new(1), 0), (ComponentId::new(2), 0),]
852        );
853
854        fn register_component<T: Component>(sets: &mut SparseSets, id: usize) {
855            let descriptor = ComponentDescriptor::new::<T>();
856            let id = ComponentId::new(id);
857            let info = ComponentInfo::new(id, descriptor);
858            sets.get_or_insert(&info);
859        }
860    }
861}
bevy_ecs/storage/sparse_set.rs

bevy_ecs/storage/
sparse_set.rs
