[Merged by Bors] - Extend EntityLocation with TableId and TableRow

crates/bevy_ecs/src/archetype.rs

@@ -15,23 +15,23 @@ use std::{
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
 #[repr(transparent)]
-pub struct ArchetypeId(usize);
+pub struct ArchetypeId(u32);
 
 impl ArchetypeId {
     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 const fn new(index: usize) -> Self {
-        ArchetypeId(index)
+        ArchetypeId(index as u32)
     }
 
     #[inline]
     pub fn index(self) -> usize {
-        self.0
+        self.0 as usize
     }
 }
 
@@ -275,8 +275,8 @@ impl Archetype {
     }
 
     #[inline]
-    pub(crate) fn set_entity_table_row(&mut self, index: usize, table_row: usize) {
-        self.entities[index].table_row = table_row;
+    pub(crate) fn set_entity_table_row(&mut self, index: u32, table_row: usize) {
+        self.entities[index as usize].table_row = table_row;
     }
 
     /// # Safety
@@ -287,7 +287,9 @@ impl Archetype {
 
         EntityLocation {
             archetype_id: self.id,
-            index: self.entities.len() - 1,
+            archetype_index: (self.entities.len() - 1) as u32,
+            table_id: self.table_id,
+            table_row: table_row as u32,
         }
     }
 
@@ -297,14 +299,14 @@ impl Archetype {
 
     /// Removes the entity at `index` by swapping it out. Returns the table row the entity is stored
     /// in.
-    pub(crate) fn swap_remove(&mut self, index: usize) -> ArchetypeSwapRemoveResult {
-        let is_last = index == self.entities.len() - 1;
-        let entity = self.entities.swap_remove(index);
+    pub(crate) fn swap_remove(&mut self, index: u32) -> ArchetypeSwapRemoveResult {
+        let is_last = index as usize == self.entities.len() - 1;
+        let entity = self.entities.swap_remove(index as usize);
         ArchetypeSwapRemoveResult {
             swapped_entity: if is_last {
                 None
             } else {
-                Some(self.entities[index].entity)
+                Some(self.entities[index as usize].entity)
             },
             table_row: entity.table_row,
         }
@@ -491,7 +493,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

@@ -528,13 +528,9 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
     pub unsafe fn insert<T: Bundle>(
         &mut self,
         entity: Entity,
-        archetype_index: usize,
+        location: EntityLocation,
         bundle: T,
     ) -> EntityLocation {
-        let location = EntityLocation {
-            index: archetype_index,
-            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,14 +544,14 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
                     self.sparse_sets,
                     add_bundle,
                     entity,
-                    self.archetype.entity_table_row(archetype_index),
+                    location.table_row as usize,
                     self.change_tick,
                     bundle,
                 );
                 location
             }
             InsertBundleResult::NewArchetypeSameTable { new_archetype } => {
-                let result = self.archetype.swap_remove(location.index);
+                let result = self.archetype.swap_remove(location.archetype_index);
                 if let Some(swapped_entity) = result.swapped_entity {
                     self.entities.meta[swapped_entity.index as usize].location = location;
                 }
@@ -583,7 +579,7 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
                 new_archetype,
                 new_table,
             } => {
-                let result = self.archetype.swap_remove(location.index);
+                let result = self.archetype.swap_remove(location.archetype_index);
                 if let Some(swapped_entity) = result.swapped_entity {
                     self.entities.meta[swapped_entity.index as usize].location = location;
                 }
@@ -611,7 +607,7 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
                     };
 
                     swapped_archetype
-                        .set_entity_table_row(swapped_location.index, result.table_row);
+                        .set_entity_table_row(swapped_location.archetype_index, result.table_row);
                 }
 
                 // PERF: this could be looked up during Inserter construction and stored (but borrowing makes this nasty)

crates/bevy_ecs/src/entity/mod.rs

@@ -34,7 +34,10 @@ mod map_entities;
 
 pub use map_entities::*;
 
-use crate::{archetype::ArchetypeId, storage::SparseSetIndex};
+use crate::{
+    archetype::ArchetypeId,
+    storage::{SparseSetIndex, TableId},
+};
 use serde::{Deserialize, Serialize};
 use std::{convert::TryFrom, fmt, mem, sync::atomic::Ordering};
 
@@ -662,12 +665,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 a given [`Entity`].
 #[derive(Copy, Clone, Debug)]
 #[repr(C)]
 pub struct EntityMeta {
+    /// The current generation of the [`Entity`].
     pub generation: u32,
+    /// The current location of the [`Entity`]
     pub location: EntityLocation,
 }
 
@@ -676,20 +684,40 @@ impl EntityMeta {
         generation: 0,
         location: EntityLocation {
             archetype_id: ArchetypeId::INVALID,
-            index: usize::MAX, // dummy value, to be filled in
+            archetype_index: u32::MAX, // dummy value, to be filled in
+            table_id: TableId::INVALID,
+            table_row: u32::MAX, // 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
-    pub index: usize,
+    /// The index of the [`Entity`] within its [`Archetype`].
+    ///
+    /// [`Archetype`]: crate::archetype::Archetype
+    pub archetype_index: u32,
+
+    /// 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: u32,
 }
 
 #[cfg(test)]

crates/bevy_ecs/src/query/iter.rs

@@ -157,7 +157,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
@@ -176,12 +176,15 @@ where
                 table,
             );
 
-            let table_row = archetype.entity_table_row(location.index);
             // SAFETY: set_archetype was called prior.
             // `location.index` 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 as usize) {
                 // SAFETY: set_archetype was called prior, `location.index` 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 as usize,
+                ));
             }
         }
         None

crates/bevy_ecs/src/query/state.rs

@@ -418,17 +418,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.index);
         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 as usize) {
+            Ok(Q::fetch(&mut fetch, entity, location.table_row as usize))
         } 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]
@@ -585,7 +600,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