[bevy_<util/ecs>] Better SparseArray performance

crates/bevy_ecs/src/storage/sparse_set.rs

@@ -3,6 +3,7 @@ use crate::{
     entity::Entity,
     storage::BlobVec,
 };
+use bevy_utils::NonMaxUsize;
 use std::{cell::UnsafeCell, marker::PhantomData};
 
 #[derive(Debug)]
@@ -89,7 +90,7 @@ pub struct ComponentSparseSet {
     dense: BlobVec,
     ticks: UnsafeCell<Vec<ComponentTicks>>,
     entities: Vec<Entity>,
-    sparse: SparseArray<Entity, usize>,
+    sparse: SparseArray<Entity, NonMaxUsize>,
 }
 
 impl ComponentSparseSet {
@@ -123,11 +124,14 @@ impl ComponentSparseSet {
         let dense = &mut self.dense;
         let entities = &mut self.entities;
         let ticks_list = self.ticks.get_mut();
-        let dense_index = *self.sparse.get_or_insert_with(entity, move || {
-            ticks_list.push(ComponentTicks::new(change_tick));
-            entities.push(entity);
-            dense.push_uninit()
-        });
+        let dense_index = self
+            .sparse
+            .get_or_insert_with(entity, move || {
+                ticks_list.push(ComponentTicks::new(change_tick));
+                entities.push(entity);
+                NonMaxUsize::new(dense.push_uninit()).unwrap()
+            })
+            .get();
         // SAFE: dense_index exists thanks to the call above
         self.dense.set_unchecked(dense_index, value);
         ((*self.ticks.get()).get_unchecked_mut(dense_index)).set_changed(change_tick);
@@ -144,7 +148,7 @@ impl ComponentSparseSet {
     pub fn get(&self, entity: Entity) -> Option<*mut u8> {
         self.sparse.get(entity).map(|dense_index| {
             // SAFE: if the sparse index points to something in the dense vec, it exists
-            unsafe { self.dense.get_unchecked(*dense_index) }
+            unsafe { self.dense.get_unchecked(dense_index.get()) }
         })
     }
 
@@ -154,7 +158,7 @@ impl ComponentSparseSet {
     pub unsafe fn get_with_ticks(&self, entity: Entity) -> Option<(*mut u8, *mut ComponentTicks)> {
         let ticks = &mut *self.ticks.get();
         self.sparse.get(entity).map(move |dense_index| {
-            let dense_index = *dense_index;
+            let dense_index = dense_index.get();
             // SAFE: if the sparse index points to something in the dense vec, it exists
             (
                 self.dense.get_unchecked(dense_index),
@@ -169,9 +173,8 @@ impl ComponentSparseSet {
     pub unsafe fn get_ticks(&self, entity: Entity) -> Option<&mut ComponentTicks> {
         let ticks = &mut *self.ticks.get();
         self.sparse.get(entity).map(move |dense_index| {
-            let dense_index = *dense_index;
             // SAFE: if the sparse index points to something in the dense vec, it exists
-            ticks.get_unchecked_mut(dense_index)
+            ticks.get_unchecked_mut(dense_index.get())
         })
     }
 
@@ -180,16 +183,20 @@ impl ComponentSparseSet {
     /// returned).
     pub fn remove_and_forget(&mut self, entity: Entity) -> Option<*mut u8> {
         self.sparse.remove(entity).map(|dense_index| {
+            let dense_index_usize = dense_index.get();
             // SAFE: unique access to ticks
             unsafe {
-                (*self.ticks.get()).swap_remove(dense_index);
+                (*self.ticks.get()).swap_remove(dense_index_usize);
             }
-            self.entities.swap_remove(dense_index);
-            let is_last = dense_index == self.dense.len() - 1;
+            self.entities.swap_remove(dense_index_usize);
+            let is_last = dense_index_usize == self.dense.len() - 1;
             // SAFE: dense_index was just removed from `sparse`, which ensures that it is valid
-            let value = unsafe { self.dense.swap_remove_and_forget_unchecked(dense_index) };
+            let value = unsafe {
+                self.dense
+                    .swap_remove_and_forget_unchecked(dense_index_usize)
+            };
             if !is_last {
-                let swapped_entity = self.entities[dense_index];
+                let swapped_entity = self.entities[dense_index_usize];
                 *self.sparse.get_mut(swapped_entity).unwrap() = dense_index;
             }
             value
@@ -198,13 +205,14 @@ impl ComponentSparseSet {
 
     pub fn remove(&mut self, entity: Entity) -> bool {
         if let Some(dense_index) = self.sparse.remove(entity) {
-            self.ticks.get_mut().swap_remove(dense_index);
-            self.entities.swap_remove(dense_index);
-            let is_last = dense_index == self.dense.len() - 1;
+            let dense_index_usize = dense_index.get();
+            self.ticks.get_mut().swap_remove(dense_index_usize);
+            self.entities.swap_remove(dense_index_usize);
+            let is_last = dense_index_usize == self.dense.len() - 1;
             // SAFE: if the sparse index points to something in the dense vec, it exists
-            unsafe { self.dense.swap_remove_and_drop_unchecked(dense_index) }
+            unsafe { self.dense.swap_remove_and_drop_unchecked(dense_index_usize) }
             if !is_last {
-                let swapped_entity = self.entities[dense_index];
+                let swapped_entity = self.entities[dense_index_usize];
                 *self.sparse.get_mut(swapped_entity).unwrap() = dense_index;
             }
             true
@@ -225,7 +233,7 @@ impl ComponentSparseSet {
 pub struct SparseSet<I, V: 'static> {
     dense: Vec<V>,
     indices: Vec<I>,
-    sparse: SparseArray<I, usize>,
+    sparse: SparseArray<I, NonMaxUsize>,
 }
 
 impl<I: SparseSetIndex, V> Default for SparseSet<I, V> {
@@ -261,10 +269,11 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
         if let Some(dense_index) = self.sparse.get(index.clone()).cloned() {
             // SAFE: dense indices stored in self.sparse always exist
             unsafe {
-                *self.dense.get_unchecked_mut(dense_index) = value;
+                *self.dense.get_unchecked_mut(dense_index.get()) = value;
             }
         } else {
-            self.sparse.insert(index.clone(), self.dense.len());
+            self.sparse
+                .insert(index.clone(), NonMaxUsize::new(self.dense.len()).unwrap());
             self.indices.push(index);
             self.dense.push(value);
         }
@@ -295,11 +304,12 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
     pub fn get_or_insert_with(&mut self, index: I, func: impl FnOnce() -> V) -> &mut V {
         if let Some(dense_index) = self.sparse.get(index.clone()).cloned() {
             // SAFE: dense indices stored in self.sparse always exist
-            unsafe { self.dense.get_unchecked_mut(dense_index) }
+            unsafe { self.dense.get_unchecked_mut(dense_index.get()) }
         } else {
             let value = func();
             let dense_index = self.dense.len();
-            self.sparse.insert(index.clone(), dense_index);
+            self.sparse
+                .insert(index.clone(), NonMaxUsize::new(dense_index).unwrap());
             self.indices.push(index);
             self.dense.push(value);
             // SAFE: dense index was just populated above
@@ -325,25 +335,26 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
     pub fn get(&self, index: I) -> Option<&V> {
         self.sparse.get(index).map(|dense_index| {
             // SAFE: if the sparse index points to something in the dense vec, it exists
-            unsafe { self.dense.get_unchecked(*dense_index) }
+            unsafe { self.dense.get_unchecked(dense_index.get()) }
         })
     }
 
     pub fn get_mut(&mut self, index: I) -> Option<&mut V> {
         let dense = &mut self.dense;
         self.sparse.get(index).map(move |dense_index| {
             // SAFE: if the sparse index points to something in the dense vec, it exists
-            unsafe { dense.get_unchecked_mut(*dense_index) }
+            unsafe { dense.get_unchecked_mut(dense_index.get()) }
         })
     }
 
     pub fn remove(&mut self, index: I) -> Option<V> {
         self.sparse.remove(index).map(|dense_index| {
-            let is_last = dense_index == self.dense.len() - 1;
-            let value = self.dense.swap_remove(dense_index);
-            self.indices.swap_remove(dense_index);
+            let dense_index_usize = dense_index.get();
+            let is_last = dense_index_usize == self.dense.len() - 1;
+            let value = self.dense.swap_remove(dense_index_usize);
+            self.indices.swap_remove(dense_index_usize);
             if !is_last {
-                let swapped_index = self.indices[dense_index].clone();
+                let swapped_index = self.indices[dense_index_usize].clone();
                 *self.sparse.get_mut(swapped_index).unwrap() = dense_index;
             }
             value

crates/bevy_utils/src/lib.rs

@@ -1,5 +1,8 @@
 mod enum_variant_meta;
+mod num;
+
 pub use enum_variant_meta::*;
+pub use num::*;
 
 pub use ahash::AHasher;
 pub use instant::{Duration, Instant};

crates/bevy_utils/src/num.rs

@@ -0,0 +1,85 @@
+macro_rules! impl_non_max_fmt {
+    (($($trait:ident),+) for $ty:ident) => {
+        $(
+            impl std::fmt::$trait for $ty {
+                #[inline]
+                fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                    self.get().fmt(f)
+                }
+            }
+        )+
+    }
+}
+
+macro_rules! impl_non_max {
+    ($nonmax:ident, $nonzero:ty, $repr:ty, $test:ident) => {
+        #[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
+        #[repr(transparent)]
+        pub struct $nonmax($nonzero);
+
+        impl $nonmax {
+            /// Creates a non-max if `n` is not the maximum value.
+            #[inline]
+            #[allow(clippy::manual_map)]
+            pub const fn new(n: $repr) -> Option<Self> {
+                if let Some(n) = <$nonzero>::new(!n) {
+                    Some(Self(n))
+                } else {
+                    None
+                }
+            }
+
+            /// Creates a non-max without checking the value.
+            ///
+            /// # Safety
+            /// `n` must not be equal to T::MAX.
+            #[inline]
+            pub const unsafe fn new_unchecked(n: $repr) -> Self {
+                Self(<$nonzero>::new_unchecked(!n))
+            }
+
+            /// Returns the value as a primitive type.
+            ///
+            /// # Note
+            /// This function is not free. Consider storing the result
+            /// into a variable instead of calling `get()` multiple times.
+            #[inline]
+            pub const fn get(self) -> $repr {
+                !self.0.get()
+            }
+        }
+
+        impl_non_max_fmt! {
+            (Debug, Display, Binary, Octal, LowerHex, UpperHex) for $nonmax
+        }
+
+        #[cfg(test)]
+        mod $test {
+            use super::*;
+
+            #[test]
+            fn test() {
+                assert!($nonmax::new(<$repr>::MAX).is_none());
+                assert_eq!($nonmax::new(0).unwrap().get(), 0);
+                assert_eq!($nonmax::new(1).unwrap().get(), 1);
+
+                // SAFE: `0` != <$repr>::MAX
+                unsafe {
+                    assert_eq!($nonmax::new_unchecked(0).get(), 0);
+                }
+
+                assert_eq!(
+                    std::mem::size_of::<$nonmax>(),
+                    std::mem::size_of::<Option<$nonmax>>()
+                );
+            }
+        }
+    };
+}
+
+impl_non_max!(
+    NonMaxUsize,
+    std::num::NonZeroUsize,
+    usize,
+    non_max_usize_test
+);