[bevy_<util/ecs>] Better SparseArray performance

NathanSWard · NathanSWard · commit da0235ae2698 · 2021-05-04T20:44:57.000-06:00
Problem:
- bevy_ecs::SparseArray uses an `Option&lt;T&gt;` internally to represent is the
  value is valid or not.
- This causes extra branching as well as extra memory overhead due to
  Option's discriminatior.

Solution:
- Implement a `PackedOption` class which uses a value of `T`
  to represent the `None` state.
- Have SparseArray use `PackedOption` internally over `Option`.
diff --git a/crates/bevy_ecs/src/archetype.rs b/crates/bevy_ecs/src/archetype.rs
@@ -4,6 +4,7 @@ use crate::{
     entity::{Entity, EntityLocation},
     storage::{Column, SparseArray, SparseSet, SparseSetIndex, TableId},
 };
+use bevy_utils::{NoneValue, PackedOption};
 use std::{
     borrow::Cow,
     collections::HashMap,
@@ -25,6 +26,11 @@ impl ArchetypeId {
         ArchetypeId(0)
     }
 
+    #[inline]
+    pub const fn invalid() -> ArchetypeId {
+        ArchetypeId(usize::MAX)
+    }
+
     #[inline]
     pub const fn resource() -> ArchetypeId {
         ArchetypeId(1)
@@ -46,11 +52,30 @@ pub struct AddBundle {
     pub bundle_status: Vec<ComponentStatus>,
 }
 
+impl NoneValue for ArchetypeId {
+    const NONE_VALUE: Self = ArchetypeId::invalid();
+
+    fn is_none_value(&self) -> bool {
+        *self == Self::NONE_VALUE
+    }
+}
+
+impl NoneValue for AddBundle {
+    const NONE_VALUE: Self = Self {
+        archetype_id: ArchetypeId::NONE_VALUE,
+        bundle_status: Vec::new(),
+    };
+
+    fn is_none_value(&self) -> bool {
+        self.archetype_id.is_none_value()
+    }
+}
+
 #[derive(Default)]
 pub struct Edges {
     pub add_bundle: SparseArray<BundleId, AddBundle>,
-    pub remove_bundle: SparseArray<BundleId, Option<ArchetypeId>>,
-    pub remove_bundle_intersection: SparseArray<BundleId, Option<ArchetypeId>>,
+    pub remove_bundle: SparseArray<BundleId, ArchetypeId>,
+    pub remove_bundle_intersection: SparseArray<BundleId, ArchetypeId>,
 }
 
 impl Edges {
@@ -76,28 +101,32 @@ impl Edges {
     }
 
     #[inline]
-    pub fn get_remove_bundle(&self, bundle_id: BundleId) -> Option<Option<ArchetypeId>> {
-        self.remove_bundle.get(bundle_id).cloned()
+    pub fn get_remove_bundle(&self, bundle_id: BundleId) -> PackedOption<ArchetypeId> {
+        self.remove_bundle.get(bundle_id).cloned().into()
     }
 
     #[inline]
-    pub fn set_remove_bundle(&mut self, bundle_id: BundleId, archetype_id: Option<ArchetypeId>) {
+    pub fn set_remove_bundle(
+        &mut self,
+        bundle_id: BundleId,
+        archetype_id: PackedOption<ArchetypeId>,
+    ) {
         self.remove_bundle.insert(bundle_id, archetype_id);
     }
 
     #[inline]
-    pub fn get_remove_bundle_intersection(
-        &self,
-        bundle_id: BundleId,
-    ) -> Option<Option<ArchetypeId>> {
-        self.remove_bundle_intersection.get(bundle_id).cloned()
+    pub fn get_remove_bundle_intersection(&self, bundle_id: BundleId) -> PackedOption<ArchetypeId> {
+        self.remove_bundle_intersection
+            .get(bundle_id)
+            .copied()
+            .into()
     }
 
     #[inline]
     pub fn set_remove_bundle_intersection(
         &mut self,
         bundle_id: BundleId,
-        archetype_id: Option<ArchetypeId>,
+        archetype_id: PackedOption<ArchetypeId>,
     ) {
         self.remove_bundle_intersection
             .insert(bundle_id, archetype_id);
@@ -119,6 +148,17 @@ pub(crate) struct ArchetypeComponentInfo {
     pub(crate) archetype_component_id: ArchetypeComponentId,
 }
 
+impl NoneValue for ArchetypeComponentInfo {
+    const NONE_VALUE: Self = Self {
+        storage_type: StorageType::Table,
+        archetype_component_id: ArchetypeComponentId::NONE_VALUE,
+    };
+
+    fn is_none_value(&self) -> bool {
+        self.archetype_component_id.is_none_value()
+    }
+}
+
 pub struct Archetype {
     id: ArchetypeId,
     entities: Vec<Entity>,
@@ -337,6 +377,11 @@ pub struct ArchetypeIdentity {
 pub struct ArchetypeComponentId(usize);
 
 impl ArchetypeComponentId {
+    #[inline]
+    pub const fn invalid() -> Self {
+        Self(usize::MAX)
+    }
+
     #[inline]
     pub const fn new(index: usize) -> Self {
         Self(index)
@@ -359,6 +404,14 @@ impl SparseSetIndex for ArchetypeComponentId {
     }
 }
 
+impl NoneValue for ArchetypeComponentId {
+    const NONE_VALUE: Self = Self::invalid();
+
+    fn is_none_value(&self) -> bool {
+        *self == Self::invalid()
+    }
+}
+
 pub struct Archetypes {
     pub(crate) archetypes: Vec<Archetype>,
     pub(crate) archetype_component_count: usize,
diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs
@@ -3,15 +3,27 @@ use crate::{
     entity::Entity,
     storage::BlobVec,
 };
+use bevy_utils::{NoneValue, PackedOption};
 use std::{cell::UnsafeCell, marker::PhantomData};
 
-#[derive(Debug)]
 pub struct SparseArray<I, V = I> {
-    values: Vec<Option<V>>,
+    values: Vec<PackedOption<V>>,
     marker: PhantomData<I>,
 }
 
-impl<I: SparseSetIndex, V> Default for SparseArray<I, V> {
+impl<I, V> std::fmt::Debug for SparseArray<I, V>
+where
+    V: NoneValue + std::fmt::Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SparseArray")
+            .field("values", &self.values)
+            .field("marker", &self.marker)
+            .finish()
+    }
+}
+
+impl<I: SparseSetIndex, V: NoneValue> Default for SparseArray<I, V> {
     fn default() -> Self {
         Self::new()
     }
@@ -27,7 +39,7 @@ impl<I, V> SparseArray<I, V> {
     }
 }
 
-impl<I: SparseSetIndex, V> SparseArray<I, V> {
+impl<I: SparseSetIndex, V: NoneValue> SparseArray<I, V> {
     pub fn with_capacity(capacity: usize) -> Self {
         Self {
             values: Vec::with_capacity(capacity),
@@ -36,39 +48,48 @@ impl<I: SparseSetIndex, V> SparseArray<I, V> {
     }
 
     #[inline]
-    pub fn insert(&mut self, index: I, value: V) {
+    pub fn insert<T: Into<PackedOption<V>>>(&mut self, index: I, value: T) {
         let index = index.sparse_set_index();
         if index >= self.values.len() {
-            self.values.resize_with(index + 1, || None);
+            self.values.resize_with(index + 1, || None.into());
         }
-        self.values[index] = Some(value);
+        self.values[index] = value.into();
     }
 
     #[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)
+        self.values
+            .get(index)
+            .map(PackedOption::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)
+        self.values
+            .get(index)
+            .map(PackedOption::as_ref)
+            .unwrap_or(None)
     }
 
     #[inline]
     pub fn get_mut(&mut self, index: I) -> Option<&mut V> {
         let index = index.sparse_set_index();
         self.values
             .get_mut(index)
-            .map(|v| v.as_mut())
+            .map(PackedOption::as_mut)
             .unwrap_or(None)
     }
 
     #[inline]
-    pub fn remove(&mut self, index: I) -> Option<V> {
+    pub fn remove(&mut self, index: I) -> PackedOption<V> {
         let index = index.sparse_set_index();
-        self.values.get_mut(index).and_then(|value| value.take())
+        self.values
+            .get_mut(index)
+            .map(PackedOption::take)
+            .unwrap_or_default()
     }
 
     #[inline]
@@ -77,9 +98,9 @@ impl<I: SparseSetIndex, V> SparseArray<I, V> {
         if index < self.values.len() {
             return self.values[index].get_or_insert_with(func);
         }
-        self.values.resize_with(index + 1, || None);
+        self.values.resize_with(index + 1, || None.into());
         let value = &mut self.values[index];
-        *value = Some(func());
+        *value = func().into();
         value.as_mut().unwrap()
     }
 }
@@ -179,7 +200,7 @@ impl ComponentSparseSet {
     /// it exists). It is the caller's responsibility to drop the returned ptr (if Some is
     /// returned).
     pub fn remove_and_forget(&mut self, entity: Entity) -> Option<*mut u8> {
-        self.sparse.remove(entity).map(|dense_index| {
+        self.sparse.remove(entity).expand().map(|dense_index| {
             // SAFE: unique access to ticks
             unsafe {
                 (*self.ticks.get()).swap_remove(dense_index);
@@ -197,7 +218,7 @@ impl ComponentSparseSet {
     }
 
     pub fn remove(&mut self, entity: Entity) -> bool {
-        if let Some(dense_index) = self.sparse.remove(entity) {
+        if let Some(dense_index) = self.sparse.remove(entity).expand() {
             self.ticks.get_mut().swap_remove(dense_index);
             self.entities.swap_remove(dense_index);
             let is_last = dense_index == self.dense.len() - 1;
@@ -233,6 +254,7 @@ impl<I: SparseSetIndex, V> Default for SparseSet<I, V> {
         Self::new()
     }
 }
+
 impl<I, V> SparseSet<I, V> {
     pub const fn new() -> Self {
         Self {
@@ -338,7 +360,7 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
     }
 
     pub fn remove(&mut self, index: I) -> Option<V> {
-        self.sparse.remove(index).map(|dense_index| {
+        self.sparse.remove(index).expand().map(|dense_index| {
             let is_last = dense_index == self.dense.len() - 1;
             let value = self.dense.swap_remove(dense_index);
             self.indices.swap_remove(dense_index);
diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs
@@ -6,6 +6,7 @@ use crate::{
     storage::{SparseSet, Storages},
     world::{Mut, World},
 };
+use bevy_utils::PackedOption;
 use std::any::TypeId;
 
 pub struct EntityRef<'w> {
@@ -309,7 +310,8 @@ impl<'w> EntityMut<'w> {
                 old_location.archetype_id,
                 bundle_info,
                 false,
-            )?
+            )
+            .expand()?
         };
 
         if new_archetype_id == old_location.archetype_id {
@@ -767,7 +769,7 @@ unsafe fn remove_bundle_from_archetype(
     archetype_id: ArchetypeId,
     bundle_info: &BundleInfo,
     intersection: bool,
-) -> Option<ArchetypeId> {
+) -> PackedOption<ArchetypeId> {
     // check the archetype graph to see if the Bundle has been removed from this archetype in the
     // past
     let remove_bundle_result = {
@@ -780,9 +782,10 @@ unsafe fn remove_bundle_from_archetype(
             current_archetype.edges().get_remove_bundle(bundle_info.id)
         }
     };
-    let result = if let Some(result) = remove_bundle_result {
+
+    let result = if remove_bundle_result.is_some() {
         // this Bundle removal result is cached. just return that!
-        result
+        remove_bundle_result
     } else {
         let mut next_table_components;
         let mut next_sparse_set_components;
@@ -805,8 +808,8 @@ unsafe fn remove_bundle_from_archetype(
                     // graph
                     current_archetype
                         .edges_mut()
-                        .set_remove_bundle(bundle_info.id, None);
-                    return None;
+                        .set_remove_bundle(bundle_info.id, None.into());
+                    return None.into();
                 }
             }
 
@@ -837,7 +840,7 @@ unsafe fn remove_bundle_from_archetype(
             next_table_components,
             next_sparse_set_components,
         );
-        Some(new_archetype_id)
+        new_archetype_id.into()
     };
     let current_archetype = &mut archetypes[archetype_id];
     // cache the result in an edge
diff --git a/crates/bevy_utils/src/lib.rs b/crates/bevy_utils/src/lib.rs
@@ -1,5 +1,8 @@
 mod enum_variant_meta;
+mod packed_option;
+
 pub use enum_variant_meta::*;
+pub use packed_option::*;
 
 pub use ahash::AHasher;
 pub use instant::{Duration, Instant};
diff --git a/crates/bevy_utils/src/packed_option.rs b/crates/bevy_utils/src/packed_option.rs
