Implement paged SparseArray

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<I, V = I> {
-    values: Vec<Option<V>>,
+    values: Vec<Option<Box<[Option<V>; PAGE_SIZE]>>>,
     marker: PhantomData<I>,
 }
 
@@ -31,59 +32,97 @@ impl<I, V> SparseArray<I, V> {
 }
 
 impl<I: SparseSetIndex, V> SparseArray<I, V> {
+    #[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<V>; 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<Option<V>>; 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<V> {
-        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) {