Immutable sparse sets for metadata storage (bevyengine#4928)

james7132 · cart · ItsDoot · commit ced100fc4109 · 2023-01-31T20:00:08.000-06:00
# Objective Make core types in ECS smaller. The column sparse set in Tables is never updated after creation. ## Solution Create `ImmutableSparseSet` which removes the capacity fields in the backing vec's and the APIs for inserting or removing elements. Drops the size of the sparse set by 3 usizes (24 bytes on 64-bit systems) ## Followup ~~After bevyengine#4809, Archetype's component SparseSet should be replaced with it.~~ This has been done. --- ## Changelog Removed: `Table::component_capacity` ## Migration Guide `Table::component_capacity()` has been removed as Tables do not support adding/removing columns after construction. Co-authored-by: Carter Anderson <mcanders1@gmail.com>
diff --git a/crates/bevy_ecs/src/archetype.rs b/crates/bevy_ecs/src/archetype.rs
@@ -5,7 +5,7 @@ use crate::{
     bundle::BundleId,
     component::{ComponentId, StorageType},
     entity::{Entity, EntityLocation},
-    storage::{SparseArray, SparseSet, SparseSetIndex, TableId},
+    storage::{ImmutableSparseSet, SparseArray, SparseSet, SparseSetIndex, TableId},
 };
 use std::{
     collections::HashMap,
@@ -182,7 +182,7 @@ pub struct Archetype {
     table_id: TableId,
     edges: Edges,
     entities: Vec<ArchetypeEntity>,
-    components: SparseSet<ComponentId, ArchetypeComponentInfo>,
+    components: ImmutableSparseSet<ComponentId, ArchetypeComponentInfo>,
 }
 
 impl Archetype {
@@ -217,8 +217,8 @@ impl Archetype {
         Self {
             id,
             table_id,
-            components,
-            entities: Default::default(),
+            entities: Vec::new(),
+            components: components.into_immutable(),
             edges: Default::default(),
         }
     }
diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs
@@ -14,6 +14,14 @@ pub(crate) struct SparseArray<I, V = I> {
     marker: PhantomData<I>,
 }
 
+/// A space-optimized version of [`SparseArray`] that cannot be changed
+/// after construction.
+#[derive(Debug)]
+pub(crate) struct ImmutableSparseArray<I, V = I> {
+    values: Box<[Option<V>]>,
+    marker: PhantomData<I>,
+}
+
 impl<I: SparseSetIndex, V> Default for SparseArray<I, V> {
     fn default() -> Self {
         Self::new()
@@ -30,6 +38,27 @@ impl<I, V> SparseArray<I, V> {
     }
 }
 
+macro_rules! impl_sparse_array {
+    ($ty:ident) => {
+        impl<I: SparseSetIndex, V> $ty<I, V> {
+            #[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)
+            }
+
+            #[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)
+            }
+        }
+    };
+}
+
+impl_sparse_array!(SparseArray);
+impl_sparse_array!(ImmutableSparseArray);
+
 impl<I: SparseSetIndex, V> SparseArray<I, V> {
     #[inline]
     pub fn insert(&mut self, index: I, value: V) {
@@ -40,18 +69,6 @@ impl<I: SparseSetIndex, V> SparseArray<I, V> {
         self.values[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)
-    }
-
-    #[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)
-    }
-
     #[inline]
     pub fn get_mut(&mut self, index: I) -> Option<&mut V> {
         let index = index.sparse_set_index();
@@ -70,6 +87,13 @@ impl<I: SparseSetIndex, V> SparseArray<I, V> {
     pub fn clear(&mut self) {
         self.values.clear();
     }
+
+    pub(crate) fn into_immutable(self) -> ImmutableSparseArray<I, V> {
+        ImmutableSparseArray {
+            values: self.values.into_boxed_slice(),
+            marker: PhantomData,
+        }
+    }
 }
 
 /// A sparse data structure of [Components](crate::component::Component)
@@ -249,11 +273,75 @@ pub struct SparseSet<I, V: 'static> {
     sparse: SparseArray<I, usize>,
 }
 
+/// A space-optimized version of [`SparseSet`] that cannot be changed
+/// after construction.
+#[derive(Debug)]
+pub(crate) struct ImmutableSparseSet<I, V: 'static> {
+    dense: Box<[V]>,
+    indices: Box<[I]>,
+    sparse: ImmutableSparseArray<I, usize>,
+}
+
+macro_rules! impl_sparse_set {
+    ($ty:ident) => {
+        impl<I: SparseSetIndex, V> $ty<I, V> {
+            #[inline]
+            pub fn len(&self) -> usize {
+                self.dense.len()
+            }
+
+            #[inline]
+            pub fn contains(&self, index: I) -> bool {
+                self.sparse.contains(index)
+            }
+
+            pub fn get(&self, index: I) -> Option<&V> {
+                self.sparse.get(index).map(|dense_index| {
+                    // SAFETY: if the sparse index points to something in the dense vec, it exists
+                    unsafe { self.dense.get_unchecked(*dense_index) }
+                })
+            }
+
+            pub fn get_mut(&mut self, index: I) -> Option<&mut V> {
+                let dense = &mut self.dense;
+                self.sparse.get(index).map(move |dense_index| {
+                    // SAFETY: if the sparse index points to something in the dense vec, it exists
+                    unsafe { dense.get_unchecked_mut(*dense_index) }
+                })
+            }
+
+            pub fn indices(&self) -> impl Iterator<Item = I> + '_ {
+                self.indices.iter().cloned()
+            }
+
+            pub fn values(&self) -> impl Iterator<Item = &V> {
+                self.dense.iter()
+            }
+
+            pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
+                self.dense.iter_mut()
+            }
+
+            pub fn iter(&self) -> impl Iterator<Item = (&I, &V)> {
+                self.indices.iter().zip(self.dense.iter())
+            }
+
+            pub fn iter_mut(&mut self) -> impl Iterator<Item = (&I, &mut V)> {
+                self.indices.iter().zip(self.dense.iter_mut())
+            }
+        }
+    };
+}
+
+impl_sparse_set!(SparseSet);
+impl_sparse_set!(ImmutableSparseSet);
+
 impl<I: SparseSetIndex, V> Default for SparseSet<I, V> {
     fn default() -> Self {
         Self::new()
     }
 }
