Skip to content

Commit 5876fae

Browse files
NathanSWardNathanSWard
committed
NonMaxUsize for SparseArrays
1 parent 883abbb commit 5876fae

File tree

3 files changed

+100
-31
lines changed

3 files changed

+100
-31
lines changed

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

Lines changed: 42 additions & 31 deletions
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‎

Lines changed: 3 additions & 0 deletions
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‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
use std::num::NonZeroUsize;
2+
3+
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
4+
#[repr(transparent)]
5+
pub struct NonMaxUsize(NonZeroUsize);
6+
7+
impl std::fmt::Debug for NonMaxUsize {
8+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
9+
f.debug_tuple("NonMaxUsize").field(&self.0.get()).finish()
10+
}
11+
}
12+
13+
impl NonMaxUsize {
14+
#[inline]
15+
pub const fn new(n: usize) -> Option<Self> {
16+
if let Some(n) = NonZeroUsize::new(n ^ usize::MAX) {
17+
Some(Self(n))
18+
} else {
19+
None
20+
}
21+
}
22+
23+
/// # Safety
24+
/// `n` must not be equal to [usize]::MAX.
25+
#[inline]
26+
pub const unsafe fn new_unchecked(n: usize) -> Self {
27+
Self(NonZeroUsize::new_unchecked(n ^ usize::MAX))
28+
}
29+
30+
#[inline(always)]
31+
pub const fn get(self) -> usize {
32+
self.0.get() ^ usize::MAX
33+
}
34+
}
35+
36+
#[cfg(test)]
37+
mod test {
38+
use super::*;
39+
40+
#[test]
41+
fn non_max_usize() {
42+
assert!(NonMaxUsize::new(usize::MAX).is_none());
43+
assert_eq!(NonMaxUsize::new(0).unwrap().get(), 0);
44+
45+
// SAFE: `0` != usize::MAX
46+
unsafe {
47+
assert_eq!(NonMaxUsize::new_unchecked(0).get(), 0);
48+
}
49+
50+
assert_eq!(
51+
std::mem::size_of::<NonMaxUsize>(),
52+
std::mem::size_of::<Option<NonMaxUsize>>()
53+
);
54+
}
55+
}

0 commit comments

Comments
 (0)