Optimize memory size of SparseArray and integer-based ECS IDs

crates/bevy_ecs/Cargo.toml

@@ -24,6 +24,7 @@ async-channel = "1.4"
 fixedbitset = "0.4"
 fxhash = "0.2"
 downcast-rs = "1.2"
+nonmax = "0.5"
 serde = "1"
 
 [dev-dependencies]

crates/bevy_ecs/src/archetype.rs

@@ -7,28 +7,43 @@ use crate::{
     entity::{Entity, EntityLocation},
     storage::{Column, SparseArray, SparseSet, SparseSetIndex, TableId},
 };
+use nonmax::NonMaxUsize;
 use std::{
     collections::HashMap,
     hash::Hash,
     ops::{Index, IndexMut},
 };
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
-pub struct ArchetypeId(usize);
+pub struct ArchetypeId(NonMaxUsize);
 
 impl ArchetypeId {
-    pub const EMPTY: ArchetypeId = ArchetypeId(0);
-    pub const RESOURCE: ArchetypeId = ArchetypeId(1);
-    pub const INVALID: ArchetypeId = ArchetypeId(usize::MAX);
+    pub const EMPTY: ArchetypeId = unsafe { ArchetypeId::new_unchecked(0) };
+    pub const RESOURCE: ArchetypeId = unsafe { ArchetypeId::new_unchecked(1) };
 
+    /// Creates a new [`ArchetypeId`].
+    ///
+    /// # Panics
+    /// This will panic if `index` is `usize::MAX`.
     #[inline]
-    pub const fn new(index: usize) -> Self {
-        ArchetypeId(index)
+    pub fn new(index: usize) -> Self {
+        Self(NonMaxUsize::new(index).expect("ArchetypeID cannot be created from usize::MAX"))
     }
 
+    /// Creates a new [`ArchetypeId`] without checking for validity of the
+    /// provided index.
+    ///
+    /// # Safety
+    /// `index` must not be `usize::MAX`
     #[inline]
-    pub fn index(self) -> usize {
-        self.0
+    pub const unsafe fn new_unchecked(index: usize) -> Self {
+        Self(NonMaxUsize::new_unchecked(index))
+    }
+
+    /// Gets the corresponding index for the ID.
+    #[inline]
+    pub const fn index(self) -> usize {
+        self.0.get()
     }
 }
 
@@ -335,28 +350,42 @@ pub struct ArchetypeIdentity {
 }
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
-pub struct ArchetypeComponentId(usize);
+pub struct ArchetypeComponentId(NonMaxUsize);
 
 impl ArchetypeComponentId {
+    /// Creates a new [`ArchetypeComponentId`] from an index without
+    /// checking for the type's invariants.
+    ///
+    /// # Safety
+    /// `index` must not be [`usize::MAX`].
     #[inline]
-    pub const fn new(index: usize) -> Self {
-        Self(index)
+    pub const unsafe fn new_unchecked(index: usize) -> Self {
+        Self(NonMaxUsize::new_unchecked(index))
+    }
+
+    /// Creates a new [`ArchetypeComponentId`] from an index.
+    ///
+    /// # Panic
+    /// This function will panic if `index` is equal to [`usize::MAX`].
+    #[inline]
+    pub fn new(index: usize) -> Self {
+        Self(NonMaxUsize::new(index).unwrap())
     }
 
     #[inline]
     pub fn index(self) -> usize {
-        self.0
+        self.0.get()
     }
 }
 
 impl SparseSetIndex for ArchetypeComponentId {
     #[inline]
     fn sparse_set_index(&self) -> usize {
-        self.0
+        self.index()
     }
 
     fn get_sparse_set_index(value: usize) -> Self {
-        Self(value)
+        Self::new(value)
     }
 }
 
@@ -486,15 +515,15 @@ impl Archetypes {
         let archetypes = &mut self.archetypes;
         let archetype_component_count = &mut self.archetype_component_count;
         let mut next_archetype_component_id = move || {
-            let id = ArchetypeComponentId(*archetype_component_count);
+            let id = ArchetypeComponentId::new(*archetype_component_count);
             *archetype_component_count += 1;
             id
         };
         *self
             .archetype_ids
             .entry(archetype_identity)
             .or_insert_with(move || {
-                let id = ArchetypeId(archetypes.len());
+                let id = ArchetypeId::new(archetypes.len());
                 let table_archetype_components = (0..table_components.len())
                     .map(|_| next_archetype_component_id())
                     .collect();
@@ -540,3 +569,16 @@ impl IndexMut<ArchetypeId> for Archetypes {
         &mut self.archetypes[index.index()]
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::ArchetypeComponentId;
+
+    #[test]
+    pub fn test_archetyype_component_id_size_optimized() {
+        assert_eq!(
+            core::mem::size_of::<ArchetypeComponentId>(),
+            core::mem::size_of::<Option<ArchetypeComponentId>>(),
+        );
+    }
+}

crates/bevy_ecs/src/bundle.rs

@@ -12,6 +12,7 @@ use crate::{
 };
 use bevy_ecs_macros::all_tuples;
 use bevy_ptr::OwningPtr;
+use nonmax::NonMaxUsize;
 use std::{any::TypeId, collections::HashMap};
 
 /// An ordered collection of [`Component`]s.
@@ -130,12 +131,31 @@ macro_rules! tuple_impl {
 all_tuples!(tuple_impl, 0, 15, C);
 
 #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
-pub struct BundleId(usize);
+pub struct BundleId(NonMaxUsize);
 
 impl BundleId {
+    /// Creates a new [`BundleId`] from an index without checking for the
+    /// type's invariants.
+    ///
+    /// # Safety
+    /// `value` must not be [`usize::MAX`].
+    #[inline]
+    pub const unsafe fn new_unchecked(value: usize) -> Self {
+        Self(NonMaxUsize::new_unchecked(value))
+    }
+
+    /// Creates a new [`BundleId`] from an index.
+    ///
+    /// # Panic
+    /// This function will panic if `value` is equal to [`usize::MAX`].
+    #[inline]
+    pub fn new(value: usize) -> Self {
+        Self(NonMaxUsize::new(value).unwrap())
+    }
+
     #[inline]
     pub fn index(self) -> usize {
-        self.0
+        self.0.get()
     }
 }
 
@@ -146,7 +166,7 @@ impl SparseSetIndex for BundleId {
     }
 
     fn get_sparse_set_index(value: usize) -> Self {
-        Self(value)
+        Self::new(value)
     }
 }
 
@@ -445,10 +465,10 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
             InsertBundleResult::NewArchetypeSameTable { new_archetype } => {
                 let result = self.archetype.swap_remove(location.index);
                 if let Some(swapped_entity) = result.swapped_entity {
-                    self.entities.meta[swapped_entity.id as usize].location = location;
+                    self.entities.meta[swapped_entity.id as usize].location = Some(location);
                 }
                 let new_location = new_archetype.allocate(entity, result.table_row);
-                self.entities.meta[entity.id as usize].location = new_location;
+                self.entities.meta[entity.id as usize].location = Some(new_location);
 
                 // PERF: this could be looked up during Inserter construction and stored (but borrowing makes this nasty)
                 let add_bundle = self
@@ -473,15 +493,15 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
             } => {
                 let result = self.archetype.swap_remove(location.index);
                 if let Some(swapped_entity) = result.swapped_entity {
-                    self.entities.meta[swapped_entity.id as usize].location = location;
+                    self.entities.meta[swapped_entity.id as usize].location = Some(location);
                 }
                 // PERF: store "non bundle" components in edge, then just move those to avoid
                 // redundant copies
                 let move_result = self
                     .table
                     .move_to_superset_unchecked(result.table_row, *new_table);
                 let new_location = new_archetype.allocate(entity, move_result.new_row);
-                self.entities.meta[entity.id as usize].location = new_location;
+                self.entities.meta[entity.id as usize].location = Some(new_location);
 
                 // if an entity was moved into this entity's table spot, update its table row
                 if let Some(swapped_entity) = move_result.swapped_entity {
@@ -557,7 +577,7 @@ impl<'a, 'b> BundleSpawner<'a, 'b> {
             self.change_tick,
             bundle,
         );
-        self.entities.meta[entity.id as usize].location = location;
+        self.entities.meta[entity.id as usize].location = Some(location);
 
         location
     }
@@ -598,7 +618,7 @@ impl Bundles {
         let bundle_infos = &mut self.bundle_infos;
         let id = self.bundle_ids.entry(TypeId::of::<T>()).or_insert_with(|| {
             let component_ids = T::component_ids(components, storages);
-            let id = BundleId(bundle_infos.len());
+            let id = BundleId::new(bundle_infos.len());
             // SAFE: T::component_id ensures info was created
             let bundle_info = unsafe {
                 initialize_bundle(std::any::type_name::<T>(), component_ids, id, components)
@@ -607,7 +627,7 @@ impl Bundles {
             id
         });
         // SAFE: index either exists, or was initialized
-        unsafe { self.bundle_infos.get_unchecked(id.0) }
+        unsafe { self.bundle_infos.get_unchecked(id.index()) }
     }
 }
 
@@ -643,3 +663,16 @@ unsafe fn initialize_bundle(
         storage_types,
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::BundleId;
+
+    #[test]
+    pub fn test_bundle_id_size_optimized() {
+        assert_eq!(
+            core::mem::size_of::<BundleId>(),
+            core::mem::size_of::<Option<BundleId>>(),
+        );
+    }
+}