Clean up Fetch code (#4800)

# Objective
Clean up code surrounding fetch by pulling out the common parts into the iteration code.

## Solution
Merge `Fetch::table_fetch` and `Fetch::archetype_fetch` into a single API: `Fetch::fetch(&mut self, entity: &Entity, table_row: &usize)`. This provides everything any fetch requires to internally decide which storage to read from and get the underlying data. All of these functions are marked as `#[inline(always)]` and the arguments are passed as references to attempt to optimize out the argument that isn't being used.

External to `Fetch`, Query iteration has been changed to keep track of the table row and entity outside of fetch, which moves a lot of the expensive bookkeeping `Fetch` structs had previously done internally into the outer loop.

~~TODO: Benchmark, docs~~ Done.

---

## Changelog
Changed: `Fetch::table_fetch` and `Fetch::archetype_fetch` have been merged into a single `Fetch::fetch` function.

## Migration Guide
TODO

Co-authored-by: Brian Merchant <bhmerchang@gmail.com>
Co-authored-by: Saverio Miroddi <saverio.pub2@gmail.com>

crates/bevy_ecs/macros/src/fetch.rs

@@ -261,9 +261,9 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
                     _fetch: &mut <Self as #path::query::WorldQueryGats<'__w>>::Fetch,
                     _state: &Self::State,
                     _archetype: &'__w #path::archetype::Archetype,
-                    _tables: &'__w #path::storage::Tables
+                    _table: &'__w #path::storage::Table
                 ) {
-                    #(<#field_types>::set_archetype(&mut _fetch.#field_idents, &_state.#field_idents, _archetype, _tables);)*
+                    #(<#field_types>::set_archetype(&mut _fetch.#field_idents, &_state.#field_idents, _archetype, _table);)*
                 }
 
                 /// SAFETY: we call `set_table` for each member that implements `Fetch`
@@ -276,40 +276,27 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
                     #(<#field_types>::set_table(&mut _fetch.#field_idents, &_state.#field_idents, _table);)*
                 }
 
-                /// SAFETY: we call `table_fetch` for each member that implements `Fetch`.
-                #[inline]
-                unsafe fn table_fetch<'__w>(
+                /// SAFETY: we call `fetch` for each member that implements `Fetch`.
+                #[inline(always)]
+                unsafe fn fetch<'__w>(
                     _fetch: &mut <Self as #path::query::WorldQueryGats<'__w>>::Fetch,
+                    _entity: Entity,
                     _table_row: usize
                 ) -> <Self as #path::query::WorldQueryGats<'__w>>::Item {
                     Self::Item {
-                        #(#field_idents: <#field_types>::table_fetch(&mut _fetch.#field_idents, _table_row),)*
+                        #(#field_idents: <#field_types>::fetch(&mut _fetch.#field_idents, _entity, _table_row),)*
                         #(#ignored_field_idents: Default::default(),)*
                     }
                 }
 
-                /// SAFETY: we call `archetype_fetch` for each member that implements `Fetch`.
-                #[inline]
-                unsafe fn archetype_fetch<'__w>(
-                    _fetch: &mut <Self as #path::query::WorldQueryGats<'__w>>::Fetch,
-                    _archetype_index: usize
-                ) -> <Self as #path::query::WorldQueryGats<'__w>>::Item {
-                    Self::Item {
-                        #(#field_idents: <#field_types>::archetype_fetch(&mut _fetch.#field_idents, _archetype_index),)*
-                        #(#ignored_field_idents: Default::default(),)*
-                    }
-                }
-
-                #[allow(unused_variables)]
-                #[inline]
-                unsafe fn table_filter_fetch<'__w>(_fetch: &mut <Self as #path::query::WorldQueryGats<'__w>>::Fetch, _table_row: usize) -> bool {
-                    true #(&& <#field_types>::table_filter_fetch(&mut _fetch.#field_idents, _table_row))*
-                }
-
                 #[allow(unused_variables)]
-                #[inline]
-                unsafe fn archetype_filter_fetch<'__w>(_fetch: &mut <Self as #path::query::WorldQueryGats<'__w>>::Fetch, _archetype_index: usize) -> bool {
-                    true #(&& <#field_types>::archetype_filter_fetch(&mut _fetch.#field_idents, _archetype_index))*
+                #[inline(always)]
+                unsafe fn filter_fetch<'__w>(
+                    _fetch: &mut <Self as #path::query::WorldQueryGats<'__w>>::Fetch,
+                    _entity: Entity,
+                    _table_row: usize
+                ) -> bool {
+                    true #(&& <#field_types>::filter_fetch(&mut _fetch.#field_idents, _entity, _table_row))*
                 }
 
                 fn update_component_access(state: &Self::State, _access: &mut #path::query::FilteredAccess<#path::component::ComponentId>) {

crates/bevy_ecs/src/archetype.rs

@@ -121,9 +121,19 @@ impl Edges {
     }
 }
 
