Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

NonMaxUsize for SparseArrays #11843

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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());
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This isn't an uncommon branch to take, and this introduces a potential panic, even if we never panic normally, this will have a negative impact on performance, and it's not something we can use debug_checked_unwrap on either since we know for a fact that it's something that we can hit.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does that mean you want me to revert that line?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I honestly am not sure how to deal with this. This branch is unavoidable if we want to avoid unsoundness. If we can find a way to address this without impacting performance, this seems feasible, but otherwise, this is potentially on the hotpath for a lot of operations within the engine, so any tangible performance regression is probably unacceptable. Definitely need to benchmark this in some way to validate.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the potential performance issue here? Is it xor operation or unwrap()?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The best way to deal with this is to improve Rust language (there're a lot of discussions like this). But I'm afraid we're not getting it in the next ten years. :(

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the potential performance issue here? Is it xor operation or unwrap()?

The unwrap, the XOR has persistently shown to be easily pipelinable in any current CPUs without adding too much here, but the extra branch and codegen from panicking has shown to have a signficant impact when working with code that is this hot.

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 {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is a nonmax crate that already implements this. We should try using that first.

(($($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
);