Component Lifecycle Hooks and a Deferred World (bevyengine#10756)

# Objective

- Provide a reliable and performant mechanism to allows users to keep
components synchronized with external sources: closing/opening sockets,
updating indexes, debugging etc.
- Implement a generic mechanism to provide mutable access to the world
without allowing structural changes; this will not only be used here but
is a foundational piece for observers, which are key for a performant
implementation of relations.

## Solution

- Implement a new type `DeferredWorld` (naming is not important,
`StaticWorld` is also suitable) that wraps a world pointer and prevents
user code from making any structural changes to the ECS; spawning
entities, creating components, initializing resources etc.
- Add component lifecycle hooks `on_add`, `on_insert` and `on_remove`
that can be assigned callbacks in user code.

---

## Changelog
- Add new `DeferredWorld` type.
- Add new world methods: `register_component::<T>` and
`register_component_with_descriptor`. These differ from `init_component`
in that they provide mutable access to the created `ComponentInfo` but
will panic if the component is already in any archetypes. These
restrictions serve two purposes:
1. Prevent users from defining hooks for components that may already
have associated hooks provided in another plugin. (a use case better
served by observers)
2. Ensure that when an `Archetype` is created it gets the appropriate
flags to early-out when triggering hooks.
- Add methods to `ComponentInfo`: `on_add`, `on_insert` and `on_remove`
to be used to register hooks of the form `fn(DeferredWorld, Entity,
ComponentId)`
- Modify `BundleInserter`, `BundleSpawner` and `EntityWorldMut` to
trigger component hooks when appropriate.
- Add bit flags to `Archetype` indicating whether or not any contained
components have each type of hook, this can be expanded for other flags
as needed.
- Add `component_hooks` example to illustrate usage. Try it out! It's
fun to mash keys.

## Safety
The changes to component insertion, removal and deletion involve a large
amount of unsafe code and it's fair for that to raise some concern. I
have attempted to document it as clearly as possible and have confirmed
that all the hooks examples are accepted by `cargo miri` as not causing
any undefined behavior. The largest issue is in ensuring there are no
outstanding references when passing a `DeferredWorld` to the hooks which
requires some use of raw pointers (as was already happening to some
degree in those places) and I have taken some time to ensure that is the
case but feel free to let me know if I've missed anything.

## Performance
These changes come with a small but measurable performance cost of
between 1-5% on `add_remove` benchmarks and between 1-3% on `insert`
benchmarks. One consideration to be made is the existence of the current
`RemovedComponents` which is on average more costly than the addition of
`on_remove` hooks due to the early-out, however hooks doesn't completely
remove the need for `RemovedComponents` as there is a chance you want to
respond to the removal of a component that already has an `on_remove`
hook defined in another plugin, so I have not removed it here. I do
intend to deprecate it with the introduction of observers in a follow up
PR.

## Discussion Questions
- Currently `DeferredWorld` implements `Deref` to `&World` which makes
sense conceptually, however it does cause some issues with rust-analyzer
providing autocomplete for `&mut World` references which is annoying.
There are alternative implementations that may address this but involve
more code churn so I have attempted them here. The other alternative is
to not implement `Deref` at all but that leads to a large amount of API
duplication.
- `DeferredWorld`, `StaticWorld`, something else?
- In adding support for hooks to `EntityWorldMut` I encountered some
unfortunate difficulties with my desired API. If commands are flushed
after each call i.e. `world.spawn() // flush commands .insert(A) //
flush commands` the entity may be despawned while `EntityWorldMut` still
exists which is invalid. An alternative was then to add
`self.world.flush_commands()` to the drop implementation for
`EntityWorldMut` but that runs into other problems for implementing
functions like `into_unsafe_entity_cell`. For now I have implemented a
`.flush()` which will flush the commands and consume `EntityWorldMut` or
users can manually run `world.flush_commands()` after using
`EntityWorldMut`.
- In order to allowing querying on a deferred world we need
implementations of `WorldQuery` to not break our guarantees of no
structural changes through their `UnsafeWorldCell`. All our
implementations do this, but there isn't currently any safety
documentation specifying what is or isn't allowed for an implementation,
just for the caller, (they also shouldn't be aliasing components they
didn't specify access for etc.) is that something we should start doing?
(see 10752)

Please check out the example `component_hooks` or the tests in
`bundle.rs` for usage examples. I will continue to expand this
description as I go.

See bevyengine#10839 for a more ergonomic API built on top of this one that isn't
subject to the same restrictions and supports `SystemParam` dependency
injection.

Cargo.toml

@@ -1403,6 +1403,17 @@ description = "Change detection on components"
 category = "ECS (Entity Component System)"
 wasm = false
 
+[[example]]
+name = "component_hooks"
+path = "examples/ecs/component_hooks.rs"
+doc-scrape-examples = true
+
+[package.metadata.example.component_hooks]
+name = "Component Hooks"
+description = "Define component hooks to manage component lifecycle events"
+category = "ECS (Entity Component System)"
+wasm = false
+
 [[example]]
 name = "custom_schedule"
 path = "examples/ecs/custom_schedule.rs"

crates/bevy_ecs/Cargo.toml

@@ -22,6 +22,7 @@ bevy_tasks = { path = "../bevy_tasks", version = "0.14.0-dev" }
 bevy_utils = { path = "../bevy_utils", version = "0.14.0-dev" }
 bevy_ecs_macros = { path = "macros", version = "0.14.0-dev" }
 
+bitflags = "2.3"
 concurrent-queue = "2.4.0"
 fixedbitset = "0.4.2"
 rustc-hash = "1.1"

crates/bevy_ecs/src/archetype.rs

@@ -21,7 +21,7 @@
 
 use crate::{
     bundle::BundleId,
-    component::{ComponentId, StorageType},
+    component::{ComponentId, Components, StorageType},
     entity::{Entity, EntityLocation},
     storage::{ImmutableSparseSet, SparseArray, SparseSet, SparseSetIndex, TableId, TableRow},
 };
@@ -107,7 +107,7 @@ impl ArchetypeId {
     }
 }
 
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, Eq, PartialEq)]
 pub(crate) enum ComponentStatus {
     Added,
     Mutated,
@@ -298,6 +298,18 @@ struct ArchetypeComponentInfo {
     archetype_component_id: ArchetypeComponentId,
 }
 
+bitflags::bitflags! {
+    /// Flags used to keep track of metadata about the component in this [`Archetype`]
+    ///
+    /// Used primarily to early-out when there are no [`ComponentHook`] registered for any contained components.
+    #[derive(Clone, Copy)]
+    pub(crate) struct ArchetypeFlags: u32 {
+        const ON_ADD_HOOK    = (1 << 0);
+        const ON_INSERT_HOOK = (1 << 1);
+        const ON_REMOVE_HOOK = (1 << 2);
+    }
+}
+
 /// Metadata for a single archetype within a [`World`].
 ///
 /// For more information, see the *[module level documentation]*.
@@ -310,20 +322,26 @@ pub struct Archetype {
     edges: Edges,
     entities: Vec<ArchetypeEntity>,
     components: ImmutableSparseSet<ComponentId, ArchetypeComponentInfo>,
+    flags: ArchetypeFlags,
 }
 
 impl Archetype {
     pub(crate) fn new(
+        components: &Components,
         id: ArchetypeId,
         table_id: TableId,
         table_components: impl Iterator<Item = (ComponentId, ArchetypeComponentId)>,
         sparse_set_components: impl Iterator<Item = (ComponentId, ArchetypeComponentId)>,
     ) -> Self {
         let (min_table, _) = table_components.size_hint();
         let (min_sparse, _) = sparse_set_components.size_hint();
-        let mut components = SparseSet::with_capacity(min_table + min_sparse);
+        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 {
-            components.insert(
+            // 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);
+            archetype_components.insert(
                 component_id,
                 ArchetypeComponentInfo {
                     storage_type: StorageType::Table,
@@ -333,7 +351,10 @@ impl Archetype {
         }
 
         for (component_id, archetype_component_id) in sparse_set_components {
-            components.insert(
+            // 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);
+            archetype_components.insert(
                 component_id,
                 ArchetypeComponentInfo {
                     storage_type: StorageType::SparseSet,
@@ -345,8 +366,9 @@ impl Archetype {
             id,
             table_id,
             entities: Vec::new(),
-            components: components.into_immutable(),
+            components: archetype_components.into_immutable(),
             edges: Default::default(),
+            flags,
         }
     }
 
@@ -356,6 +378,12 @@ impl Archetype {
         self.id
     }
 
+    /// Fetches the flags for the archetype.
+    #[inline]
+    pub(crate) fn flags(&self) -> ArchetypeFlags {
+        self.flags
+    }
+
     /// Fetches the archetype's [`Table`] ID.
     ///
     /// [`Table`]: crate::storage::Table
@@ -542,6 +570,24 @@ impl Archetype {
     pub(crate) fn clear_entities(&mut self) {
         self.entities.clear();
     }
+
+    /// Returns true if any of the components in this archetype have `on_add` hooks
+    #[inline]
+    pub(crate) fn has_on_add(&self) -> bool {
+        self.flags().contains(ArchetypeFlags::ON_ADD_HOOK)
+    }
+
+    /// Returns true if any of the components in this archetype have `on_insert` hooks
+    #[inline]
+    pub(crate) fn has_on_insert(&self) -> bool {
+        self.flags().contains(ArchetypeFlags::ON_INSERT_HOOK)
+    }
+
+    /// Returns true if any of the components in this archetype have `on_remove` hooks
+    #[inline]
+    pub(crate) fn has_on_remove(&self) -> bool {
+        self.flags().contains(ArchetypeFlags::ON_REMOVE_HOOK)
+    }
 }
 
 /// The next [`ArchetypeId`] in an [`Archetypes`] collection.
@@ -624,7 +670,15 @@ impl Archetypes {
             by_components: Default::default(),
             archetype_component_count: 0,
         };
-        archetypes.get_id_or_insert(TableId::empty(), Vec::new(), Vec::new());
+        // SAFETY: Empty archetype has no components
+        unsafe {
+            archetypes.get_id_or_insert(
+                &Components::default(),
+                TableId::empty(),
+                Vec::new(),
+                Vec::new(),
+            );
+        }
         archetypes
     }
 
@@ -717,8 +771,10 @@ impl Archetypes {
     ///
     /// # Safety
     /// [`TableId`] must exist in tables
-    pub(crate) fn get_id_or_insert(
+    /// `table_components` and `sparse_set_components` must exist in `components`
+    pub(crate) unsafe fn get_id_or_insert(
         &mut self,
+        components: &Components,
         table_id: TableId,
         table_components: Vec<ComponentId>,
         sparse_set_components: Vec<ComponentId>,
@@ -744,6 +800,7 @@ impl Archetypes {
                 let sparse_set_archetype_components =
                     (sparse_start..*archetype_component_count).map(ArchetypeComponentId);
                 archetypes.push(Archetype::new(
+                    components,
                     id,
                     table_id,
                     table_components.into_iter().zip(table_archetype_components),