remove unsafe get_unchecked (and mut variant) from Tables and Archety…

…pes (#1614)

Removes `get_unchecked` and `get_unchecked_mut` from `Tables` and `Archetypes` collections in favor of safe Index implementations. This fixes a safety error in `Archetypes::get_id_or_insert()` (which previously relied on TableId being valid to be safe ... the alternative was to make that method unsafe too). It also cuts down on a lot of unsafe and makes the code easier to look at. I'm not sure what changed since the last benchmark, but these numbers are more favorable than my last tests of similar changes. I didn't include the Components collection as those severely killed perf last time I tried. But this does inspire me to try again (just in a separate pr)! 

Note that the `simple_insert/bevy_unbatched` benchmark fluctuates a lot on both branches (this was also true for prior versions of bevy). It seems like the allocator has more variance for many small allocations. And `sparse_frag_iter/bevy` operates on such a small scale that 10% fluctuations are common.

Some benches do take a small hit here, but I personally think its worth it.

This also fixes a safety error in Query::for_each_mut, which needed to mutably borrow Query (aaahh!).  

![image](https://user-images.githubusercontent.com/2694663/110726926-2b52eb80-81cf-11eb-9ea3-bff951060c7c.png)
![image](https://user-images.githubusercontent.com/2694663/110726991-4c1b4100-81cf-11eb-9199-ca79bef0b9bd.png)

crates/bevy_ecs/src/archetype.rs

@@ -4,7 +4,12 @@ use crate::{
     entity::{Entity, EntityLocation},
     storage::{Column, SparseArray, SparseSet, SparseSetIndex, TableId},
 };
-use std::{borrow::Cow, collections::HashMap, hash::Hash};
+use std::{
+    borrow::Cow,
+    collections::HashMap,
+    hash::Hash,
+    ops::{Index, IndexMut},
+};
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
 pub struct ArchetypeId(usize);
@@ -443,27 +448,11 @@ impl Archetypes {
         self.archetypes.get(id.index())
     }
 
-    /// # Safety
-    /// `id` must be valid
-    #[inline]
-    pub unsafe fn get_unchecked(&self, id: ArchetypeId) -> &Archetype {
-        debug_assert!(id.index() < self.archetypes.len());
-        self.archetypes.get_unchecked(id.index())
-    }
-
     #[inline]
     pub fn get_mut(&mut self, id: ArchetypeId) -> Option<&mut Archetype> {
         self.archetypes.get_mut(id.index())
     }
 
-    /// # Safety
-    /// `id` must be valid
-    #[inline]
-    pub unsafe fn get_unchecked_mut(&mut self, id: ArchetypeId) -> &mut Archetype {
-        debug_assert!(id.index() < self.archetypes.len());
-        self.archetypes.get_unchecked_mut(id.index())
-    }
-
     #[inline]
     pub fn iter(&self) -> impl Iterator<Item = &Archetype> {
         self.archetypes.iter()
@@ -522,3 +511,19 @@ impl Archetypes {
         self.archetype_component_count
     }
 }
+
+impl Index<ArchetypeId> for Archetypes {
+    type Output = Archetype;
+
+    #[inline]
+    fn index(&self, index: ArchetypeId) -> &Self::Output {
+        &self.archetypes[index.index()]
+    }
+}
+
+impl IndexMut<ArchetypeId> for Archetypes {
+    #[inline]
+    fn index_mut(&mut self, index: ArchetypeId) -> &mut Self::Output {
+        &mut self.archetypes[index.index()]
+    }
+}

crates/bevy_ecs/src/query/fetch.rs

@@ -271,9 +271,9 @@ impl<'w, T: Component> Fetch<'w> for ReadFetch<T> {
         match state.storage_type {
             StorageType::Table => {
                 self.entity_table_rows = archetype.entity_table_rows().as_ptr();
-                // SAFE: archetype tables always exist
-                let table = tables.get_unchecked(archetype.table_id());
-                let column = table.get_column(state.component_id).unwrap();
+                let column = tables[archetype.table_id()]
+                    .get_column(state.component_id)
+                    .unwrap();
                 self.table_components = column.get_ptr().cast::<T>();
             }
             StorageType::SparseSet => self.entities = archetype.entities().as_ptr(),
@@ -411,9 +411,9 @@ impl<'w, T: Component> Fetch<'w> for WriteFetch<T> {
         match state.storage_type {
             StorageType::Table => {
                 self.entity_table_rows = archetype.entity_table_rows().as_ptr();
-                // SAFE: archetype tables always exist
-                let table = tables.get_unchecked(archetype.table_id());
-                let column = table.get_column(state.component_id).unwrap();
+                let column = tables[archetype.table_id()]
+                    .get_column(state.component_id)
+                    .unwrap();
                 self.table_components = column.get_ptr().cast::<T>();
                 self.table_flags = column.get_flags_mut_ptr();
             }
@@ -698,9 +698,9 @@ impl<'w, T: Component> Fetch<'w> for FlagsFetch<T> {
         match state.storage_type {
             StorageType::Table => {
                 self.entity_table_rows = archetype.entity_table_rows().as_ptr();
-                // SAFE: archetype tables always exist
-                let table = tables.get_unchecked(archetype.table_id());
-                let column = table.get_column(state.component_id).unwrap();
+                let column = tables[archetype.table_id()]
+                    .get_column(state.component_id)
+                    .unwrap();
                 self.table_flags = column.get_flags_mut_ptr().cast::<ComponentFlags>();
             }
             StorageType::SparseSet => self.entities = archetype.entities().as_ptr(),

crates/bevy_ecs/src/query/filter.rs

@@ -546,8 +546,7 @@ macro_rules! impl_flag_filter {
                 match state.storage_type {
                     StorageType::Table => {
                         self.entity_table_rows = archetype.entity_table_rows().as_ptr();
-                        // SAFE: archetype tables always exist
-                        let table = tables.get_unchecked(archetype.table_id());
+                        let table = &tables[archetype.table_id()];
                         self.table_flags = table
                             .get_column(state.component_id).unwrap()
                             .get_flags_mut_ptr();

crates/bevy_ecs/src/query/iter.rs

@@ -58,7 +58,7 @@ where
                 loop {
                     if self.current_index == self.current_len {
                         let table_id = self.table_id_iter.next()?;
-                        let table = self.tables.get_unchecked(*table_id);
+                        let table = &self.tables[*table_id];
                         self.fetch.set_table(&self.query_state.fetch_state, table);
                         self.filter.set_table(&self.query_state.filter_state, table);
                         self.current_len = table.len();
@@ -80,7 +80,7 @@ where
                 loop {
                     if self.current_index == self.current_len {
                         let archetype_id = self.archetype_id_iter.next()?;
-                        let archetype = self.archetypes.get_unchecked(*archetype_id);
+                        let archetype = &self.archetypes[*archetype_id];
                         self.fetch.set_archetype(
                             &self.query_state.fetch_state,
                             archetype,
@@ -120,12 +120,7 @@ impl<'w, 's, Q: WorldQuery> ExactSizeIterator for QueryIter<'w, 's, Q, ()> {
         self.query_state
             .matched_archetypes
             .ones()
-            .map(|index| {
-                // SAFE: matched archetypes always exist
-                let archetype =
-                    unsafe { self.world.archetypes.get_unchecked(ArchetypeId::new(index)) };
-                archetype.len()
-            })
+            .map(|index| self.world.archetypes[ArchetypeId::new(index)].len())
             .sum()
     }
 }

crates/bevy_ecs/src/query/state.rs

@@ -75,9 +75,7 @@ where
             }
         };
         for archetype_index in archetype_index_range {
-            // SAFE: archetype indices less than the archetype generation are guaranteed to exist
-            let archetype = unsafe { archetypes.get_unchecked(ArchetypeId::new(archetype_index)) };
-            self.new_archetype(archetype);
+            self.new_archetype(&archetypes[ArchetypeId::new(archetype_index)]);
         }
     }
 
@@ -155,8 +153,7 @@ where
         {
             return Err(QueryEntityError::QueryDoesNotMatch);
         }
-        // SAFE: live entities always exist in an archetype
-        let archetype = world.archetypes.get_unchecked(location.archetype_id);
+        let archetype = &world.archetypes[location.archetype_id];
         let mut fetch = <Q::Fetch as Fetch>::init(world, &self.fetch_state);
         let mut filter = <F::Fetch as Fetch>::init(world, &self.filter_state);
 
@@ -308,7 +305,7 @@ where
         if fetch.is_dense() && filter.is_dense() {
             let tables = &world.storages().tables;
             for table_id in self.matched_table_ids.iter() {
-                let table = tables.get_unchecked(*table_id);
+                let table = &tables[*table_id];
                 fetch.set_table(&self.fetch_state, table);
                 filter.set_table(&self.filter_state, table);
 
@@ -324,7 +321,7 @@ where
             let archetypes = &world.archetypes;
             let tables = &world.storages().tables;
             for archetype_id in self.matched_archetype_ids.iter() {
-                let archetype = archetypes.get_unchecked(*archetype_id);
+                let archetype = &archetypes[*archetype_id];
                 fetch.set_archetype(&self.fetch_state, archetype, tables);
                 filter.set_archetype(&self.filter_state, archetype, tables);
 
@@ -357,15 +354,15 @@ where
             if fetch.is_dense() && filter.is_dense() {
                 let tables = &world.storages().tables;
                 for table_id in self.matched_table_ids.iter() {
-                    let table = tables.get_unchecked(*table_id);
+                    let table = &tables[*table_id];
                     let mut offset = 0;
                     while offset < table.len() {
                         let func = func.clone();
                         scope.spawn(async move {
                             let mut fetch = <Q::Fetch as Fetch>::init(world, &self.fetch_state);
                             let mut filter = <F::Fetch as Fetch>::init(world, &self.filter_state);
                             let tables = &world.storages().tables;
-                            let table = tables.get_unchecked(*table_id);
+                            let table = &tables[*table_id];
                             fetch.set_table(&self.fetch_state, table);
                             filter.set_table(&self.filter_state, table);
                             let len = batch_size.min(table.len() - offset);
@@ -384,14 +381,14 @@ where
                 let archetypes = &world.archetypes;
                 for archetype_id in self.matched_archetype_ids.iter() {
                     let mut offset = 0;
-                    let archetype = archetypes.get_unchecked(*archetype_id);
+                    let archetype = &archetypes[*archetype_id];
                     while offset < archetype.len() {
                         let func = func.clone();
                         scope.spawn(async move {
                             let mut fetch = <Q::Fetch as Fetch>::init(world, &self.fetch_state);
                             let mut filter = <F::Fetch as Fetch>::init(world, &self.filter_state);
                             let tables = &world.storages().tables;
-                            let archetype = world.archetypes.get_unchecked(*archetype_id);
+                            let archetype = &world.archetypes[*archetype_id];
                             fetch.set_archetype(&self.fetch_state, archetype, tables);
                             filter.set_archetype(&self.filter_state, archetype, tables);
 

crates/bevy_ecs/src/storage/table.rs

@@ -8,6 +8,7 @@ use bevy_utils::{AHasher, HashMap};
 use std::{
     cell::UnsafeCell,
     hash::{Hash, Hasher},
+    ops::{Index, IndexMut},
     ptr::NonNull,
 };
 
@@ -407,30 +408,14 @@ impl Tables {
         self.tables.is_empty()
     }
 
-    #[inline]
-    pub fn get_mut(&mut self, id: TableId) -> Option<&mut Table> {
-        self.tables.get_mut(id.index())
-    }
-
     #[inline]
     pub fn get(&self, id: TableId) -> Option<&Table> {
         self.tables.get(id.index())
     }
 
-    /// # Safety
-    /// `id` must be a valid table
     #[inline]
-    pub unsafe fn get_unchecked_mut(&mut self, id: TableId) -> &mut Table {
-        debug_assert!(id.index() < self.tables.len());
-        self.tables.get_unchecked_mut(id.index())
-    }
-
-    /// # Safety
-    /// `id` must be a valid table
-    #[inline]
-    pub unsafe fn get_unchecked(&self, id: TableId) -> &Table {
-        debug_assert!(id.index() < self.tables.len());
-        self.tables.get_unchecked(id.index())
+    pub fn get_mut(&mut self, id: TableId) -> Option<&mut Table> {
+        self.tables.get_mut(id.index())
     }
 
     #[inline]
@@ -476,6 +461,22 @@ impl Tables {
     }
 }
 
+impl Index<TableId> for Tables {
+    type Output = Table;
+
+    #[inline]
+    fn index(&self, index: TableId) -> &Self::Output {
+        &self.tables[index.index()]
+    }
+}
+
+impl IndexMut<TableId> for Tables {
+    #[inline]
+    fn index_mut(&mut self, index: TableId) -> &mut Self::Output {
+        &mut self.tables[index.index()]
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{

crates/bevy_ecs/src/system/query.rs

@@ -77,7 +77,7 @@ where
     /// Runs `f` on each query result. This is faster than the equivalent iter() method, but cannot
     /// be chained like a normal iterator.
     #[inline]
-    pub fn for_each_mut(&self, f: impl FnMut(<Q::Fetch as Fetch<'w>>::Item)) {
+    pub fn for_each_mut(&mut self, f: impl FnMut(<Q::Fetch as Fetch<'w>>::Item)) {
         // SAFE: system runs without conflicts with other systems. same-system queries have runtime
         // borrow checks when they conflict
         unsafe { self.state.for_each_unchecked_manual(self.world, f) };