NonMaxUsize for SparseArrays

NathanSWard · May 7, 2021 · 5876fae · 5876fae
1 parent 883abbb
commit 5876fae
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 31 deletions.
diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/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

diff --git a/crates/bevy_utils/src/lib.rs b/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};

diff --git a/crates/bevy_utils/src/num.rs b/crates/bevy_utils/src/num.rs
@@ -0,0 +1,55 @@
+use std::num::NonZeroUsize;
+
+#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
+#[repr(transparent)]
+pub struct NonMaxUsize(NonZeroUsize);
+
+impl std::fmt::Debug for NonMaxUsize {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("NonMaxUsize").field(&self.0.get()).finish()
+    }
+}
+
+impl NonMaxUsize {
+    #[inline]
+    pub const fn new(n: usize) -> Option<Self> {
+        if let Some(n) = NonZeroUsize::new(n ^ usize::MAX) {
+            Some(Self(n))
+        } else {
+            None
+        }
+    }
+
+    /// # Safety
+    /// `n` must not be equal to [usize]::MAX.
+    #[inline]
+    pub const unsafe fn new_unchecked(n: usize) -> Self {
+        Self(NonZeroUsize::new_unchecked(n ^ usize::MAX))
+    }
+
+    #[inline(always)]
+    pub const fn get(self) -> usize {
+        self.0.get() ^ usize::MAX
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn non_max_usize() {
+        assert!(NonMaxUsize::new(usize::MAX).is_none());
+        assert_eq!(NonMaxUsize::new(0).unwrap().get(), 0);
+
+        // SAFE: `0` != usize::MAX
+        unsafe {
+            assert_eq!(NonMaxUsize::new_unchecked(0).get(), 0);
+        }
+
+        assert_eq!(
+            std::mem::size_of::<NonMaxUsize>(),
+            std::mem::size_of::<Option<NonMaxUsize>>()
+        );
+    }
+}