From e3a66673657f90c0b93156774f3509fc49b5cbba Mon Sep 17 00:00:00 2001 From: NathanW Date: Fri, 7 May 2021 15:52:38 -0600 Subject: [PATCH] NonMaxUsize for SparseArrays --- crates/bevy_ecs/src/storage/sparse_set.rs | 73 +++++++++++--------- crates/bevy_utils/src/lib.rs | 3 + crates/bevy_utils/src/num.rs | 84 +++++++++++++++++++++++ 3 files changed, 129 insertions(+), 31 deletions(-) create mode 100644 crates/bevy_utils/src/num.rs diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index 280c73c06de4f..a30e5cf5a1dc9 100644 --- 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>, entities: Vec, - sparse: SparseArray, + sparse: SparseArray, } 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 { dense: Vec, indices: Vec, - sparse: SparseArray, + sparse: SparseArray, } impl Default for SparseSet { @@ -261,10 +269,11 @@ impl SparseSet { 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 SparseSet { 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,7 +335,7 @@ impl SparseSet { 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()) } }) } @@ -333,17 +343,18 @@ impl SparseSet { 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 { 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 index 9d66179df301c..8a75b00064969 100644 --- 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 new file mode 100644 index 0000000000000..f610a72069756 --- /dev/null +++ b/crates/bevy_utils/src/num.rs @@ -0,0 +1,84 @@ +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)] + 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 { + 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::>() + ); + } + } + }; +} + +impl_non_max!( + NonMaxUsize, + std::num::NonZeroUsize, + usize, + non_max_usize_test +);