Extend EntityLocation with TableId and TableRow (#6681)

# Objective
`Query::get` and other random access methods require looking up `EntityLocation` for every provided entity, then always looking up the `Archetype` to get the table ID and table row. This requires 4 total random fetches from memory: the `Entities` lookup, the `Archetype` lookup, the table row lookup, and the final fetch from table/sparse sets. If `EntityLocation` contains the table ID and table row, only the `Entities` lookup and the final storage fetch are required.

## Solution
Add `TableId` and table row to `EntityLocation`. Ensure it's updated whenever entities are moved around. To ensure `EntityMeta` does not grow bigger, both `TableId` and `ArchetypeId` have been shrunk to u32, and the archetype index and table row are stored as u32s instead of as usizes. This should shrink `EntityMeta` by 4 bytes, from 24 to 20 bytes, as there is no padding anymore due to the change in alignment.

This idea was partially concocted by @BoxyUwU. 

## Performance
This should restore the `Query::get` "gains" lost to #6625 that were introduced in #4800 without being unsound, and also incorporates some of the memory usage reductions seen in #3678.

This also removes the same lookups during add/remove/spawn commands, so there may be a bit of a speedup in commands and `Entity{Ref,Mut}`.

---

## Changelog
Added: `EntityLocation::table_id`
Added: `EntityLocation::table_row`.
Changed: `World`s can now only hold a maximum of 2<sup>32</sup>- 1 archetypes.
Changed: `World`s can now only hold a maximum of 2<sup>32</sup> - 1 tables.

## Migration Guide

A `World` can only hold a maximum of 2<sup>32</sup> - 1 archetypes and tables now. If your use case requires more than this, please file an issue explaining your use case.

Author: james7132

crates/bevy_ecs/src/archetype.rs

@@ -41,20 +41,20 @@ use std::{
 /// [`Entities::get`]: crate::entity::Entities
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 #[repr(transparent)]
-pub struct ArchetypeRow(usize);
+pub struct ArchetypeRow(u32);
 
 impl ArchetypeRow {
-    pub const INVALID: ArchetypeRow = ArchetypeRow(usize::MAX);
+    pub const INVALID: ArchetypeRow = ArchetypeRow(u32::MAX);
 
     /// Creates a `ArchetypeRow`.
     pub const fn new(index: usize) -> Self {
-        Self(index)
+        Self(index as u32)
     }
 
     /// Gets the index of the row.
     #[inline]
     pub const fn index(self) -> usize {
-        self.0
+        self.0 as usize
     }
 }
 
@@ -69,24 +69,24 @@ impl ArchetypeRow {
 /// [`EMPTY`]: crate::archetype::ArchetypeId::EMPTY
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
 #[repr(transparent)]
-pub struct ArchetypeId(usize);
+pub struct ArchetypeId(u32);
 
 impl ArchetypeId {
     /// The ID for the [`Archetype`] without any components.
     pub const EMPTY: ArchetypeId = ArchetypeId(0);
     /// # Safety:
     ///
     /// This must always have an all-1s bit pattern to ensure soundness in fast entity id space allocation.
-    pub const INVALID: ArchetypeId = ArchetypeId(usize::MAX);
+    pub const INVALID: ArchetypeId = ArchetypeId(u32::MAX);
 
     #[inline]
     pub(crate) const fn new(index: usize) -> Self {
-        ArchetypeId(index)
+        ArchetypeId(index as u32)
     }
 
     #[inline]
     pub(crate) fn index(self) -> usize {
-        self.0
+        self.0 as usize
     }
 }
 
@@ -407,18 +407,18 @@ impl Archetype {
     /// Fetches the row in the [`Table`] where the components for the entity at `index`
     /// is stored.
     ///
-    /// An entity's archetype index can be fetched from [`EntityLocation::archetype_row`], which
+    /// An entity's archetype row  can be fetched from [`EntityLocation::archetype_row`], which
     /// can be retrieved from [`Entities::get`].
     ///
     /// # Panics
     /// This function will panic if `index >= self.len()`.
     ///
     /// [`Table`]: crate::storage::Table
-    /// [`EntityLocation`]: crate::entity::EntityLocation::archetype_row
+    /// [`EntityLocation::archetype_row`]: crate::entity::EntityLocation::archetype_row
     /// [`Entities::get`]: crate::entity::Entities::get
     #[inline]
-    pub fn entity_table_row(&self, index: ArchetypeRow) -> TableRow {
-        self.entities[index.0].table_row
+    pub fn entity_table_row(&self, row: ArchetypeRow) -> TableRow {
+        self.entities[row.index()].table_row
     }
 
     /// Updates if the components for the entity at `index` can be found
@@ -427,8 +427,8 @@ impl Archetype {
     /// # Panics
     /// This function will panic if `index >= self.len()`.
     #[inline]
-    pub(crate) fn set_entity_table_row(&mut self, index: ArchetypeRow, table_row: TableRow) {
-        self.entities[index.0].table_row = table_row;
+    pub(crate) fn set_entity_table_row(&mut self, row: ArchetypeRow, table_row: TableRow) {
+        self.entities[row.index()].table_row = table_row;
     }
 
     /// Allocates an entity to the archetype.
@@ -441,11 +441,14 @@ impl Archetype {
         entity: Entity,
         table_row: TableRow,
     ) -> EntityLocation {
+        let archetype_row = ArchetypeRow::new(self.entities.len());
         self.entities.push(ArchetypeEntity { entity, table_row });
 
         EntityLocation {
             archetype_id: self.id,
-            archetype_row: ArchetypeRow(self.entities.len() - 1),
+            archetype_row,
+            table_id: self.table_id,
+            table_row,
         }
     }
 
@@ -458,14 +461,14 @@ impl Archetype {
     ///
     /// # Panics
     /// This function will panic if `index >= self.len()`
-    pub(crate) fn swap_remove(&mut self, index: ArchetypeRow) -> ArchetypeSwapRemoveResult {
-        let is_last = index.0 == self.entities.len() - 1;
-        let entity = self.entities.swap_remove(index.0);
+    pub(crate) fn swap_remove(&mut self, row: ArchetypeRow) -> ArchetypeSwapRemoveResult {
+        let is_last = row.index() == self.entities.len() - 1;
+        let entity = self.entities.swap_remove(row.index());
         ArchetypeSwapRemoveResult {
             swapped_entity: if is_last {
                 None
             } else {
-                Some(self.entities[index.0].entity)
+                Some(self.entities[row.index()].entity)
             },
             table_row: entity.table_row,
         }
@@ -691,7 +694,7 @@ impl Archetypes {
             .archetype_ids
             .entry(archetype_identity)
             .or_insert_with(move || {
-                let id = ArchetypeId(archetypes.len());
+                let id = ArchetypeId::new(archetypes.len());
                 let table_start = *archetype_component_count;
                 *archetype_component_count += table_components.len();
                 let table_archetype_components =

crates/bevy_ecs/src/bundle.rs

@@ -6,7 +6,7 @@ pub use bevy_ecs_macros::Bundle;
 
 use crate::{
     archetype::{
-        Archetype, ArchetypeId, ArchetypeRow, Archetypes, BundleComponentStatus, ComponentStatus,
+        Archetype, ArchetypeId, Archetypes, BundleComponentStatus, ComponentStatus,
         SpawnBundleStatus,
     },
     component::{Component, ComponentId, ComponentStorage, Components, StorageType, Tick},
@@ -528,13 +528,9 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
     pub unsafe fn insert<T: Bundle>(
         &mut self,
         entity: Entity,
-        archetype_row: ArchetypeRow,
+        location: EntityLocation,
         bundle: T,
     ) -> EntityLocation {
-        let location = EntityLocation {
-            archetype_row,
-            archetype_id: self.archetype.id(),
-        };
         match &mut self.result {
             InsertBundleResult::SameArchetype => {
                 // PERF: this could be looked up during Inserter construction and stored (but borrowing makes this nasty)
@@ -548,7 +544,7 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
                     self.sparse_sets,
                     add_bundle,
                     entity,
-                    self.archetype.entity_table_row(archetype_row),
+                    location.table_row,
                     self.change_tick,
                     bundle,
                 );

crates/bevy_ecs/src/entity/mod.rs

@@ -36,7 +36,7 @@ pub use map_entities::*;
 
 use crate::{
     archetype::{ArchetypeId, ArchetypeRow},
-    storage::SparseSetIndex,
+    storage::{SparseSetIndex, TableId, TableRow},
 };
 use serde::{Deserialize, Serialize};
 use std::{convert::TryFrom, fmt, mem, sync::atomic::Ordering};
@@ -716,12 +716,17 @@ impl Entities {
     }
 }
 
+// This type is repr(C) to ensure that the layout and values within it can be safe to fully fill
+// with u8::MAX, as required by [`Entities::flush_and_reserve_invalid_assuming_no_entities`].
 // Safety:
 // This type must not contain any pointers at any level, and be safe to fully fill with u8::MAX.
+/// Metadata for an [`Entity`].
 #[derive(Copy, Clone, Debug)]
 #[repr(C)]
 struct EntityMeta {
+    /// The current generation of the [`Entity`].
     pub generation: u32,
+    /// The current location of the [`Entity`]
     pub location: EntityLocation,
 }
 
@@ -731,19 +736,39 @@ impl EntityMeta {
         location: EntityLocation {
             archetype_id: ArchetypeId::INVALID,
             archetype_row: ArchetypeRow::INVALID, // dummy value, to be filled in
+            table_id: TableId::INVALID,
+            table_row: TableRow::INVALID, // dummy value, to be filled in
         },
     };
 }
 
+// This type is repr(C) to ensure that the layout and values within it can be safe to fully fill
+// with u8::MAX, as required by [`Entities::flush_and_reserve_invalid_assuming_no_entities`].
+// SAFETY:
+// This type must not contain any pointers at any level, and be safe to fully fill with u8::MAX.
 /// A location of an entity in an archetype.
 #[derive(Copy, Clone, Debug)]
 #[repr(C)]
 pub struct EntityLocation {
-    /// The archetype index
+    /// The ID of the [`Archetype`] the [`Entity`] belongs to.
+    ///
+    /// [`Archetype`]: crate::archetype::Archetype
     pub archetype_id: ArchetypeId,
 
-    /// The index of the entity in the archetype
+    /// The index of the [`Entity`] within its [`Archetype`].
+    ///
+    /// [`Archetype`]: crate::archetype::Archetype
     pub archetype_row: ArchetypeRow,
+
+    /// The ID of the [`Table`] the [`Entity`] belongs to.
+    ///
+    /// [`Table`]: crate::storage::Table
+    pub table_id: TableId,
+
+    /// The index of the [`Entity`] within its [`Table`].
+    ///
+    /// [`Table`]: crate::storage::Table
+    pub table_row: TableRow,
 }
 
 #[cfg(test)]

crates/bevy_ecs/src/query/iter.rs

@@ -151,7 +151,7 @@ where
                 .archetypes
                 .get(location.archetype_id)
                 .debug_checked_unwrap();
-            let table = self.tables.get(archetype.table_id()).debug_checked_unwrap();
+            let table = self.tables.get(location.table_id).debug_checked_unwrap();
 
             // SAFETY: `archetype` is from the world that `fetch/filter` were created for,
             // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
@@ -170,12 +170,11 @@ where
                 table,
             );
 
-            let table_row = archetype.entity_table_row(location.archetype_row);
             // SAFETY: set_archetype was called prior.
             // `location.archetype_row` is an archetype index row in range of the current archetype, because if it was not, the match above would have `continue`d
-            if F::filter_fetch(&mut self.filter, entity, table_row) {
+            if F::filter_fetch(&mut self.filter, entity, location.table_row) {
                 // SAFETY: set_archetype was called prior, `location.archetype_row` is an archetype index in range of the current archetype
-                return Some(Q::fetch(&mut self.fetch, entity, table_row));
+                return Some(Q::fetch(&mut self.fetch, entity, location.table_row));
             }
         }
         None

crates/bevy_ecs/src/query/state.rs

@@ -408,17 +408,16 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
         let mut fetch = Q::init_fetch(world, &self.fetch_state, last_change_tick, change_tick);
         let mut filter = F::init_fetch(world, &self.filter_state, last_change_tick, change_tick);
 
-        let table_row = archetype.entity_table_row(location.archetype_row);
         let table = world
             .storages()
             .tables
-            .get(archetype.table_id())
+            .get(location.table_id)
             .debug_checked_unwrap();
         Q::set_archetype(&mut fetch, &self.fetch_state, archetype, table);
         F::set_archetype(&mut filter, &self.filter_state, archetype, table);
 
-        if F::filter_fetch(&mut filter, entity, table_row) {
-            Ok(Q::fetch(&mut fetch, entity, table_row))
+        if F::filter_fetch(&mut filter, entity, location.table_row) {
+            Ok(Q::fetch(&mut fetch, entity, location.table_row))
         } else {
             Err(QueryEntityError::QueryDoesNotMatch(entity))
         }

crates/bevy_ecs/src/storage/table.rs

@@ -12,18 +12,33 @@ use std::{
     ops::{Index, IndexMut},
 };
 
+/// An opaque unique ID for a [`Table`] within a [`World`].
+///
+/// Can be used with [`Tables::get`] to fetch the corresponding
+/// table.
+///
+/// Each [`Archetype`] always points to a table via [`Archetype::table_id`].
+/// Multiple archetypes can point to the same table so long as the components
+/// stored in the table are identical, but do not share the same sparse set
+/// components.
+///
+/// [`World`]: crate::world::World
+/// [`Archetype`]: crate::archetype::Archetype
+/// [`Archetype::table_id`]: crate::archetype::Archetype::table_id
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct TableId(usize);
+pub struct TableId(u32);
 
 impl TableId {
+    pub(crate) const INVALID: TableId = TableId(u32::MAX);
+
     #[inline]
     pub fn new(index: usize) -> Self {
-        TableId(index)
+        TableId(index as u32)
     }
 
     #[inline]
     pub fn index(self) -> usize {
-        self.0
+        self.0 as usize
     }
 
     #[inline]
@@ -49,19 +64,21 @@ impl TableId {
 /// [`Archetype::table_id`]: crate::archetype::Archetype::table_id
 /// [`Entity`]: crate::entity::Entity
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct TableRow(usize);
+pub struct TableRow(u32);
 
 impl TableRow {
+    pub const INVALID: TableRow = TableRow(u32::MAX);
+
     /// Creates a `TableRow`.
     #[inline]
     pub const fn new(index: usize) -> Self {
-        Self(index)
+        Self(index as u32)
     }
 
     /// Gets the index of the row.
     #[inline]
     pub const fn index(self) -> usize {
-        self.0
+        self.0 as usize
     }
 }
 
@@ -568,7 +585,7 @@ impl Table {
             column.added_ticks.push(UnsafeCell::new(Tick::new(0)));
             column.changed_ticks.push(UnsafeCell::new(Tick::new(0)));
         }
-        TableRow(index)
+        TableRow::new(index)
     }
 
     #[inline]
@@ -677,7 +694,7 @@ impl Tables {
                     table.add_column(components.get_info_unchecked(*component_id));
                 }
                 tables.push(table.build());
-                (component_ids.to_vec(), TableId(tables.len() - 1))
+                (component_ids.to_vec(), TableId::new(tables.len() - 1))
             });
 
         *value