Support component hooks closures

crates/bevy_ecs/macros/src/component.rs

@@ -179,10 +179,10 @@ pub const ON_REMOVE: &str = "on_remove";
 struct Attrs {
     storage: StorageTy,
     requires: Option<Punctuated<Require, Comma>>,
-    on_add: Option<ExprPath>,
-    on_insert: Option<ExprPath>,
-    on_replace: Option<ExprPath>,
-    on_remove: Option<ExprPath>,
+    on_add: Option<HookFunc>,
+    on_insert: Option<HookFunc>,
+    on_replace: Option<HookFunc>,
+    on_remove: Option<HookFunc>,
 }
 
 #[derive(Clone, Copy)]
@@ -201,6 +201,35 @@ enum RequireFunc {
     Closure(ExprClosure),
 }
 
+enum HookFunc {
+    Path(ExprPath),
+    Closure(ExprClosure),
+}
+
+impl ToTokens for HookFunc {
+    fn to_tokens(&self, tokens: &mut TokenStream2) {
+        match self {
+            HookFunc::Path(path) => {
+                quote!(#path).to_tokens(tokens);
+            }
+            HookFunc::Closure(expr_closure) => {
+                quote!(#expr_closure).to_tokens(tokens);
+            }
+        }
+    }
+}
+
+impl Parse for HookFunc {
+    fn parse(input: syn::parse::ParseStream) -> Result<Self> {
+        if let Ok(func) = input.parse::<ExprClosure>() {
+            Ok(HookFunc::Closure(func))
+        } else {
+            let func = input.parse::<ExprPath>()?;
+            Ok(HookFunc::Path(func))
+        }
+    }
+}
+
 // values for `storage` attribute
 const TABLE: &str = "Table";
 const SPARSE_SET: &str = "SparseSet";
@@ -231,16 +260,16 @@ fn parse_component_attr(ast: &DeriveInput) -> Result<Attrs> {
                     };
                     Ok(())
                 } else if nested.path.is_ident(ON_ADD) {
-                    attrs.on_add = Some(nested.value()?.parse::<ExprPath>()?);
+                    attrs.on_add = Some(nested.value()?.parse::<HookFunc>()?);
                     Ok(())
                 } else if nested.path.is_ident(ON_INSERT) {
-                    attrs.on_insert = Some(nested.value()?.parse::<ExprPath>()?);
+                    attrs.on_insert = Some(nested.value()?.parse::<HookFunc>()?);
                     Ok(())
                 } else if nested.path.is_ident(ON_REPLACE) {
-                    attrs.on_replace = Some(nested.value()?.parse::<ExprPath>()?);
+                    attrs.on_replace = Some(nested.value()?.parse::<HookFunc>()?);
                     Ok(())
                 } else if nested.path.is_ident(ON_REMOVE) {
-                    attrs.on_remove = Some(nested.value()?.parse::<ExprPath>()?);
+                    attrs.on_remove = Some(nested.value()?.parse::<HookFunc>()?);
                     Ok(())
                 } else {
                     Err(nested.error("Unsupported attribute"))
@@ -298,7 +327,7 @@ fn storage_path(bevy_ecs_path: &Path, ty: StorageTy) -> TokenStream2 {
 
 fn hook_register_function_call(
     hook: TokenStream2,
-    function: Option<ExprPath>,
+    function: Option<HookFunc>,
 ) -> Option<TokenStream2> {
     function.map(|meta| quote! { hooks. #hook (#meta); })
 }

examples/ecs/component_hooks.rs

@@ -12,32 +12,53 @@
 //!
 //! - Enforcing structural rules: When you have systems that depend on specific relationships
 //!     between components (like hierarchies or parent-child links) and need to maintain correctness.
+//!
+//! Hooks can be registered during component initialization by using [`Component`] derive macro,
+//! where both functions and closures can be used:
+//!
+//! ```no_run
+//! #[derive(Component)]
+//! #[component(on_add = on_add_hook, on_remove = |deferred_world, entity, component_id| { /* do stuff */ }, on_insert = |_,_,_| { /* ignore parameters and do stuff */ })]
+//! struct MyComponent;
+//!
+//! fn on_add_hook(deferred_world: DeferredWorld<'_>, entity: Entity, component_id: ComponentId) {
+//!     // Do stuff...
+//! }
+//! ```
+//!
+//! Hooks can also be registered during component initialization by manually implementing [`Component`]
+//! and `register_component_hooks`
+//!
+//! ```no_run
+//! struct MyComponent;
+//!
+//! impl Component for MyComponent {
+//!     const STORAGE_TYPE: StorageType = StorageType::Table;
+//!
+//!     fn register_component_hooks(hooks: &mut ComponentHooks) {
+//!         hooks
+//!             .on_add(|_, _, _| { /* do stuff */ })
+//!             .on_insert(|deferred_world, entity, component_id| { /* do stuff */ })
+//!             .on_remove(on_remove_hook)
+//!             .on_replace(on_replace_hook);
+//!     }
+//! }
+//!
+//! fn on_remove_hook(deferred_world: DeferredWorld<'_>, entity: Entity, component_id: ComponentId) {
+//!     // Do stuff...
+//! }
+//!
+//! fn on_replace_hook(deferred_world: DeferredWorld<'_>, entity: Entity, component_id: ComponentId) {
+//!     // Do stuff...
+//! }
+//! ```
 
-use bevy::{
-    ecs::component::{ComponentHooks, StorageType},
-    prelude::*,
-};
+use bevy::prelude::*;
 use std::collections::HashMap;
 
-#[derive(Debug)]
-/// Hooks can also be registered during component initialization by
-/// using [`Component`] derive macro:
-/// ```no_run
-/// #[derive(Component)]
-/// #[component(on_add = ..., on_insert = ..., on_replace = ..., on_remove = ...)]
-/// ```
+#[derive(Component)]
 struct MyComponent(KeyCode);
 
-impl Component for MyComponent {
-    const STORAGE_TYPE: StorageType = StorageType::Table;
-
-    /// Hooks can also be registered during component initialization by
-    /// implementing `register_component_hooks`
-    fn register_component_hooks(_hooks: &mut ComponentHooks) {
-        // Register hooks...
-    }
-}
-
 #[derive(Resource, Default, Debug, Deref, DerefMut)]
 struct MyComponentIndex(HashMap<KeyCode, Entity>);