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.
-#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
-pub struct ArchetypeGeneration(usize);
+/// The next [`ArchetypeId`] in an [`Archetypes`] collection.
+///
+/// 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 {
- ArchetypeGeneration(0)
- }
-
- #[inline]
- pub(crate) fn value(self) -> usize {
- self.0
+ pub const fn initial() -> Self {
+ ArchetypeGeneration(ArchetypeId::EMPTY)
  }
 }
 
 #[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
- /// that are registered with the world.
+ /// Returns the "generation", a handle to the current highest archetype ID.
+ ///
+ /// 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,15 +708,15 @@ 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(),
  };
 
  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 = std::mem::replace(&mut self.archetype_generation, new_generation);
- let archetype_index_range = old_generation.value()..new_generation.value();
+ let old_generation =
+ std::mem::replace(&mut self.archetype_generation, archetypes.generation());
 
- for archetype_index in archetype_index_range {
- self.new_archetype(&archetypes[ArchetypeId::new(archetype_index)]);
+ for archetype in &archetypes[old_generation..] {
+ 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 = std::mem::replace(&mut self.archetype_generation, new_generation);
- let archetype_index_range = old_generation.value()..new_generation.value();
-
- for archetype_index in archetype_index_range {
- Param::new_archetype(
- &mut self.param_state,
- &archetypes[ArchetypeId::new(archetype_index)],
- &mut self.meta,
- );
+ let old_generation =
+ std::mem::replace(&mut self.archetype_generation, archetypes.generation());
+
+ for archetype in &archetypes[old_generation..] {
+ Param::new_archetype(&mut self.param_state, archetype, &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 = std::mem::replace(&mut self.archetype_generation, new_generation);
- let archetype_index_range = old_generation.value()..new_generation.value();
+ let old_generation =
+ std::mem::replace(&mut self.archetype_generation, archetypes.generation());
 
- 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(
- param_state,
- &archetypes[ArchetypeId::new(archetype_index)],
- &mut self.system_meta,
- );
+ F::Param::new_archetype(param_state, archetype, &mut self.system_meta);
  }
  }