Add way to insert components to loaded scenes

This adds the `hook` field to `SceneBundle`. The `hook` is a
`SceneHook`, which accepts a closure, that is ran once per entity in the
scene. The closure can use the `EntityRef` argument to check components
the entity has, and use the `EntityCommands` to add components and
children to the entity.

This code is adapted from https://github.com/nicopap/bevy-scene-hook

Co-authored-by: Hennadii Chernyshchyk <genaloner@gmail.com>

crates/bevy_scene/src/bundle.rs

@@ -9,7 +9,7 @@ use bevy_ecs::{
 };
 use bevy_transform::components::{GlobalTransform, Transform};
 
-use crate::{DynamicScene, InstanceId, Scene, SceneSpawner};
+use crate::{DynamicScene, InstanceId, Scene, SceneHook, SceneSpawner};
 
 /// [`InstanceId`] of a spawned scene. It can be used with the [`SceneSpawner`] to
 /// interact with the spawned scene.
@@ -26,6 +26,7 @@ pub struct SceneBundle {
     pub scene: Handle<Scene>,
     pub transform: Transform,
     pub global_transform: GlobalTransform,
+    pub hook: SceneHook,
 }
 
 /// A component bundle for a [`DynamicScene`] root.
@@ -38,6 +39,7 @@ pub struct DynamicSceneBundle {
     pub scene: Handle<DynamicScene>,
     pub transform: Transform,
     pub global_transform: GlobalTransform,
+    pub hook: SceneHook,
 }
 
 /// System that will spawn scenes from [`SceneBundle`].

crates/bevy_scene/src/hook.rs

