Add a ComponentIndex and update QueryState creation/update to use it (#13460)

# Objective

To implement relations we will need to add a `ComponentIndex`, which is
a map from a Component to the list of archetypes that contain this
component.
One of the reasons is that with fragmenting relations the number of
archetypes will explode, so it will become inefficient to create and
update the query caches by iterating through the list of all archetypes.

In this PR, we introduce the `ComponentIndex`, and we update the
`QueryState` to make use of it:
- if a query has at least 1 required component (i.e. something other
than `()`, `Entity` or `Option<>`, etc.): for each of the required
components we find the list of archetypes that contain it (using the
ComponentIndex). Then, we select the smallest list among these. This
gives a small subset of archetypes to iterate through compared with
iterating through all new archetypes
- if it doesn't, then we keep using the current approach of iterating
through all new archetypes


# Implementation
- This breaks query iteration order, in the sense that we are not
guaranteed anymore to return results in the order in which the
archetypes were created. I think this should be fine because this wasn't
an explicit bevy guarantee so users should not be relying on this. I
updated a bunch of unit tests that were failing because of this.

- I had an issue with the borrow checker because iterating the list of
potential archetypes requires access to `&state.component_access`, which
was conflicting with the calls to
```
  if state.new_archetype_internal(archetype) {
      state.update_archetype_component_access(archetype, access);
  }
```
which need a mutable access to the state.

The solution I chose was to introduce a `QueryStateView` which is a
temporary view into the `QueryState` which enables a "split-borrows"
kind of approach. It is described in detail in this blog post:
https://smallcultfollowing.com/babysteps/blog/2018/11/01/after-nll-interprocedural-conflicts/

# Test

The unit tests pass.

Benchmark results:
```
❯ critcmp main pr
group                                  main                                   pr
-----                                  ----                                   --
iter_fragmented/base                   1.00   342.2±25.45ns        ? ?/sec    1.02   347.5±16.24ns        ? ?/sec
iter_fragmented/foreach                1.04   165.4±11.29ns        ? ?/sec    1.00    159.5±4.27ns        ? ?/sec
iter_fragmented/foreach_wide           1.03      3.3±0.04µs        ? ?/sec    1.00      3.2±0.06µs        ? ?/sec
iter_fragmented/wide                   1.03      3.1±0.06µs        ? ?/sec    1.00      3.0±0.08µs        ? ?/sec
iter_fragmented_sparse/base            1.00      6.5±0.14ns        ? ?/sec    1.02      6.6±0.08ns        ? ?/sec
iter_fragmented_sparse/foreach         1.00      6.3±0.08ns        ? ?/sec    1.04      6.6±0.08ns        ? ?/sec
iter_fragmented_sparse/foreach_wide    1.00     43.8±0.15ns        ? ?/sec    1.02     44.6±0.53ns        ? ?/sec
iter_fragmented_sparse/wide            1.00     29.8±0.44ns        ? ?/sec    1.00     29.8±0.26ns        ? ?/sec
iter_simple/base                       1.00      8.2±0.10µs        ? ?/sec    1.00      8.2±0.09µs        ? ?/sec
iter_simple/foreach                    1.00      3.8±0.02µs        ? ?/sec    1.02      3.9±0.03µs        ? ?/sec
iter_simple/foreach_sparse_set         1.00     19.0±0.26µs        ? ?/sec    1.01     19.3±0.16µs        ? ?/sec
iter_simple/foreach_wide               1.00     17.8±0.24µs        ? ?/sec    1.00     17.9±0.31µs        ? ?/sec
iter_simple/foreach_wide_sparse_set    1.06     95.6±6.23µs        ? ?/sec    1.00     90.6±0.59µs        ? ?/sec
iter_simple/sparse_set                 1.00     19.3±1.63µs        ? ?/sec    1.01     19.5±0.29µs        ? ?/sec
iter_simple/system                     1.00      8.1±0.10µs        ? ?/sec    1.00      8.1±0.09µs        ? ?/sec
iter_simple/wide                       1.05     37.7±2.53µs        ? ?/sec    1.00     35.8±0.57µs        ? ?/sec
iter_simple/wide_sparse_set            1.00     95.7±1.62µs        ? ?/sec    1.00     95.9±0.76µs        ? ?/sec
par_iter_simple/with_0_fragment        1.04     35.0±2.51µs        ? ?/sec    1.00     33.7±0.49µs        ? ?/sec
par_iter_simple/with_1000_fragment     1.00     50.4±2.52µs        ? ?/sec    1.01     51.0±3.84µs        ? ?/sec
par_iter_simple/with_100_fragment      1.02     40.3±2.23µs        ? ?/sec    1.00     39.5±1.32µs        ? ?/sec
par_iter_simple/with_10_fragment       1.14     38.8±7.79µs        ? ?/sec    1.00     34.0±0.78µs        ? ?/sec
```

Author: cBournhonesque

benches/benches/bevy_render/torus.rs

@@ -4,12 +4,9 @@ use bevy_render::mesh::TorusMeshBuilder;
 
 fn torus(c: &mut Criterion) {
     c.bench_function("build_torus", |b| {
-        b.iter(|| black_box(TorusMeshBuilder::new(black_box(0.5),black_box(1.0))));
+        b.iter(|| black_box(TorusMeshBuilder::new(black_box(0.5), black_box(1.0))));
     });
 }
 
-criterion_group!(
-    benches,
-    torus,
-);
+criterion_group!(benches, torus,);
 criterion_main!(benches);

crates/bevy_ecs/src/archetype.rs

@@ -26,6 +26,7 @@ use crate::{
     observer::Observers,
     storage::{ImmutableSparseSet, SparseArray, SparseSet, SparseSetIndex, TableId, TableRow},
 };
+use bevy_utils::HashMap;
 use std::{
     hash::Hash,
     ops::{Index, IndexMut, RangeFrom},
@@ -345,8 +346,10 @@ pub struct Archetype {
 }
 
 impl Archetype {
+    /// `table_components` and `sparse_set_components` must be sorted
     pub(crate) fn new(
         components: &Components,
+        component_index: &mut ComponentIndex,
         observers: &Observers,
         id: ArchetypeId,
         table_id: TableId,
@@ -357,7 +360,7 @@ impl Archetype {
         let (min_sparse, _) = sparse_set_components.size_hint();
         let mut flags = ArchetypeFlags::empty();
         let mut archetype_components = SparseSet::with_capacity(min_table + min_sparse);
-        for (component_id, archetype_component_id) in table_components {
+        for (idx, (component_id, archetype_component_id)) in table_components.enumerate() {
             // SAFETY: We are creating an archetype that includes this component so it must exist
             let info = unsafe { components.get_info_unchecked(component_id) };
             info.update_archetype_flags(&mut flags);
@@ -369,6 +372,13 @@ impl Archetype {
                     archetype_component_id,
                 },
             );
+            // NOTE: the `table_components` are sorted AND they were inserted in the `Table` in the same
+            // sorted order, so the index of the `Column` in the `Table` is the same as the index of the
+            // component in the `table_components` vector
+            component_index
+                .entry(component_id)
+                .or_insert_with(HashMap::new)
+                .insert(id, ArchetypeRecord { column: Some(idx) });
         }
 
         for (component_id, archetype_component_id) in sparse_set_components {
@@ -383,6 +393,10 @@ impl Archetype {
                     archetype_component_id,
                 },
             );
+            component_index
+                .entry(component_id)
+                .or_insert_with(HashMap::new)
+                .insert(id, ArchetypeRecord { column: None });
         }
         Self {
             id,
@@ -460,7 +474,7 @@ impl Archetype {
         self.components.len()
     }
 
-    /// Fetches a immutable reference to the archetype's [`Edges`], a cache of
+    /// Fetches an immutable reference to the archetype's [`Edges`], a cache of
     /// archetypal relationships.
     #[inline]
     pub fn edges(&self) -> &Edges {
@@ -655,7 +669,7 @@ impl Archetype {
 /// This is used in archetype update methods to limit archetype updates to the
 /// ones added since the last time the method ran.
 #[derive(Debug, Copy, Clone, PartialEq)]
-pub struct ArchetypeGeneration(ArchetypeId);
+pub struct ArchetypeGeneration(pub(crate) ArchetypeId);
 
 impl ArchetypeGeneration {
     /// The first archetype.
@@ -711,6 +725,10 @@ impl SparseSetIndex for ArchetypeComponentId {
     }
 }
 
+/// Maps a [`ComponentId`] to the list of [`Archetypes`]([`Archetype`]) that contain the [`Component`](crate::component::Component),
+/// along with an [`ArchetypeRecord`] which contains some metadata about how the component is stored in the archetype.
+pub type ComponentIndex = HashMap<ComponentId, HashMap<ArchetypeId, ArchetypeRecord>>;
+
 /// The backing store of all [`Archetype`]s within a [`World`].
 ///
 /// For more information, see the *[module level documentation]*.
@@ -720,14 +738,26 @@ impl SparseSetIndex for ArchetypeComponentId {
 pub struct Archetypes {
     pub(crate) archetypes: Vec<Archetype>,
     archetype_component_count: usize,
-    by_components: bevy_utils::HashMap<ArchetypeComponents, ArchetypeId>,
+    /// find the archetype id by the archetype's components
+    by_components: HashMap<ArchetypeComponents, ArchetypeId>,
+    /// find all the archetypes that contain a component
+    by_component: ComponentIndex,
+}
+
+/// Metadata about how a component is stored in an [`Archetype`].
+pub struct ArchetypeRecord {
+    /// Index of the component in the archetype's [`Table`](crate::storage::Table),
+    /// or None if the component is a sparse set component.
+    #[allow(dead_code)]
+    pub(crate) column: Option<usize>,
 }
 
 impl Archetypes {
     pub(crate) fn new() -> Self {
         let mut archetypes = Archetypes {
             archetypes: Vec::new(),
             by_components: Default::default(),
+            by_component: Default::default(),
             archetype_component_count: 0,
         };
         // SAFETY: Empty archetype has no components
@@ -770,7 +800,7 @@ impl Archetypes {
         unsafe { self.archetypes.get_unchecked(ArchetypeId::EMPTY.index()) }
     }
 
-    /// Fetches an mutable reference to the archetype without any components.
+    /// Fetches a mutable reference to the archetype without any components.
     #[inline]
     pub(crate) fn empty_mut(&mut self) -> &mut Archetype {
         // SAFETY: empty archetype always exists
@@ -848,28 +878,27 @@ impl Archetypes {
 
         let archetypes = &mut self.archetypes;
         let archetype_component_count = &mut self.archetype_component_count;
-        *self
+        let component_index = &mut self.by_component;
+        let archetype_id = *self
             .by_components
             .entry(archetype_identity)
             .or_insert_with_key(move |identity| {
                 let ArchetypeComponents {
                     table_components,
                     sparse_set_components,
                 } = identity;
-
                 let id = ArchetypeId::new(archetypes.len());
                 let table_start = *archetype_component_count;
                 *archetype_component_count += table_components.len();
                 let table_archetype_components =
                     (table_start..*archetype_component_count).map(ArchetypeComponentId);
-
                 let sparse_start = *archetype_component_count;
                 *archetype_component_count += sparse_set_components.len();
                 let sparse_set_archetype_components =
                     (sparse_start..*archetype_component_count).map(ArchetypeComponentId);
-
                 archetypes.push(Archetype::new(
                     components,
+                    component_index,
                     observers,
                     id,
                     table_id,
@@ -883,7 +912,8 @@ impl Archetypes {
                         .zip(sparse_set_archetype_components),
                 ));
                 id
-            })
+            });
+        archetype_id
     }
 
     /// Returns the number of components that are stored in archetypes.
@@ -900,16 +930,25 @@ impl Archetypes {
         }
     }
 
+    /// Get the component index
+    pub(crate) fn component_index(&self) -> &ComponentIndex {
+        &self.by_component
+    }
+
     pub(crate) fn update_flags(
         &mut self,
         component_id: ComponentId,
         flags: ArchetypeFlags,
         set: bool,
     ) {
-        // TODO: Refactor component index to speed this up.
-        for archetype in &mut self.archetypes {
-            if archetype.contains(component_id) {
-                archetype.flags.set(flags, set);
+        if let Some(archetypes) = self.by_component.get(&component_id) {
+            for archetype_id in archetypes.keys() {
+                // SAFETY: the component index only contains valid archetype ids
+                self.archetypes
+                    .get_mut(archetype_id.index())
+                    .unwrap()
+                    .flags
+                    .set(flags, set);
             }
         }
     }