+
 impl<I, V> SparseSet<I, V> {
     pub const fn new() -> Self {
         Self {
@@ -306,36 +394,11 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
         }
     }
 
-    #[inline]
-    pub fn len(&self) -> usize {
-        self.dense.len()
-    }
-
     #[inline]
     pub fn is_empty(&self) -> bool {
         self.dense.len() == 0
     }
 
-    #[inline]
-    pub fn contains(&self, index: I) -> bool {
-        self.sparse.contains(index)
-    }
-
-    pub fn get(&self, index: I) -> Option<&V> {
-        self.sparse.get(index).map(|dense_index| {
-            // SAFETY: if the sparse index points to something in the dense vec, it exists
-            unsafe { self.dense.get_unchecked(*dense_index) }
-        })
-    }
-
-    pub fn get_mut(&mut self, index: I) -> Option<&mut V> {
-        let dense = &mut self.dense;
-        self.sparse.get(index).map(move |dense_index| {
-            // SAFETY: if the sparse index points to something in the dense vec, it exists
-            unsafe { dense.get_unchecked_mut(*dense_index) }
-        })
-    }
-
     pub fn remove(&mut self, index: I) -> Option<V> {
         self.sparse.remove(index).map(|dense_index| {
             let is_last = dense_index == self.dense.len() - 1;
@@ -349,24 +412,12 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
         })
     }
 
-    pub fn indices(&self) -> impl Iterator<Item = I> + '_ {
-        self.indices.iter().cloned()
-    }
-
-    pub fn values(&self) -> impl Iterator<Item = &V> {
-        self.dense.iter()
-    }
-
-    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
-        self.dense.iter_mut()
-    }
-
-    pub fn iter(&self) -> impl Iterator<Item = (&I, &V)> {
-        self.indices.iter().zip(self.dense.iter())
-    }
-
-    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&I, &mut V)> {
-        self.indices.iter().zip(self.dense.iter_mut())
+    pub(crate) fn into_immutable(self) -> ImmutableSparseSet<I, V> {
+        ImmutableSparseSet {
+            dense: self.dense.into_boxed_slice(),
+            indices: self.indices.into_boxed_slice(),
+            sparse: self.sparse.into_immutable(),
+        }
     }
 }
 
diff --git a/crates/bevy_ecs/src/storage/table.rs b/crates/bevy_ecs/src/storage/table.rs
@@ -2,7 +2,7 @@ use crate::{
     component::{ComponentId, ComponentInfo, ComponentTicks, Components},
     entity::Entity,
     query::DebugCheckedUnwrap,
-    storage::{blob_vec::BlobVec, SparseSet},
+    storage::{blob_vec::BlobVec, ImmutableSparseSet, SparseSet},
 };
 use bevy_ptr::{OwningPtr, Ptr, PtrMut};
 use bevy_utils::HashMap;
@@ -262,31 +262,68 @@ impl Column {
     }
 }
 
