From 3f2e22a57b6f0be005a3fc7beb69f007ba0a7f6a Mon Sep 17 00:00:00 2001 From: james7132 Date: Sun, 30 Jan 2022 09:45:30 -0800 Subject: [PATCH 1/7] Page SparseArray --- crates/bevy_ecs/src/storage/sparse_set.rs | 78 ++++++++++++++++------- 1 file changed, 54 insertions(+), 24 deletions(-) diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index 0b4efbe9bf339..56c7b5639df33 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -3,11 +3,13 @@ use crate::{ entity::Entity, storage::BlobVec, }; -use std::{cell::UnsafeCell, marker::PhantomData}; +use std::{cell::UnsafeCell, marker::PhantomData, mem::MaybeUninit}; + +const PAGE_SIZE: usize = 64; #[derive(Debug)] pub struct SparseArray { - values: Vec>, + values: Vec; PAGE_SIZE]>>>, marker: PhantomData, } @@ -28,59 +30,87 @@ impl SparseArray { } impl SparseArray { + fn split_index(index: I) -> (usize, usize) { + let idx = index.sparse_set_index(); + (idx / PAGE_SIZE, idx % PAGE_SIZE) + } + + fn make_page() -> Box<[Option; PAGE_SIZE]> { + // SAFE: The memory is all initialized to None upon return. + unsafe { + let mut page: MaybeUninit<[Option; PAGE_SIZE]> = MaybeUninit::uninit(); + let slice = page.as_mut_ptr(); + for idx in 0..PAGE_SIZE { + (*slice)[idx] = None; + } + // TODO: Initialize with Box::assume_uninit when https://github.com/rust-lang/rust/issues/63291 lands in stable. + Box::new(page.assume_init()) + } + } + 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_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_index(index); + self.values + .get(page) + .and_then(|p| p.as_ref()) + .and_then(|p| p.get(index)) + .map(Option::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_index(index); + self.values + .get(page) + .and_then(|p| p.as_ref()) + .and_then(|p| p[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_index(index); self.values - .get_mut(index) - .map(|v| v.as_mut()) - .unwrap_or(None) + .get_mut(page) + .and_then(|page| page.as_mut()) + .and_then(|page| page[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_index(index); + self.values + .get_mut(page) + .and_then(|page| page.as_mut()) + .and_then(|page| page[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_index(index); + if page > self.values.len() { + self.values.resize_with(page + 1, || None); } - self.values.resize_with(index + 1, || None); - let value = &mut self.values[index]; - *value = Some(func()); - value.as_mut().unwrap() + let page = self.values[page].get_or_insert_with(Self::make_page); + page[index].get_or_insert_with(func) } pub fn clear(&mut self) { From 7ec837c597d1c2c6fdd8bf975afbbedf5b99828f Mon Sep 17 00:00:00 2001 From: james7132 Date: Sun, 30 Jan 2022 10:40:04 -0800 Subject: [PATCH 2/7] Fix tests/examples --- crates/bevy_ecs/src/storage/sparse_set.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index 56c7b5639df33..ca99f10680f5f 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -58,7 +58,7 @@ impl SparseArray { #[inline] pub fn insert(&mut self, index: I, value: V) { let (page, index) = Self::split_index(index); - if page > self.values.len() { + if page >= self.values.len() { self.values.resize_with(page + 1, || None); } let page = self.values[page].get_or_insert_with(Self::make_page); @@ -106,7 +106,7 @@ impl SparseArray { #[inline] pub fn get_or_insert_with(&mut self, index: I, func: impl FnOnce() -> V) -> &mut V { let (page, index) = Self::split_index(index); - if page > self.values.len() { + if page >= self.values.len() { self.values.resize_with(page + 1, || None); } let page = self.values[page].get_or_insert_with(Self::make_page); From 70bd6170b3286df3a56bed575196bba9d180ce43 Mon Sep 17 00:00:00 2001 From: james7132 Date: Sun, 30 Jan 2022 10:56:08 -0800 Subject: [PATCH 3/7] Be more safe --- crates/bevy_ecs/src/storage/sparse_set.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index ca99f10680f5f..7d3104c7aa90e 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -39,9 +39,9 @@ impl SparseArray { // SAFE: The memory is all initialized to None upon return. unsafe { let mut page: MaybeUninit<[Option; PAGE_SIZE]> = MaybeUninit::uninit(); - let slice = page.as_mut_ptr(); - for idx in 0..PAGE_SIZE { - (*slice)[idx] = None; + let array = page.as_mut_ptr().as_mut().unwrap(); + for item in &mut array[..] { + *item = None; } // TODO: Initialize with Box::assume_uninit when https://github.com/rust-lang/rust/issues/63291 lands in stable. Box::new(page.assume_init()) From 93ec5595d35e87fa6c549cf6ea92b72760d54ece Mon Sep 17 00:00:00 2001 From: james7132 Date: Sun, 30 Jan 2022 12:59:16 -0800 Subject: [PATCH 4/7] Fix initialization soundness issues --- crates/bevy_ecs/src/storage/sparse_set.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index 7d3104c7aa90e..d2c30df343b7c 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -36,16 +36,16 @@ impl SparseArray { } 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. - unsafe { - let mut page: MaybeUninit<[Option; PAGE_SIZE]> = MaybeUninit::uninit(); - let array = page.as_mut_ptr().as_mut().unwrap(); - for item in &mut array[..] { - *item = None; - } - // TODO: Initialize with Box::assume_uninit when https://github.com/rust-lang/rust/issues/63291 lands in stable. - Box::new(page.assume_init()) + 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 { From 2dce3ca7de3b0d323b1782208c3b3c8f25ca8e0d Mon Sep 17 00:00:00 2001 From: james7132 Date: Wed, 1 Jun 2022 04:53:36 -0700 Subject: [PATCH 5/7] Add unsafe gets as the intra-page index is always valid --- crates/bevy_ecs/src/storage/sparse_set.rs | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index 75d722eac5d35..3a78850d503e5 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -73,8 +73,8 @@ impl SparseArray { self.values .get(page) .and_then(|p| p.as_ref()) - .and_then(|p| p.get(index)) - .map(Option::is_some) + // 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) } @@ -84,7 +84,8 @@ impl SparseArray { self.values .get(page) .and_then(|p| p.as_ref()) - .and_then(|p| p[index].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] @@ -93,7 +94,8 @@ impl SparseArray { self.values .get_mut(page) .and_then(|page| page.as_mut()) - .and_then(|page| page[index].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] @@ -102,7 +104,8 @@ impl SparseArray { self.values .get_mut(page) .and_then(|page| page.as_mut()) - .and_then(|page| page[index].take()) + // SAFETY: Index is always valid as it must be less than page size. + .and_then(|page| unsafe { page.get_unchecked_mut(index).take() }) } #[inline] @@ -111,8 +114,14 @@ impl SparseArray { if page >= self.values.len() { self.values.resize_with(page + 1, || None); } - let page = self.values[page].get_or_insert_with(Self::make_page); - page[index].get_or_insert_with(func) + // 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) + } } pub fn clear(&mut self) { From dfc574fcfa51dd1e8963b98e87a751119a072151 Mon Sep 17 00:00:00 2001 From: james7132 Date: Wed, 1 Jun 2022 04:54:12 -0700 Subject: [PATCH 6/7] Inline split_index --- crates/bevy_ecs/src/storage/sparse_set.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index 3a78850d503e5..dab4d312ff597 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -32,6 +32,7 @@ impl SparseArray { } impl SparseArray { + #[inline] fn split_index(index: I) -> (usize, usize) { let idx = index.sparse_set_index(); (idx / PAGE_SIZE, idx % PAGE_SIZE) From 5d0f2bded597ad9ac504ebea7c1e56ddf5b3913a Mon Sep 17 00:00:00 2001 From: james7132 Date: Wed, 1 Jun 2022 04:55:09 -0700 Subject: [PATCH 7/7] split_index -> split_at_index --- crates/bevy_ecs/src/storage/sparse_set.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index dab4d312ff597..4dcb110b78751 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -33,7 +33,7 @@ impl SparseArray { impl SparseArray { #[inline] - fn split_index(index: I) -> (usize, usize) { + fn split_at_index(index: I) -> (usize, usize) { let idx = index.sparse_set_index(); (idx / PAGE_SIZE, idx % PAGE_SIZE) } @@ -60,7 +60,7 @@ impl SparseArray { #[inline] pub fn insert(&mut self, index: I, value: V) { - let (page, index) = Self::split_index(index); + let (page, index) = Self::split_at_index(index); if page >= self.values.len() { self.values.resize_with(page + 1, || None); } @@ -70,7 +70,7 @@ impl SparseArray { #[inline] pub fn contains(&self, index: I) -> bool { - let (page, index) = Self::split_index(index); + let (page, index) = Self::split_at_index(index); self.values .get(page) .and_then(|p| p.as_ref()) @@ -81,7 +81,7 @@ impl SparseArray { #[inline] pub fn get(&self, index: I) -> Option<&V> { - let (page, index) = Self::split_index(index); + let (page, index) = Self::split_at_index(index); self.values .get(page) .and_then(|p| p.as_ref()) @@ -91,7 +91,7 @@ impl SparseArray { #[inline] pub fn get_mut(&mut self, index: I) -> Option<&mut V> { - let (page, index) = Self::split_index(index); + let (page, index) = Self::split_at_index(index); self.values .get_mut(page) .and_then(|page| page.as_mut()) @@ -101,7 +101,7 @@ impl SparseArray { #[inline] pub fn remove(&mut self, index: I) -> Option { - let (page, index) = Self::split_index(index); + let (page, index) = Self::split_at_index(index); self.values .get_mut(page) .and_then(|page| page.as_mut()) @@ -111,7 +111,7 @@ impl SparseArray { #[inline] pub fn get_or_insert_with(&mut self, index: I, func: impl FnOnce() -> V) -> &mut V { - let (page, index) = Self::split_index(index); + let (page, index) = Self::split_at_index(index); if page >= self.values.len() { self.values.resize_with(page + 1, || None); }