diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index 74e99d4853502..4dcb110b78751 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -4,13 +4,14 @@ use crate::{ storage::BlobVec, }; use bevy_ptr::{OwningPtr, Ptr}; -use std::{cell::UnsafeCell, hash::Hash, marker::PhantomData}; +use std::{cell::UnsafeCell, hash::Hash, marker::PhantomData, mem::MaybeUninit}; +const PAGE_SIZE: usize = 64; type EntityId = u32; #[derive(Debug)] pub struct SparseArray { - values: Vec>, + values: Vec; PAGE_SIZE]>>>, marker: PhantomData, } @@ -31,59 +32,97 @@ impl SparseArray { } impl SparseArray { + #[inline] + fn split_at_index(index: I) -> (usize, usize) { + let idx = index.sparse_set_index(); + (idx / PAGE_SIZE, idx % PAGE_SIZE) + } + + fn make_page() -> Box<[Option; PAGE_SIZE]> { + // TODO: Initialize with Box::assume_uninit when https://github.com/rust-lang/rust/issues/63291 lands in stable. + // SAFE: The memory is all initialized to None upon return. + let mut page: Box<[MaybeUninit>; PAGE_SIZE]> = + Box::new(unsafe { MaybeUninit::uninit().assume_init() }); + // SAFE: Nothing in the following code can panic or be dropped leaving any uninitialized state. + for item in &mut page[..] { + item.write(None); + } + // SAFE: This transmutation has the same ABI as the uninitialized boxed page. + unsafe { std::mem::transmute(page) } + } + pub fn with_capacity(capacity: usize) -> Self { Self { - values: Vec::with_capacity(capacity), + values: Vec::with_capacity(capacity / PAGE_SIZE), marker: PhantomData, } } #[inline] pub fn insert(&mut self, index: I, value: V) { - let index = index.sparse_set_index(); - if index >= self.values.len() { - self.values.resize_with(index + 1, || None); + let (page, index) = Self::split_at_index(index); + if page >= self.values.len() { + self.values.resize_with(page + 1, || None); } - self.values[index] = Some(value); + let page = self.values[page].get_or_insert_with(Self::make_page); + page[index] = Some(value); } #[inline] pub fn contains(&self, index: I) -> bool { - let index = index.sparse_set_index(); - self.values.get(index).map(|v| v.is_some()).unwrap_or(false) + let (page, index) = Self::split_at_index(index); + self.values + .get(page) + .and_then(|p| p.as_ref()) + // SAFETY: Index is always valid as it must be less than page size. + .map(|p| unsafe { p.get_unchecked(index).is_some() }) + .unwrap_or(false) } #[inline] pub fn get(&self, index: I) -> Option<&V> { - let index = index.sparse_set_index(); - self.values.get(index).map(|v| v.as_ref()).unwrap_or(None) + let (page, index) = Self::split_at_index(index); + self.values + .get(page) + .and_then(|p| p.as_ref()) + // SAFETY: Index is always valid as it must be less than page size. + .and_then(|p| unsafe { p.get_unchecked(index).as_ref() }) } #[inline] pub fn get_mut(&mut self, index: I) -> Option<&mut V> { - let index = index.sparse_set_index(); + let (page, index) = Self::split_at_index(index); self.values - .get_mut(index) - .map(|v| v.as_mut()) - .unwrap_or(None) + .get_mut(page) + .and_then(|page| page.as_mut()) + // SAFETY: Index is always valid as it must be less than page size. + .and_then(|page| unsafe { page.get_unchecked_mut(index).as_mut() }) } #[inline] pub fn remove(&mut self, index: I) -> Option { - let index = index.sparse_set_index(); - self.values.get_mut(index).and_then(|value| value.take()) + let (page, index) = Self::split_at_index(index); + self.values + .get_mut(page) + .and_then(|page| page.as_mut()) + // SAFETY: Index is always valid as it must be less than page size. + .and_then(|page| unsafe { page.get_unchecked_mut(index).take() }) } #[inline] pub fn get_or_insert_with(&mut self, index: I, func: impl FnOnce() -> V) -> &mut V { - let index = index.sparse_set_index(); - if index < self.values.len() { - return self.values[index].get_or_insert_with(func); + let (page, index) = Self::split_at_index(index); + if page >= self.values.len() { + self.values.resize_with(page + 1, || None); + } + // SAFETY: Page and index both must after the resize. + unsafe { + let page = self + .values + .get_unchecked_mut(page) + .get_or_insert_with(Self::make_page); + page.get_unchecked_mut(index).get_or_insert_with(func) } - self.values.resize_with(index + 1, || None); - let value = &mut self.values[index]; - *value = Some(func()); - value.as_mut().unwrap() } pub fn clear(&mut self) {