Component Lifecycle Hook & Observer Trigger for replaced values (#14212)

# Objective

Fixes #14202

## Solution

Add `on_replaced` component hook and `OnReplaced` observer trigger

## Testing

- Did you test these changes? If so, how?
  - Updated & added unit tests

---

## Changelog

- Added new `on_replaced` component hook and `OnReplaced` observer
trigger for performing cleanup on component values when they are
overwritten with `.insert()`

Author: BigWingBeat

crates/bevy_ecs/macros/src/component.rs

@@ -58,6 +58,7 @@ pub fn derive_component(input: TokenStream) -> TokenStream {
 
     let on_add = hook_register_function_call(quote! {on_add}, attrs.on_add);
     let on_insert = hook_register_function_call(quote! {on_insert}, attrs.on_insert);
+    let on_replace = hook_register_function_call(quote! {on_replace}, attrs.on_replace);
     let on_remove = hook_register_function_call(quote! {on_remove}, attrs.on_remove);
 
     ast.generics
@@ -76,6 +77,7 @@ pub fn derive_component(input: TokenStream) -> TokenStream {
             fn register_component_hooks(hooks: &mut #bevy_ecs_path::component::ComponentHooks) {
                 #on_add
                 #on_insert
+                #on_replace
                 #on_remove
             }
         }
@@ -86,12 +88,14 @@ pub const COMPONENT: &str = "component";
 pub const STORAGE: &str = "storage";
 pub const ON_ADD: &str = "on_add";
 pub const ON_INSERT: &str = "on_insert";
+pub const ON_REPLACE: &str = "on_replace";
 pub const ON_REMOVE: &str = "on_remove";
 
 struct Attrs {
     storage: StorageTy,
     on_add: Option<ExprPath>,
     on_insert: Option<ExprPath>,
+    on_replace: Option<ExprPath>,
     on_remove: Option<ExprPath>,
 }
 
@@ -110,6 +114,7 @@ fn parse_component_attr(ast: &DeriveInput) -> Result<Attrs> {
         storage: StorageTy::Table,
         on_add: None,
         on_insert: None,
+        on_replace: None,
         on_remove: None,
     };
 
@@ -132,6 +137,9 @@ fn parse_component_attr(ast: &DeriveInput) -> Result<Attrs> {
             } else if nested.path.is_ident(ON_INSERT) {
                 attrs.on_insert = Some(nested.value()?.parse::<ExprPath>()?);
                 Ok(())
+            } else if nested.path.is_ident(ON_REPLACE) {
+                attrs.on_replace = Some(nested.value()?.parse::<ExprPath>()?);
+                Ok(())
             } else if nested.path.is_ident(ON_REMOVE) {
                 attrs.on_remove = Some(nested.value()?.parse::<ExprPath>()?);
                 Ok(())

crates/bevy_ecs/src/archetype.rs

@@ -121,6 +121,7 @@ pub(crate) struct AddBundle {
     /// indicate if the component is newly added to the target archetype or if it already existed
     pub bundle_status: Vec<ComponentStatus>,
     pub added: Vec<ComponentId>,
+    pub mutated: Vec<ComponentId>,
 }
 
 /// This trait is used to report the status of [`Bundle`](crate::bundle::Bundle) components
@@ -205,13 +206,15 @@ impl Edges {
         archetype_id: ArchetypeId,
         bundle_status: Vec<ComponentStatus>,
         added: Vec<ComponentId>,
+        mutated: Vec<ComponentId>,
     ) {
         self.add_bundle.insert(
             bundle_id,
             AddBundle {
                 archetype_id,
                 bundle_status,
                 added,
+                mutated,
             },
         );
     }
@@ -317,10 +320,12 @@ bitflags::bitflags! {
     pub(crate) struct ArchetypeFlags: u32 {
         const ON_ADD_HOOK    = (1 << 0);
         const ON_INSERT_HOOK = (1 << 1);
-        const ON_REMOVE_HOOK = (1 << 2);
-        const ON_ADD_OBSERVER = (1 << 3);
-        const ON_INSERT_OBSERVER = (1 << 4);
-        const ON_REMOVE_OBSERVER = (1 << 5);
+        const ON_REPLACE_HOOK = (1 << 2);
+        const ON_REMOVE_HOOK = (1 << 3);
+        const ON_ADD_OBSERVER = (1 << 4);
+        const ON_INSERT_OBSERVER = (1 << 5);
+        const ON_REPLACE_OBSERVER = (1 << 6);
+        const ON_REMOVE_OBSERVER = (1 << 7);
     }
 }
 
@@ -600,6 +605,12 @@ impl Archetype {
         self.flags().contains(ArchetypeFlags::ON_INSERT_HOOK)
     }
 
+    /// Returns true if any of the components in this archetype have `on_replace` hooks
+    #[inline]
+    pub fn has_replace_hook(&self) -> bool {
+        self.flags().contains(ArchetypeFlags::ON_REPLACE_HOOK)
+    }
+
     /// Returns true if any of the components in this archetype have `on_remove` hooks
     #[inline]
     pub fn has_remove_hook(&self) -> bool {
@@ -622,6 +633,14 @@ impl Archetype {
         self.flags().contains(ArchetypeFlags::ON_INSERT_OBSERVER)
     }
 
+    /// Returns true if any of the components in this archetype have at least one [`OnReplace`] observer
+    ///
+    /// [`OnReplace`]: crate::world::OnReplace
+    #[inline]
+    pub fn has_replace_observer(&self) -> bool {
+        self.flags().contains(ArchetypeFlags::ON_REPLACE_OBSERVER)
+    }
+
     /// Returns true if any of the components in this archetype have at least one [`OnRemove`] observer
     ///
     /// [`OnRemove`]: crate::world::OnRemove

crates/bevy_ecs/src/bundle.rs

@@ -17,7 +17,7 @@ use crate::{
     prelude::World,
     query::DebugCheckedUnwrap,
     storage::{SparseSetIndex, SparseSets, Storages, Table, TableRow},
-    world::{unsafe_world_cell::UnsafeWorldCell, ON_ADD, ON_INSERT},
+    world::{unsafe_world_cell::UnsafeWorldCell, ON_ADD, ON_INSERT, ON_REPLACE},
 };
 
 use bevy_ptr::{ConstNonNull, OwningPtr};
@@ -456,11 +456,13 @@ impl BundleInfo {
         let mut new_sparse_set_components = Vec::new();
         let mut bundle_status = Vec::with_capacity(self.component_ids.len());
         let mut added = Vec::new();
+        let mut mutated = Vec::new();
 
         let current_archetype = &mut archetypes[archetype_id];
         for component_id in self.component_ids.iter().cloned() {
             if current_archetype.contains(component_id) {
                 bundle_status.push(ComponentStatus::Mutated);
+                mutated.push(component_id);
             } else {
                 bundle_status.push(ComponentStatus::Added);
                 added.push(component_id);
@@ -476,7 +478,7 @@ impl BundleInfo {
         if new_table_components.is_empty() && new_sparse_set_components.is_empty() {
             let edges = current_archetype.edges_mut();
             // the archetype does not change when we add this bundle
-            edges.insert_add_bundle(self.id, archetype_id, bundle_status, added);
+            edges.insert_add_bundle(self.id, archetype_id, bundle_status, added, mutated);
             archetype_id
         } else {
             let table_id;
@@ -526,6 +528,7 @@ impl BundleInfo {
                 new_archetype_id,
                 bundle_status,
                 added,
+                mutated,
             );
             new_archetype_id
         }
@@ -665,6 +668,30 @@ impl<'w> BundleInserter<'w> {
         let bundle_info = self.bundle_info.as_ref();
         let add_bundle = self.add_bundle.as_ref();
         let table = self.table.as_mut();
+        let archetype = self.archetype.as_ref();
+
+        // SAFETY: All components in the bundle are guaranteed to exist in the World
+        // as they must be initialized before creating the BundleInfo.
+        unsafe {
+            // SAFETY: Mutable references do not alias and will be dropped after this block
+            let mut deferred_world = self.world.into_deferred();
+
+            deferred_world.trigger_on_replace(
+                archetype,
+                entity,
+                add_bundle.mutated.iter().copied(),
+            );
+            if archetype.has_replace_observer() {
+                deferred_world.trigger_observers(
+                    ON_REPLACE,
+                    entity,
+                    add_bundle.mutated.iter().copied(),
+                );
+            }
+        }
+
+        // SAFETY: Archetype gets borrowed when running the on_replace observers above,
+        // so this reference can only be promoted from shared to &mut down here, after they have been ran
         let archetype = self.archetype.as_mut();
 
         let (new_archetype, new_location) = match &mut self.result {
@@ -1132,7 +1159,7 @@ mod tests {
     struct A;
 
     #[derive(Component)]
-    #[component(on_add = a_on_add, on_insert = a_on_insert, on_remove = a_on_remove)]
+    #[component(on_add = a_on_add, on_insert = a_on_insert, on_replace = a_on_replace, on_remove = a_on_remove)]
     struct AMacroHooks;
 
     fn a_on_add(mut world: DeferredWorld, _: Entity, _: ComponentId) {
@@ -1143,10 +1170,14 @@ mod tests {
         world.resource_mut::<R>().assert_order(1);
     }
 
-    fn a_on_remove<T1, T2>(mut world: DeferredWorld, _: T1, _: T2) {
+    fn a_on_replace<T1, T2>(mut world: DeferredWorld, _: T1, _: T2) {
         world.resource_mut::<R>().assert_order(2);
     }
 
+    fn a_on_remove<T1, T2>(mut world: DeferredWorld, _: T1, _: T2) {
+        world.resource_mut::<R>().assert_order(3);
+    }
+
     #[derive(Component)]
     struct B;
 
@@ -1173,15 +1204,14 @@ mod tests {
         world.init_resource::<R>();
         world
             .register_component_hooks::<A>()
-            .on_add(|mut world, _, _| {
-                world.resource_mut::<R>().assert_order(0);
-            })
+            .on_add(|mut world, _, _| world.resource_mut::<R>().assert_order(0))
             .on_insert(|mut world, _, _| world.resource_mut::<R>().assert_order(1))
-            .on_remove(|mut world, _, _| world.resource_mut::<R>().assert_order(2));
+            .on_replace(|mut world, _, _| world.resource_mut::<R>().assert_order(2))
+            .on_remove(|mut world, _, _| world.resource_mut::<R>().assert_order(3));
 
         let entity = world.spawn(A).id();
         world.despawn(entity);
-        assert_eq!(3, world.resource::<R>().0);
+        assert_eq!(4, world.resource::<R>().0);
     }
 
     #[test]
@@ -1192,7 +1222,7 @@ mod tests {
         let entity = world.spawn(AMacroHooks).id();
         world.despawn(entity);
 
-        assert_eq!(3, world.resource::<R>().0);
+        assert_eq!(4, world.resource::<R>().0);
     }
 
     #[test]
@@ -1201,21 +1231,36 @@ mod tests {
         world.init_resource::<R>();
         world
             .register_component_hooks::<A>()
-            .on_add(|mut world, _, _| {
-                world.resource_mut::<R>().assert_order(0);
-            })
-            .on_insert(|mut world, _, _| {
-                world.resource_mut::<R>().assert_order(1);
-            })
-            .on_remove(|mut world, _, _| {
-                world.resource_mut::<R>().assert_order(2);
-            });
+            .on_add(|mut world, _, _| world.resource_mut::<R>().assert_order(0))
+            .on_insert(|mut world, _, _| world.resource_mut::<R>().assert_order(1))
+            .on_replace(|mut world, _, _| world.resource_mut::<R>().assert_order(2))
+            .on_remove(|mut world, _, _| world.resource_mut::<R>().assert_order(3));
 
         let mut entity = world.spawn_empty();
         entity.insert(A);
         entity.remove::<A>();
         entity.flush();
-        assert_eq!(3, world.resource::<R>().0);
+        assert_eq!(4, world.resource::<R>().0);
+    }
+
+    #[test]
+    fn component_hook_order_replace() {
+        let mut world = World::new();
+        world
+            .register_component_hooks::<A>()
+            .on_replace(|mut world, _, _| world.resource_mut::<R>().assert_order(0))
+            .on_insert(|mut world, _, _| {
+                if let Some(mut r) = world.get_resource_mut::<R>() {
+                    r.assert_order(1);
+                }
+            });
+
+        let entity = world.spawn(A).id();
+        world.init_resource::<R>();
+        let mut entity = world.entity_mut(entity);
+        entity.insert(A);
+        entity.flush();
+        assert_eq!(2, world.resource::<R>().0);
     }
 
     #[test]

crates/bevy_ecs/src/component.rs

@@ -100,6 +100,7 @@ use std::{
 /// Alternatively to the example shown in [`ComponentHooks`]' documentation, hooks can be configured using following attributes:
 /// - `#[component(on_add = on_add_function)]`
 /// - `#[component(on_insert = on_insert_function)]`
+/// - `#[component(on_replace = on_replace_function)]`
 /// - `#[component(on_remove = on_remove_function)]`
 ///
 /// ```
@@ -114,8 +115,8 @@ use std::{
 /// // Another possible way of configuring hooks:
 /// // #[component(on_add = my_on_add_hook, on_insert = my_on_insert_hook)]
 /// //
-/// // We don't have a remove hook, so we can leave it out:
-/// // #[component(on_remove = my_on_remove_hook)]
+/// // We don't have a replace or remove hook, so we can leave them out:
+/// // #[component(on_replace = my_on_replace_hook, on_remove = my_on_remove_hook)]
 /// struct ComponentA;
 ///
 /// fn my_on_add_hook(world: DeferredWorld, entity: Entity, id: ComponentId) {
@@ -280,6 +281,7 @@ pub type ComponentHook = for<'w> fn(DeferredWorld<'w>, Entity, ComponentId);
 pub struct ComponentHooks {
     pub(crate) on_add: Option<ComponentHook>,
     pub(crate) on_insert: Option<ComponentHook>,
+    pub(crate) on_replace: Option<ComponentHook>,
     pub(crate) on_remove: Option<ComponentHook>,
 }
 
@@ -314,6 +316,28 @@ impl ComponentHooks {
             .expect("Component id: {:?}, already has an on_insert hook")
     }
 
+    /// Register a [`ComponentHook`] that will be run when this component is about to be dropped,
+    /// such as being replaced (with `.insert`) or removed.
+    ///
+    /// If this component is inserted onto an entity that already has it, this hook will run before the value is replaced,
+    /// allowing access to the previous data just before it is dropped.
+    /// This hook does *not* run if the entity did not already have this component.
+    ///
+    /// An `on_replace` hook always runs before any `on_remove` hooks (if the component is being removed from the entity).
+    ///
+    /// # Warning
+    ///
+    /// The hook won't run if the component is already present and is only mutated, such as in a system via a query.
+    /// As a result, this is *not* an appropriate mechanism for reliably updating indexes and other caches.
+    ///
+    /// # Panics
+    ///
+    /// Will panic if the component already has an `on_replace` hook
+    pub fn on_replace(&mut self, hook: ComponentHook) -> &mut Self {
+        self.try_on_replace(hook)
+            .expect("Component id: {:?}, already has an on_replace hook")
+    }
+
     /// Register a [`ComponentHook`] that will be run when this component is removed from an entity.
     /// Despawning an entity counts as removing all of its components.
     ///
@@ -351,6 +375,19 @@ impl ComponentHooks {
         Some(self)
     }
 
+    /// Attempt to register a [`ComponentHook`] that will be run when this component is replaced (with `.insert`) or removed
+    ///
+    /// This is a fallible version of [`Self::on_replace`].
+    ///
+    /// Returns `None` if the component already has an `on_replace` hook.
+    pub fn try_on_replace(&mut self, hook: ComponentHook) -> Option<&mut Self> {
+        if self.on_replace.is_some() {
+            return None;
+        }
+        self.on_replace = Some(hook);
+        Some(self)
+    }
+
     /// Attempt to register a [`ComponentHook`] that will be run when this component is removed from an entity.
     ///
     /// This is a fallible version of [`Self::on_remove`].
@@ -442,6 +479,9 @@ impl ComponentInfo {
         if self.hooks().on_insert.is_some() {
             flags.insert(ArchetypeFlags::ON_INSERT_HOOK);
         }
+        if self.hooks().on_replace.is_some() {
+            flags.insert(ArchetypeFlags::ON_REPLACE_HOOK);
+        }
         if self.hooks().on_remove.is_some() {
             flags.insert(ArchetypeFlags::ON_REMOVE_HOOK);
         }

crates/bevy_ecs/src/lib.rs

@@ -61,7 +61,8 @@ pub mod prelude {
             SystemParamFunction,
         },
         world::{
-            EntityMut, EntityRef, EntityWorldMut, FromWorld, OnAdd, OnInsert, OnRemove, World,
+            EntityMut, EntityRef, EntityWorldMut, FromWorld, OnAdd, OnInsert, OnRemove, OnReplace,
+            World,
         },
     };
 }