Newtype ArchetypeRow and TableRow (#4878)

# Objective
Prevent future unsoundness that was seen in #6623.

## Solution
Newtype both indexes in `Archetype` and `Table` as `ArchetypeRow` and `TableRow`. This avoids weird numerical manipulation on the indices, and can be stored and treated opaquely. Also enforces the source and destination of where these indices at a type level.

---

## Changelog
Changed: `Archetype` indices and `Table` rows have been newtyped as `ArchetypeRow` and `TableRow`.

Author: james7132

crates/bevy_ecs/macros/src/fetch.rs

@@ -283,7 +283,7 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
                 unsafe fn fetch<'__w>(
                     _fetch: &mut <Self as #path::query::WorldQuery>::Fetch<'__w>,
                     _entity: #path::entity::Entity,
-                    _table_row: usize
+                    _table_row: #path::storage::TableRow,
                 ) -> <Self as #path::query::WorldQuery>::Item<'__w> {
                     Self::Item {
                         #(#field_idents: <#field_types>::fetch(&mut _fetch.#field_idents, _entity, _table_row),)*
@@ -296,7 +296,7 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
                 unsafe fn filter_fetch<'__w>(
                     _fetch: &mut <Self as #path::query::WorldQuery>::Fetch<'__w>,
                     _entity: #path::entity::Entity,
-                    _table_row: usize
+                    _table_row: #path::storage::TableRow,
                 ) -> bool {
                     true #(&& <#field_types>::filter_fetch(&mut _fetch.#field_idents, _entity, _table_row))*
                 }

crates/bevy_ecs/src/archetype.rs

@@ -23,14 +23,41 @@ use crate::{
     bundle::BundleId,
     component::{ComponentId, StorageType},
     entity::{Entity, EntityLocation},
-    storage::{ImmutableSparseSet, SparseArray, SparseSet, SparseSetIndex, TableId},
+    storage::{ImmutableSparseSet, SparseArray, SparseSet, SparseSetIndex, TableId, TableRow},
 };
 use std::{
     collections::HashMap,
     hash::Hash,
     ops::{Index, IndexMut},
 };
 
+/// An opaque location within a [`Archetype`].
+///
+/// This can be used in conjunction with [`ArchetypeId`] to find the exact location
+/// of an [`Entity`] within a [`World`]. An entity's archetype and index can be
+/// retrieved via [`Entities::get`].
+///
+/// [`World`]: crate::world::World
+/// [`Entities::get`]: crate::entity::Entities
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+#[repr(transparent)]
+pub struct ArchetypeRow(usize);
+
+impl ArchetypeRow {
+    pub const INVALID: ArchetypeRow = ArchetypeRow(usize::MAX);
+
+    /// Creates a `ArchetypeRow`.
+    pub const fn new(index: usize) -> Self {
+        Self(index)
+    }
+
+    /// Gets the index of the row.
+    #[inline]
+    pub const fn index(self) -> usize {
+        self.0
+    }
+}
+
 /// An opaque unique ID for a single [`Archetype`] within a [`World`].
 ///
 /// Archetype IDs are only valid for a given World, and are not globally unique.