@@ -0,0 +1,174 @@
+//! Systems to insert components on loaded scenes.
+//!
+//! Please see the [`SceneHook`] documentation for detailed examples.
+
+use bevy_ecs::{
+    component::Component,
+    entity::Entity,
+    prelude::{Without, World},
+    system::{Commands, EntityCommands, Query, Res},
+    world::EntityRef,
+};
+
+use crate::{SceneInstance, SceneSpawner};
+
+/// Marker Component for scenes that were succesfully hooked.
+#[derive(Component)]
+#[non_exhaustive]
+pub struct SceneHooked;
+
+/// Add this as a component to any entity to trigger `hook`'s
+/// [`Hook::hook_entity`] method when the scene is loaded.
+///
+/// You can use it to add your own non-serializable components to entites
+/// present in a scene file.
+///
+/// A typical usage is adding animation, physics collision data or marker
+/// components to a scene spawned from a file format that do not support it.
+///
+/// # Example
+///
+///  ```rust
+/// # use bevy_ecs::{system::Res, component::Component, system::Commands};
+/// # use bevy_asset::AssetServer;
+/// # use bevy_utils::default;
+/// use bevy_scene::{Hook, SceneHook, SceneBundle};
+/// # #[derive(Component)]
+/// # struct Name; impl Name { fn as_str(&self) -> &str { todo!() } }
+/// enum PileType { Drawing }
+///
+/// #[derive(Component)]
+/// struct Pile(PileType);
+///
+/// #[derive(Component)]
+/// struct Card;
+///
+/// fn load_scene(mut cmds: Commands, asset_server: Res<AssetServer>) {
+///     cmds.spawn_bundle(SceneBundle {
+///         scene: asset_server.load("scene.glb#Scene0"),
+///         hook: SceneHook::new_fn(|entity, cmds| {
+///             match entity.get::<Name>().map(|t|t.as_str()) {
+///                 Some("Pile") => cmds.insert(Pile(PileType::Drawing)),
+///                 Some("Card") => cmds.insert(Card),
+///                 _ => cmds,
+///             };
+///         }),
+///         ..default()
+///     });
+/// }
+/// ```
+#[derive(Component)]
+pub struct SceneHook {
+    hook: Box<dyn Hook>,
+}
+impl Default for SceneHook {
+    fn default() -> Self {
+        Self { hook: Box::new(()) }
+    }
+}
+impl<T: Hook> From<T> for SceneHook {
+    fn from(hook: T) -> Self {
+        Self::new(hook)
+    }
+}
+impl SceneHook {
+    /// Add a hook to a scene, to run for each entities when the scene is
+    /// loaded, closures implement `Hook`.
+    ///
+    ///  You can also implement [`Hook`] on your own types and provide one. Note
+    ///  that strictly speaking, you might as well pass a closure. Please check
+    ///  the [`Hook`] trait documentation for details.
+    pub fn new<T: Hook>(hook: T) -> Self {
+        let hook = Box::new(hook);
+        Self { hook }
+    }
+
+    /// Same as [`Self::new`] but with type bounds to make it easier to
+    /// use a closure.
+    pub fn new_fn<F: Fn(&EntityRef, &mut EntityCommands) + Send + Sync + 'static>(hook: F) -> Self {
+        Self::new(hook)
+    }
+
+    /// Add a closure with component parameter as hook.
+    ///
+    /// This is useful if you only care about a specific component to identify
+    /// individual entities of your scene, rather than every possible components.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// # use bevy_ecs::{
+    ///     world::EntityRef, component::Component,
+    ///     system::{Commands, Res, EntityCommands}
+    /// # };
+    /// # use bevy_asset::{AssetServer, Handle};
+    /// # use bevy_utils::default;
+    /// # use bevy_scene::Scene;
+    /// use bevy_scene::{Hook, SceneHook, SceneBundle};
+    /// # #[derive(Component)] struct Name;
+    /// # type DeckData = Scene;
+    /// #[derive(Clone)]
+    /// struct DeckAssets { player: Handle<DeckData>, oppo: Handle<DeckData> }
+    ///
+    /// fn hook(decks: &DeckAssets, name: &Name, cmds: &mut EntityCommands) {}
+    /// fn load_scene(mut cmds: Commands, decks: Res<DeckAssets>, assets: Res<AssetServer>) {
+    ///     let decks = decks.clone();
+    ///     cmds.spawn_bundle(SceneBundle {
+    ///         scene: assets.load("scene.glb#Scene0"),
+    ///         hook: SceneHook::new_comp(move |name, cmds| hook(&decks, name, cmds)),
+    ///         ..default()
+    ///     });
+    /// }
+    /// ```
+    pub fn new_comp<C, F>(hook: F) -> Self
+    where
+        F: Fn(&C, &mut EntityCommands) + Send + Sync + 'static,
+        C: Component,
+    {
+        let hook = move |e: &EntityRef, cmds: &mut EntityCommands| match e.get::<C>() {
+            Some(comp) => hook(comp, cmds),
+            None => {}
+        };
+        Self::new(hook)
+    }
+}
+
+/// Handle adding components to entites named in a loaded scene.
+///
+/// The [`hook_entity`][Hook::hook_entity] method is called once per Entity
+/// added in a scene in the [`run_hooks`] system.
+pub trait Hook: Send + Sync + 'static {
+    /// Add [`Component`]s or do anything with entity in the spawned scene
+    /// refered by `entity_ref`.
+    ///
+    /// This runs once for all entities in the spawned scene, once loaded.
+    fn hook_entity(&self, entity_ref: &EntityRef, commands: &mut EntityCommands);
+}
+
+/// Run once [`SceneHook`]s added to [`SceneBundle`](crate::SceneBundle) or
+/// [`DynamicSceneBundle`](crate::DynamicSceneBundle) when the scenes are loaded.
+pub fn run_hooks(
+    unloaded_instances: Query<(Entity, &SceneInstance, &SceneHook), Without<SceneHooked>>,
+    scene_manager: Res<SceneSpawner>,
+    world: &World,
+    mut cmds: Commands,
+) {
+    for (entity, instance, hooked) in unloaded_instances.iter() {
+        if let Some(entities) = scene_manager.iter_instance_entities(**instance) {
+            for entity_ref in entities.filter_map(|e| world.get_entity(e)) {
+                let mut cmd = cmds.entity(entity_ref.id());
+                hooked.hook.hook_entity(&entity_ref, &mut cmd);
+            }
+            cmds.entity(entity).insert(SceneHooked);
+        }
+    }
+}
+impl<F: Fn(&EntityRef, &mut EntityCommands) + Send + Sync + 'static> Hook for F {
+    fn hook_entity(&self, entity_ref: &EntityRef, commands: &mut EntityCommands) {
+        (self)(entity_ref, commands)
+    }
+}
+// This is useful for the `Default` implementation of `SceneBundle`
+impl Hook for () {
+    fn hook_entity(&self, _: &EntityRef, _: &mut EntityCommands) {}
+}

crates/bevy_scene/src/lib.rs

@@ -1,12 +1,14 @@
 mod bundle;
 mod dynamic_scene;
+mod hook;
 mod scene;
 mod scene_loader;
 mod scene_spawner;
 pub mod serde;
 
 pub use bundle::*;
 pub use dynamic_scene::*;
+pub use hook::*;
 pub use scene::*;
 pub use scene_loader::*;
 pub use scene_spawner::*;
@@ -33,6 +35,7 @@ impl Plugin for ScenePlugin {
                 CoreStage::PreUpdate,
                 scene_spawner_system.exclusive_system().at_end(),
             )
+            .add_system_to_stage(CoreStage::PreUpdate, run_hooks)
             // Systems `*_bundle_spawner` must run before `scene_spawner_system`
             .add_system_to_stage(CoreStage::PreUpdate, scene_spawner);
     }