-pub struct Table {
+/// A builder type for constructing [`Table`]s.
+///
+///  - Use [`with_capacity`] to initialize the builder.
+///  - Repeatedly call [`add_column`] to add columns for components.
+///  - Finalize with [`build`] to get the constructed [`Table`].
+///
+/// [`with_capacity`]: Self::with_capacity
+/// [`add_column`]: Self::add_column
+/// [`build`]: Self::build
+pub(crate) struct TableBuilder {
     columns: SparseSet<ComponentId, Column>,
-    entities: Vec<Entity>,
+    capacity: usize,
 }
 
-impl Table {
-    pub(crate) fn with_capacity(capacity: usize, column_capacity: usize) -> Table {
+impl TableBuilder {
+    /// Creates a blank [`Table`], allocating space for `column_capacity` columns
+    /// with the capacity to hold `capacity` entities worth of components each.
+    pub fn with_capacity(capacity: usize, column_capacity: usize) -> Self {
         Self {
             columns: SparseSet::with_capacity(column_capacity),
-            entities: Vec::with_capacity(capacity),
+            capacity,
         }
     }
 
-    #[inline]
-    pub fn entities(&self) -> &[Entity] {
-        &self.entities
-    }
-
-    pub(crate) fn add_column(&mut self, component_info: &ComponentInfo) {
+    pub fn add_column(&mut self, component_info: &ComponentInfo) {
         self.columns.insert(
             component_info.id(),
-            Column::with_capacity(component_info, self.entities.capacity()),
+            Column::with_capacity(component_info, self.capacity),
         );
     }
 
+    pub fn build(self) -> Table {
+        Table {
+            columns: self.columns.into_immutable(),
+            entities: Vec::with_capacity(self.capacity),
+        }
+    }
+}
+
+/// A column-oriented [structure-of-arrays] based storage for [`Component`]s of entities
+/// in a [`World`].
+///
+/// Conceptually, a `Table` can be thought of as an `HashMap<ComponentId, Column>`, where
+/// each `Column` is a type-erased `Vec<T: Component>`. Each row corresponds to a single entity
+/// (i.e. index 3 in Column A and index 3 in Column B point to different components on the same
+/// entity). Fetching components from a table involves fetching the associated column for a
+/// component type (via it's [`ComponentId`]), then fetching the entity's row within that column.
+///
+/// [structure-of-arrays]: https://en.wikipedia.org/wiki/AoS_and_SoA#Structure_of_arrays
+/// [`Component`]: crate::component::Component
+/// [`World`]: crate::world::World
+pub struct Table {
+    columns: ImmutableSparseSet<ComponentId, Column>,
+    entities: Vec<Entity>,
+}
+
+impl Table {
+    #[inline]
+    pub fn entities(&self) -> &[Entity] {
+        &self.entities
+    }
+
     /// Removes the entity at the given row and returns the entity swapped in to replace it (if an
     /// entity was swapped in)
     ///
@@ -457,11 +494,6 @@ impl Table {
         self.entities.capacity()
     }
 
-    #[inline]
-    pub fn component_capacity(&self) -> usize {
-        self.columns.capacity()
-    }
-
     #[inline]
     pub fn is_empty(&self) -> bool {
         self.entities.is_empty()
@@ -495,7 +527,7 @@ pub struct Tables {
 
 impl Default for Tables {
     fn default() -> Self {
-        let empty_table = Table::with_capacity(0, 0);
+        let empty_table = TableBuilder::with_capacity(0, 0).build();
         Tables {
             tables: vec![empty_table],
             table_ids: HashMap::default(),
@@ -548,11 +580,11 @@ impl Tables {
             .raw_entry_mut()
             .from_key(component_ids)
             .or_insert_with(|| {
-                let mut table = Table::with_capacity(0, component_ids.len());
+                let mut table = TableBuilder::with_capacity(0, component_ids.len());
                 for component_id in component_ids {
                     table.add_column(components.get_info_unchecked(*component_id));
                 }
-                tables.push(table);
+                tables.push(table.build());
                 (component_ids.to_vec(), TableId(tables.len() - 1))
             });
 
@@ -601,7 +633,7 @@ mod tests {
     use crate::{
         component::{ComponentTicks, Components},
         entity::Entity,
-        storage::Table,
+        storage::TableBuilder,
     };
     #[derive(Component)]
     struct W<T>(T);
@@ -612,8 +644,9 @@ mod tests {
         let mut storages = Storages::default();
         let component_id = components.init_component::<W<usize>>(&mut storages);
         let columns = &[component_id];
-        let mut table = Table::with_capacity(0, columns.len());
-        table.add_column(components.get_info(component_id).unwrap());
+        let mut builder = TableBuilder::with_capacity(0, columns.len());
+        builder.add_column(components.get_info(component_id).unwrap());
+        let mut table = builder.build();
         let entities = (0..200).map(Entity::from_raw).collect::<Vec<_>>();
         for entity in &entities {
             // SAFETY: we allocate and immediately set data afterwards