@@ -226,7 +253,7 @@ impl Edges {
 /// Metadata about an [`Entity`] in a [`Archetype`].
 pub struct ArchetypeEntity {
     entity: Entity,
-    table_row: usize,
+    table_row: TableRow,
 }
 
 impl ArchetypeEntity {
@@ -240,14 +267,14 @@ impl ArchetypeEntity {
     ///
     /// [`Table`]: crate::storage::Table
     #[inline]
-    pub const fn table_row(&self) -> usize {
+    pub const fn table_row(&self) -> TableRow {
         self.table_row
     }
 }
 
 pub(crate) struct ArchetypeSwapRemoveResult {
     pub(crate) swapped_entity: Option<Entity>,
-    pub(crate) table_row: usize,
+    pub(crate) table_row: TableRow,
 }
 
 /// Internal metadata for a [`Component`] within a given [`Archetype`].
@@ -380,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::index`], which
+    /// An entity's archetype index 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::index
+    /// [`EntityLocation`]: crate::entity::EntityLocation::archetype_row
     /// [`Entities::get`]: crate::entity::Entities::get
     #[inline]
-    pub fn entity_table_row(&self, index: usize) -> usize {
-        self.entities[index].table_row
+    pub fn entity_table_row(&self, index: ArchetypeRow) -> TableRow {
+        self.entities[index.0].table_row
     }
 
     /// Updates if the components for the entity at `index` can be found
@@ -400,21 +427,25 @@ impl Archetype {
     /// # Panics
     /// This function will panic if `index >= self.len()`.
     #[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: ArchetypeRow, table_row: TableRow) {
+        self.entities[index.0].table_row = table_row;
     }
 
     /// Allocates an entity to the archetype.
     ///
     /// # Safety
     /// valid component values must be immediately written to the relevant storages
     /// `table_row` must be valid
-    pub(crate) unsafe fn allocate(&mut self, entity: Entity, table_row: usize) -> EntityLocation {
+    pub(crate) unsafe fn allocate(
+        &mut self,
+        entity: Entity,
+        table_row: TableRow,
+    ) -> EntityLocation {
         self.entities.push(ArchetypeEntity { entity, table_row });
 
         EntityLocation {
             archetype_id: self.id,
-            index: self.entities.len() - 1,
+            archetype_row: ArchetypeRow(self.entities.len() - 1),
         }
     }
 
@@ -427,14 +458,14 @@ impl Archetype {
     ///
     /// # Panics
     /// This function will panic if `index >= self.len()`
-    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: ArchetypeRow) -> ArchetypeSwapRemoveResult {
+        let is_last = index.0 == self.entities.len() - 1;
+        let entity = self.entities.swap_remove(index.0);
         ArchetypeSwapRemoveResult {
             swapped_entity: if is_last {
                 None
             } else {
-                Some(self.entities[index].entity)
+                Some(self.entities[index.0].entity)
             },
             table_row: entity.table_row,
         }

crates/bevy_ecs/src/bundle.rs

@@ -6,12 +6,12 @@ pub use bevy_ecs_macros::Bundle;
 
 use crate::{
     archetype::{
-        Archetype, ArchetypeId, Archetypes, BundleComponentStatus, ComponentStatus,
+        Archetype, ArchetypeId, ArchetypeRow, Archetypes, BundleComponentStatus, ComponentStatus,
         SpawnBundleStatus,
     },
     component::{Component, ComponentId, Components, StorageType, Tick},
     entity::{Entities, Entity, EntityLocation},
-    storage::{SparseSetIndex, SparseSets, Storages, Table},
+    storage::{SparseSetIndex, SparseSets, Storages, Table, TableRow},
 };
 use bevy_ecs_macros::all_tuples;
 use bevy_ptr::OwningPtr;
@@ -379,7 +379,7 @@ impl BundleInfo {
         sparse_sets: &mut SparseSets,
         bundle_component_status: &S,
         entity: Entity,
-        table_row: usize,
+        table_row: TableRow,
         change_tick: u32,
         bundle: T,
     ) {
@@ -518,17 +518,17 @@ pub(crate) enum InsertBundleResult<'a> {
 
 impl<'a, 'b> BundleInserter<'a, 'b> {
     /// # Safety
-    /// `entity` must currently exist in the source archetype for this inserter. `archetype_index`
+    /// `entity` must currently exist in the source archetype for this inserter. `archetype_row`
     /// must be `entity`'s location in the archetype. `T` must match this [`BundleInfo`]'s type
     #[inline]
     pub unsafe fn insert<T: Bundle>(
         &mut self,
         entity: Entity,
-        archetype_index: usize,
+        archetype_row: ArchetypeRow,
         bundle: T,
     ) -> EntityLocation {
         let location = EntityLocation {
-            index: archetype_index,
+            archetype_row,
             archetype_id: self.archetype.id(),
         };
         match &mut self.result {
@@ -544,14 +544,14 @@ impl<'a, 'b> BundleInserter<'a, 'b> {
                     self.sparse_sets,
                     add_bundle,
                     entity,
-                    self.archetype.entity_table_row(archetype_index),
+                    self.archetype.entity_table_row(archetype_row),
                     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_row);
                 if let Some(swapped_entity) = result.swapped_entity {
                     self.entities.set(swapped_entity.index(), location);
                 }
@@ -579,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_row);
                 if let Some(swapped_entity) = result.swapped_entity {
                     self.entities.set(swapped_entity.index(), location);
                 }
@@ -607,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_row, 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, ArchetypeRow},
+    storage::SparseSetIndex,
+};
 use serde::{Deserialize, Serialize};
 use std::{convert::TryFrom, fmt, mem, sync::atomic::Ordering};
 
@@ -727,7 +730,7 @@ impl EntityMeta {
         generation: 0,
         location: EntityLocation {
             archetype_id: ArchetypeId::INVALID,
-            index: usize::MAX, // dummy value, to be filled in
+            archetype_row: ArchetypeRow::INVALID, // dummy value, to be filled in
         },
     };
 }
@@ -740,7 +743,7 @@ pub struct EntityLocation {
     pub archetype_id: ArchetypeId,
 
     /// The index of the entity in the archetype
-    pub index: usize,
+    pub archetype_row: ArchetypeRow,
 }
 
 #[cfg(test)]