Add a public API to ArchetypeGeneration/Id (#9825)

Objective
---------

- Since #6742, It is not possible to build an `ArchetypeId` from a
`ArchetypeGeneration`
- This was useful to 3rd party crate extending the base bevy ECS
capabilities, such as [`bevy_ecs_dynamic`] and now
[`bevy_mod_dynamic_query`]
- Making `ArchetypeGeneration` opaque this way made it completely
useless, and removed the ability to limit archetype updates to a subset
of archetypes.
- Making the `index` method on `ArchetypeId` private prevented the use
of bitfields and other optimized data structure to store sets of
archetype ids. (without `transmute`)

This PR is not a simple reversal of the change. It exposes a different
API, rethought to keep the private stuff private and the public stuff
less error-prone.

- Add a `StartRange<ArchetypeGeneration>` `Index` implementation to
`Archetypes`
- Instead of converting the generation into an index, then creating a
ArchetypeId from that index, and indexing `Archetypes` with it, use
directly the old `ArchetypeGeneration` to get the range of new
archetypes.

From careful benchmarking, it seems to also be a performance improvement
(~0-5%) on add_archetypes.

---

Changelog
---------

- Added `impl Index<RangeFrom<ArchetypeGeneration>> for Archetypes` this
allows you to get a slice of newly added archetypes since the last
recorded generation.
- Added `ArchetypeId::index` and `ArchetypeId::new` methods. It should
enable 3rd party crates to use the `Archetypes` API in a meaningful way.

[`bevy_ecs_dynamic`]:
https://github.com/jakobhellermann/bevy_ecs_dynamic/tree/main
[`bevy_mod_dynamic_query`]:
https://github.com/nicopap/bevy_mod_dynamic_query/

---------

Co-authored-by: vero <email@atlasdostal.com>

crates/bevy_ecs/src/archetype.rs

@@ -27,7 +27,7 @@ use crate::{
 };
 use std::{
     hash::Hash,
-    ops::{Index, IndexMut},
+    ops::{Index, IndexMut, RangeFrom},
 };
 
 /// An opaque location within a [`Archetype`].
@@ -70,7 +70,7 @@ impl ArchetypeRow {
 ///
 /// [`World`]: crate::world::World
 /// [`EMPTY`]: crate::archetype::ArchetypeId::EMPTY
-#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, PartialOrd, Ord)]
 // SAFETY: Must be repr(transparent) due to the safety requirements on EntityLocation
 #[repr(transparent)]
 pub struct ArchetypeId(u32);
@@ -83,13 +83,26 @@ impl ArchetypeId {
     /// This must always have an all-1s bit pattern to ensure soundness in fast entity id space allocation.
     pub const INVALID: ArchetypeId = ArchetypeId(u32::MAX);
 
+    /// Create an `ArchetypeId` from a plain value.
+    ///
+    /// This is useful if you need to store the `ArchetypeId` as a plain value,
+    /// for example in a specialized data structure such as a bitset.
+    ///
+    /// While it doesn't break any safety invariants, you should ensure the
+    /// values comes from a pre-existing [`ArchetypeId::index`] in this world
+    /// to avoid panics and other unexpected behaviors.
     #[inline]
-    pub(crate) const fn new(index: usize) -> Self {
+    pub const fn new(index: usize) -> Self {
         ArchetypeId(index as u32)
     }
 
+    /// The plain value of this `ArchetypeId`.
+    ///
+    /// In bevy, this is mostly used to store archetype ids in [`FixedBitSet`]s.
+    ///
+    /// [`FixedBitSet`]: fixedbitset::FixedBitSet
     #[inline]
-    pub(crate) fn index(self) -> usize {
+    pub fn index(self) -> usize {
         self.0 as usize
     }
 }
@@ -525,24 +538,23 @@ impl Archetype {
     }
 }
 
-/// An opaque generational id that changes every time the set of [`Archetypes`] changes.
+/// The next [`ArchetypeId`] in an [`Archetypes`] collection.
-#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
+///
-pub struct ArchetypeGeneration(usize);
+/// 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)]
+pub struct ArchetypeGeneration(ArchetypeId);
 
 impl ArchetypeGeneration {
+    /// The first archetype.
     #[inline]
-    pub(crate) const fn initial() -> Self {
+    pub const fn initial() -> Self {
-        ArchetypeGeneration(0)
+        ArchetypeGeneration(ArchetypeId::EMPTY)
-    }
-
-    #[inline]
-    pub(crate) fn value(self) -> usize {
-        self.0
     }
 }
 
 #[derive(Hash, PartialEq, Eq)]
-struct ArchetypeIdentity {
+struct ArchetypeComponents {
     table_components: Box<[ComponentId]>,
     sparse_set_components: Box<[ComponentId]>,
 }
@@ -603,25 +615,29 @@ impl SparseSetIndex for ArchetypeComponentId {
 pub struct Archetypes {
     pub(crate) archetypes: Vec<Archetype>,
     pub(crate) archetype_component_count: usize,
-    archetype_ids: bevy_utils::HashMap<ArchetypeIdentity, ArchetypeId>,
+    by_components: bevy_utils::HashMap<ArchetypeComponents, ArchetypeId>,
 }
 
 impl Archetypes {
     pub(crate) fn new() -> Self {
         let mut archetypes = Archetypes {
             archetypes: Vec::new(),
-            archetype_ids: Default::default(),
+            by_components: Default::default(),
             archetype_component_count: 0,
         };
         archetypes.get_id_or_insert(TableId::empty(), Vec::new(), Vec::new());
         archetypes
     }
 
-    /// Returns the current archetype generation. This is an ID indicating the current set of archetypes
+    /// Returns the "generation", a handle to the current highest archetype ID.
-    /// that are registered with the world.
+    ///
+    /// This can be used with the `Index` [`Archetypes`] implementation to
+    /// iterate over newly introduced [`Archetype`]s since the last time this
+    /// function was called.
     #[inline]
     pub fn generation(&self) -> ArchetypeGeneration {
-        ArchetypeGeneration(self.archetypes.len())
+        let id = ArchetypeId::new(self.archetypes.len());
+        ArchetypeGeneration(id)
     }
 
     /// Fetches the total number of [`Archetype`]s within the world.
@@ -692,7 +708,7 @@ impl Archetypes {
         table_components: Vec<ComponentId>,
         sparse_set_components: Vec<ComponentId>,
     ) -> ArchetypeId {
-        let archetype_identity = ArchetypeIdentity {
+        let archetype_identity = ArchetypeComponents {
             sparse_set_components: sparse_set_components.clone().into_boxed_slice(),
             table_components: table_components.clone().into_boxed_slice(),
         };
@@ -700,7 +716,7 @@ impl Archetypes {
         let archetypes = &mut self.archetypes;
         let archetype_component_count = &mut self.archetype_component_count;
         *self
-            .archetype_ids
+            .by_components
             .entry(archetype_identity)
             .or_insert_with(move || {
                 let id = ArchetypeId::new(archetypes.len());
@@ -739,6 +755,14 @@ impl Archetypes {
     }
 }
 
+impl Index<RangeFrom<ArchetypeGeneration>> for Archetypes {
+    type Output = [Archetype];
+
+    #[inline]
+    fn index(&self, index: RangeFrom<ArchetypeGeneration>) -> &Self::Output {
+        &self.archetypes[index.start.0.index()..]
+    }
+}
 impl Index<ArchetypeId> for Archetypes {
     type Output = Archetype;
 

crates/bevy_ecs/src/query/state.rs

@@ -216,12 +216,11 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
     pub fn update_archetypes_unsafe_world_cell(&mut self, world: UnsafeWorldCell) {
         self.validate_world(world.id());
         let archetypes = world.archetypes();
-        let new_generation = archetypes.generation();
+        let old_generation =
-        let old_generation = std::mem::replace(&mut self.archetype_generation, new_generation);
+            std::mem::replace(&mut self.archetype_generation, archetypes.generation());
-        let archetype_index_range = old_generation.value()..new_generation.value();
 
-        for archetype_index in archetype_index_range {
+        for archetype in &archetypes[old_generation..] {
-            self.new_archetype(&archetypes[ArchetypeId::new(archetype_index)]);
+            self.new_archetype(archetype);
         }
     }
 

crates/bevy_ecs/src/system/function_system.rs

@@ -1,5 +1,5 @@
 use crate::{
-    archetype::{ArchetypeComponentId, ArchetypeGeneration, ArchetypeId},
+    archetype::{ArchetypeComponentId, ArchetypeGeneration},
     component::{ComponentId, Tick},
     prelude::FromWorld,
     query::{Access, FilteredAccessSet},
@@ -270,16 +270,11 @@ impl<Param: SystemParam> SystemState<Param> {
     #[inline]
     pub fn update_archetypes_unsafe_world_cell(&mut self, world: UnsafeWorldCell) {
         let archetypes = world.archetypes();
-        let new_generation = archetypes.generation();
+        let old_generation =
-        let old_generation = std::mem::replace(&mut self.archetype_generation, new_generation);
+            std::mem::replace(&mut self.archetype_generation, archetypes.generation());
-        let archetype_index_range = old_generation.value()..new_generation.value();
 
-        for archetype_index in archetype_index_range {
+        for archetype in &archetypes[old_generation..] {
-            Param::new_archetype(
+            Param::new_archetype(&mut self.param_state, archetype, &mut self.meta);
-                &mut self.param_state,
-                &archetypes[ArchetypeId::new(archetype_index)],
-                &mut self.meta,
-            );
         }
     }
 
@@ -515,17 +510,12 @@ where
     fn update_archetype_component_access(&mut self, world: UnsafeWorldCell) {
         assert!(self.world_id == Some(world.id()), "Encountered a mismatched World. A System cannot be used with Worlds other than the one it was initialized with.");
         let archetypes = world.archetypes();
-        let new_generation = archetypes.generation();
+        let old_generation =
-        let old_generation = std::mem::replace(&mut self.archetype_generation, new_generation);
+            std::mem::replace(&mut self.archetype_generation, archetypes.generation());
-        let archetype_index_range = old_generation.value()..new_generation.value();
 
-        for archetype_index in archetype_index_range {
+        for archetype in &archetypes[old_generation..] {
             let param_state = self.param_state.as_mut().unwrap();
-            F::Param::new_archetype(
+            F::Param::new_archetype(param_state, archetype, &mut self.system_meta);
-                param_state,
-                &archetypes[ArchetypeId::new(archetype_index)],
-                &mut self.system_meta,
-            );
         }
     }