-struct TableInfo {
-    id: TableId,
-    entity_rows: Vec<usize>,
+pub struct ArchetypeEntity {
+    pub(crate) entity: Entity,
+    pub(crate) table_row: usize,
+}
+
+impl ArchetypeEntity {
+    pub fn entity(&self) -> Entity {
+        self.entity
+    }
+
+    pub fn table_row(&self) -> usize {
+        self.table_row
+    }
 }
 
 pub(crate) struct ArchetypeSwapRemoveResult {
@@ -138,9 +148,9 @@ pub(crate) struct ArchetypeComponentInfo {
 
 pub struct Archetype {
     id: ArchetypeId,
-    entities: Vec<Entity>,
+    table_id: TableId,
     edges: Edges,
-    table_info: TableInfo,
+    entities: Vec<ArchetypeEntity>,
     table_components: Box<[ComponentId]>,
     sparse_set_components: Box<[ComponentId]>,
     components: SparseSet<ComponentId, ArchetypeComponentInfo>,
@@ -183,14 +193,11 @@ impl Archetype {
         }
         Self {
             id,
-            table_info: TableInfo {
-                id: table_id,
-                entity_rows: Default::default(),
-            },
+            table_id,
+            entities: Vec::new(),
             components,
             table_components,
             sparse_set_components,
-            entities: Default::default(),
             edges: Default::default(),
         }
     }
@@ -202,19 +209,14 @@ impl Archetype {
 
     #[inline]
     pub fn table_id(&self) -> TableId {
-        self.table_info.id
+        self.table_id
     }
 
     #[inline]
-    pub fn entities(&self) -> &[Entity] {
+    pub fn entities(&self) -> &[ArchetypeEntity] {
         &self.entities
     }
 
-    #[inline]
-    pub fn entity_table_rows(&self) -> &[usize] {
-        &self.table_info.entity_rows
-    }
-
     #[inline]
     pub fn table_components(&self) -> &[ComponentId] {
         &self.table_components
@@ -242,20 +244,19 @@ impl Archetype {
 
     #[inline]
     pub fn entity_table_row(&self, index: usize) -> usize {
-        self.table_info.entity_rows[index]
+        self.entities[index].table_row
     }
 
     #[inline]
     pub(crate) fn set_entity_table_row(&mut self, index: usize, table_row: usize) {
-        self.table_info.entity_rows[index] = table_row;
+        self.entities[index].table_row = table_row;
     }
 
     /// # 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 {
-        self.entities.push(entity);
-        self.table_info.entity_rows.push(table_row);
+        self.entities.push(ArchetypeEntity { entity, table_row });
 
         EntityLocation {
             archetype_id: self.id,
@@ -265,21 +266,20 @@ impl Archetype {
 
     pub(crate) fn reserve(&mut self, additional: usize) {
         self.entities.reserve(additional);
-        self.table_info.entity_rows.reserve(additional);
     }
 
     /// 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;
-        self.entities.swap_remove(index);
+        let entity = self.entities.swap_remove(index);
         ArchetypeSwapRemoveResult {
             swapped_entity: if is_last {
                 None
             } else {
-                Some(self.entities[index])
+                Some(self.entities[index].entity)
             },
-            table_row: self.table_info.entity_rows.swap_remove(index),
+            table_row: entity.table_row,
         }
     }
 
@@ -317,7 +317,6 @@ impl Archetype {
 
     pub(crate) fn clear_entities(&mut self) {
         self.entities.clear();
-        self.table_info.entity_rows.clear();
     }
 }
 

crates/bevy_ecs/src/lib.rs

@@ -540,8 +540,8 @@ mod tests {
         let f = world
             .spawn((TableStored("def"), A(456), SparseStored(1)))
             .id();
-        // // this should be skipped
-        // SparseStored(1).spawn("abc");
+        // this should be skipped
+        // world.spawn(SparseStored(1));
         let ents = world
             .query::<(Entity, Option<&SparseStored>, &A)>()
             .iter(&world)