Skip to content

Commit 6edd244

Browse files
NathanSWardNathanSWard
committed
NonMaxUsize for SparseArrays
1 parent 883abbb commit 6edd244

File tree

3 files changed

+128
-31
lines changed

3 files changed

+128
-31
lines changed

‎crates/bevy_ecs/src/storage/sparse_set.rs

+42-31
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use crate::{
33
entity::Entity,
44
storage::BlobVec,
55
};
6+
use bevy_utils::NonMaxUsize;
67
use std::{cell::UnsafeCell, marker::PhantomData};
78

89
#[derive(Debug)]
@@ -89,7 +90,7 @@ pub struct ComponentSparseSet {
8990
dense: BlobVec,
9091
ticks: UnsafeCell<Vec<ComponentTicks>>,
9192
entities: Vec<Entity>,
92-
sparse: SparseArray<Entity, usize>,
93+
sparse: SparseArray<Entity, NonMaxUsize>,
9394
}
9495

9596
impl ComponentSparseSet {
@@ -123,11 +124,14 @@ impl ComponentSparseSet {
123124
let dense = &mut self.dense;
124125
let entities = &mut self.entities;
125126
let ticks_list = self.ticks.get_mut();
126-
let dense_index = *self.sparse.get_or_insert_with(entity, move || {
127-
ticks_list.push(ComponentTicks::new(change_tick));
128-
entities.push(entity);
129-
dense.push_uninit()
130-
});
127+
let dense_index = self
128+
.sparse
129+
.get_or_insert_with(entity, move || {
130+
ticks_list.push(ComponentTicks::new(change_tick));
131+
entities.push(entity);
132+
NonMaxUsize::new(dense.push_uninit()).unwrap()
133+
})
134+
.get();
131135
// SAFE: dense_index exists thanks to the call above
132136
self.dense.set_unchecked(dense_index, value);
133137
((*self.ticks.get()).get_unchecked_mut(dense_index)).set_changed(change_tick);
@@ -144,7 +148,7 @@ impl ComponentSparseSet {
144148
pub fn get(&self, entity: Entity) -> Option<*mut u8> {
145149
self.sparse.get(entity).map(|dense_index| {
146150
// SAFE: if the sparse index points to something in the dense vec, it exists
147-
unsafe { self.dense.get_unchecked(*dense_index) }
151+
unsafe { self.dense.get_unchecked(dense_index.get()) }
148152
})
149153
}
150154

@@ -154,7 +158,7 @@ impl ComponentSparseSet {
154158
pub unsafe fn get_with_ticks(&self, entity: Entity) -> Option<(*mut u8, *mut ComponentTicks)> {
155159
let ticks = &mut *self.ticks.get();
156160
self.sparse.get(entity).map(move |dense_index| {
157-
let dense_index = *dense_index;
161+
let dense_index = dense_index.get();
158162
// SAFE: if the sparse index points to something in the dense vec, it exists
159163
(
160164
self.dense.get_unchecked(dense_index),
@@ -169,9 +173,8 @@ impl ComponentSparseSet {
169173
pub unsafe fn get_ticks(&self, entity: Entity) -> Option<&mut ComponentTicks> {
170174
let ticks = &mut *self.ticks.get();
171175
self.sparse.get(entity).map(move |dense_index| {
172-
let dense_index = *dense_index;
173176
// SAFE: if the sparse index points to something in the dense vec, it exists
174-
ticks.get_unchecked_mut(dense_index)
177+
ticks.get_unchecked_mut(dense_index.get())
175178
})
176179
}
177180

@@ -180,16 +183,20 @@ impl ComponentSparseSet {
180183
/// returned).
181184
pub fn remove_and_forget(&mut self, entity: Entity) -> Option<*mut u8> {
182185
self.sparse.remove(entity).map(|dense_index| {
186+
let dense_index_usize = dense_index.get();
183187
// SAFE: unique access to ticks
184188
unsafe {
185-
(*self.ticks.get()).swap_remove(dense_index);
189+
(*self.ticks.get()).swap_remove(dense_index_usize);
186190
}
187-
self.entities.swap_remove(dense_index);
188-
let is_last = dense_index == self.dense.len() - 1;
191+
self.entities.swap_remove(dense_index_usize);
192+
let is_last = dense_index_usize == self.dense.len() - 1;
189193
// SAFE: dense_index was just removed from `sparse`, which ensures that it is valid
190-
let value = unsafe { self.dense.swap_remove_and_forget_unchecked(dense_index) };
194+
let value = unsafe {
195+
self.dense
196+
.swap_remove_and_forget_unchecked(dense_index_usize)
197+
};
191198
if !is_last {
192-
let swapped_entity = self.entities[dense_index];
199+
let swapped_entity = self.entities[dense_index_usize];
193200
*self.sparse.get_mut(swapped_entity).unwrap() = dense_index;
194201
}
195202
value
@@ -198,13 +205,14 @@ impl ComponentSparseSet {
198205

199206
pub fn remove(&mut self, entity: Entity) -> bool {
200207
if let Some(dense_index) = self.sparse.remove(entity) {
201-
self.ticks.get_mut().swap_remove(dense_index);
202-
self.entities.swap_remove(dense_index);
203-
let is_last = dense_index == self.dense.len() - 1;
208+
let dense_index_usize = dense_index.get();
209+
self.ticks.get_mut().swap_remove(dense_index_usize);
210+
self.entities.swap_remove(dense_index_usize);
211+
let is_last = dense_index_usize == self.dense.len() - 1;
204212
// SAFE: if the sparse index points to something in the dense vec, it exists
205-
unsafe { self.dense.swap_remove_and_drop_unchecked(dense_index) }
213+
unsafe { self.dense.swap_remove_and_drop_unchecked(dense_index_usize) }
206214
if !is_last {
207-
let swapped_entity = self.entities[dense_index];
215+
let swapped_entity = self.entities[dense_index_usize];
208216
*self.sparse.get_mut(swapped_entity).unwrap() = dense_index;
209217
}
210218
true
@@ -225,7 +233,7 @@ impl ComponentSparseSet {
225233
pub struct SparseSet<I, V: 'static> {
226234
dense: Vec<V>,
227235
indices: Vec<I>,
228-
sparse: SparseArray<I, usize>,
236+
sparse: SparseArray<I, NonMaxUsize>,
229237
}
230238

231239
impl<I: SparseSetIndex, V> Default for SparseSet<I, V> {
@@ -261,10 +269,11 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
261269
if let Some(dense_index) = self.sparse.get(index.clone()).cloned() {
262270
// SAFE: dense indices stored in self.sparse always exist
263271
unsafe {
264-
*self.dense.get_unchecked_mut(dense_index) = value;
272+
*self.dense.get_unchecked_mut(dense_index.get()) = value;
265273
}
266274
} else {
267-
self.sparse.insert(index.clone(), self.dense.len());
275+
self.sparse
276+
.insert(index.clone(), NonMaxUsize::new(self.dense.len()).unwrap());
268277
self.indices.push(index);
269278
self.dense.push(value);
270279
}
@@ -295,11 +304,12 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
295304
pub fn get_or_insert_with(&mut self, index: I, func: impl FnOnce() -> V) -> &mut V {
296305
if let Some(dense_index) = self.sparse.get(index.clone()).cloned() {
297306
// SAFE: dense indices stored in self.sparse always exist
298-
unsafe { self.dense.get_unchecked_mut(dense_index) }
307+
unsafe { self.dense.get_unchecked_mut(dense_index.get()) }
299308
} else {
300309
let value = func();
301310
let dense_index = self.dense.len();
302-
self.sparse.insert(index.clone(), dense_index);
311+
self.sparse
312+
.insert(index.clone(), NonMaxUsize::new(dense_index).unwrap());
303313
self.indices.push(index);
304314
self.dense.push(value);
305315
// SAFE: dense index was just populated above
@@ -325,25 +335,26 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
325335
pub fn get(&self, index: I) -> Option<&V> {
326336
self.sparse.get(index).map(|dense_index| {
327337
// SAFE: if the sparse index points to something in the dense vec, it exists
328-
unsafe { self.dense.get_unchecked(*dense_index) }
338+
unsafe { self.dense.get_unchecked(dense_index.get()) }
329339
})
330340
}
331341

332342
pub fn get_mut(&mut self, index: I) -> Option<&mut V> {
333343
let dense = &mut self.dense;
334344
self.sparse.get(index).map(move |dense_index| {
335345
// SAFE: if the sparse index points to something in the dense vec, it exists
336-
unsafe { dense.get_unchecked_mut(*dense_index) }
346+
unsafe { dense.get_unchecked_mut(dense_index.get()) }
337347
})
338348
}
339349

340350
pub fn remove(&mut self, index: I) -> Option<V> {
341351
self.sparse.remove(index).map(|dense_index| {
342-
let is_last = dense_index == self.dense.len() - 1;
343-
let value = self.dense.swap_remove(dense_index);
344-
self.indices.swap_remove(dense_index);
352+
let dense_index_usize = dense_index.get();
353+
let is_last = dense_index_usize == self.dense.len() - 1;
354+
let value = self.dense.swap_remove(dense_index_usize);
355+
self.indices.swap_remove(dense_index_usize);
345356
if !is_last {
346-
let swapped_index = self.indices[dense_index].clone();
357+
let swapped_index = self.indices[dense_index_usize].clone();
347358
*self.sparse.get_mut(swapped_index).unwrap() = dense_index;
348359
}
349360
value

‎crates/bevy_utils/src/lib.rs

+3
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
mod enum_variant_meta;
2+
mod num;
3+
24
pub use enum_variant_meta::*;
5+
pub use num::*;
36

47
pub use ahash::AHasher;
58
pub use instant::{Duration, Instant};

‎crates/bevy_utils/src/num.rs

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
macro_rules! impl_non_max_fmt {
2+
(($($trait:ident),+) for $ty:ident) => {
3+
$(
4+
impl std::fmt::$trait for $ty {
5+
#[inline]
6+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
7+
self.get().fmt(f)
8+
}
9+
}
10+
)+
11+
}
12+
}
13+
14+
macro_rules! impl_non_max {
15+
($nonmax:ident, $nonzero:ty, $repr:ty, $test:ident) => {
16+
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
17+
pub struct $nonmax($nonzero);
18+
19+
impl $nonmax {
20+
/// Creates a non-max if `n` is not the maximum value.
21+
#[inline]
22+
pub const fn new(n: $repr) -> Option<Self> {
23+
if let Some(n) = <$nonzero>::new(!n) {
24+
Some(Self(n))
25+
} else {
26+
None
27+
}
28+
}
29+
30+
/// Creates a non-max without checking the value.
31+
///
32+
/// # Safety
33+
/// `n` must not be equal to T::MAX.
34+
#[inline]
35+
pub const unsafe fn new_unchecked(n: $repr) -> Self {
36+
Self(<$nonzero>::new_unchecked(!n))
37+
}
38+
39+
/// Returns the value as a primitive type.
40+
///
41+
/// # Note
42+
/// This function is not free. Consider storing the result
43+
/// into a variable instead of calling `get()` multiple times.
44+
#[inline]
45+
pub const fn get(self) -> $repr {
46+
!self.0.get()
47+
}
48+
}
49+
50+
impl_non_max_fmt! {
51+
(Debug, Display, Binary, Octal, LowerHex, UpperHex) for $nonmax
52+
}
53+
54+
#[cfg(test)]
55+
mod $test {
56+
use super::*;
57+
58+
#[test]
59+
fn test() {
60+
assert!($nonmax::new(<$repr>::MAX).is_none());
61+
assert_eq!($nonmax::new(0).unwrap().get(), 0);
62+
assert_eq!($nonmax::new(1).unwrap().get(), 1);
63+
64+
// SAFE: `0` != <$repr>::MAX
65+
unsafe {
66+
assert_eq!($nonmax::new_unchecked(0).get(), 0);
67+
}
68+
69+
assert_eq!(
70+
std::mem::size_of::<$nonmax>(),
71+
std::mem::size_of::<Option<$nonmax>>()
72+
);
73+
}
74+
}
75+
};
76+
}
77+
78+
impl_non_max!(
79+
NonMaxUsize,
80+
std::num::NonZeroUsize,
81+
usize,
82+
non_max_usize_test
83+
);

0 commit comments

Comments
 (0)