Skip to content

Commit

Permalink
NonMaxUsize for SparseArrays
Browse files Browse the repository at this point in the history
  • Loading branch information
NathanSWard authored and tygyh committed Feb 13, 2024
1 parent 173bb48 commit 4adb1ae
Show file tree
Hide file tree
Showing 3 changed files with 131 additions and 32 deletions.
75 changes: 43 additions & 32 deletions crates/bevy_ecs/src/storage/sparse_set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use crate::{
entity::Entity,
storage::BlobVec,
};
use bevy_utils::NonMaxUsize;
use std::{cell::UnsafeCell, marker::PhantomData};

#[derive(Debug)]
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -123,14 +124,17 @@ 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);
(*self.ticks.get()).get_unchecked_mut(dense_index).set_changed(change_tick);
}

#[inline]
Expand All @@ -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()) }
})
}

Expand All @@ -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),
Expand All @@ -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())
})
}

Expand All @@ -180,13 +183,17 @@ impl ComponentSparseSet {
/// returned).
pub fn remove_and_forget(&mut self, entity: Entity) -> Option<*mut u8> {
self.sparse.remove(entity).map(|dense_index| {
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: 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
Expand All @@ -195,13 +202,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
Expand All @@ -222,7 +230,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> {
Expand Down Expand Up @@ -258,10 +266,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);
}
Expand Down Expand Up @@ -292,11 +301,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
Expand All @@ -322,25 +332,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
Expand Down
3 changes: 3 additions & 0 deletions crates/bevy_utils/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down
85 changes: 85 additions & 0 deletions crates/bevy_utils/src/num.rs
Original file line number Diff line number Diff line change
@@ -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
);

0 comments on commit 4adb1ae

Please sign in to comment.