Tag scene skeleton if it exists

studio/src/algo.rs

@@ -0,0 +1,89 @@
+use bevy::{
+    prelude::*,
+    utils::{EntityHashMap, EntityHashSet},
+};
+
+/// Find the parent entity of `target`. Returns `None` if `target` has no
+/// parents.
+///
+/// - Time: Constant
+/// - Space: Constant
+pub(crate) fn find_parent(target: Entity, q_parents: &Query<Option<&Parent>>) -> Option<Entity> {
+    q_parents
+        .get(target)
+        .ok()
+        .and_then(|parent| parent.map(|p| **p))
+}
+
+/// Find the nearest common ancestor entity of a set of `entities`. Returns
+/// `None` if the provided `entities` do not have a common ancestor.
+///
+/// - Time: O(D) where D = the depth of the deepest member of `entities`
+/// - Space: O(N) where N = the length of `entities`
+pub(crate) fn nearest_common_ancestor(
+    entities: &[Entity],
+    q_parents: &Query<Option<&Parent>>,
+) -> Option<Entity> {
+    // Find the depth of each entity in the tree
+    let mut depths: EntityHashMap<Entity, usize> = EntityHashMap::default();
+    for &ent in entities {
+        depths.insert(ent, find_depth(ent, q_parents));
+    }
+
+    // Find the lowest depth of all the entities
+    let Some(lowest_depth) = depths.values().copied().reduce(usize::min) else {
+        return None;
+    };
+
+    // Traverse up the tree from each entity until we find each one's ancestor
+    // at the level of lowest_depth
+    let mut equalized: EntityHashSet<Entity> = EntityHashSet::default();
+    for ent in entities {
+        let mut depth = depths[ent];
+        let mut result = *ent;
+
+        while depth > lowest_depth {
+            result = find_parent(result, q_parents).unwrap();
+            depth -= 1;
+        }
+
+        equalized.insert(result);
+    }
+
+    // Walk up the tree from each of those entities until we find the common ancestor
+    let mut parents: EntityHashSet<Entity> = EntityHashSet::default();
+    loop {
+        for &ent in equalized.iter() {
+            if let Some(parent) = find_parent(ent, q_parents) {
+                parents.insert(parent);
+            } else {
+                return None;
+            }
+        }
+
+        if parents.len() == 1 {
+            break parents.iter().copied().next();
+        }
+
+        // Reset the loop:
+        //   equalized = parents
+        //   parents = {}
+        equalized.clear();
+        equalized.extend(parents.iter());
+
+        parents.clear();
+    }
+}
+
+fn find_depth(mut target: Entity, q_parents: &Query<Option<&Parent>>) -> usize {
+    let mut result = 0;
+    let mut parent = find_parent(target, q_parents);
+
+    while let Some(p) = parent {
+        result += 1;
+        target = p;
+        parent = find_parent(target, q_parents);
+    }
+
+    result
+}

studio/src/lib.rs

@@ -8,6 +8,7 @@ use space_editor::prelude::space_undo::UndoPlugin;
 #[cfg(feature = "wasm")]
 use wasm_bindgen::prelude::*;
 
+mod algo;
 mod bvh;
 mod camera_controller;
 mod debug;

studio/src/scene.rs

@@ -18,6 +18,7 @@ use space_editor::prelude::{space_undo::OneFrameUndoIgnore, PrefabMarker};
 use wasm_bindgen::prelude::wasm_bindgen;
 
 use crate::{
+    algo,
     bvh::{Bvh, CommandsSpawnBvhExt},
     camera_controller::{CameraController, CameraControllerBundle},
     debug::DebugPoint,
@@ -58,7 +59,7 @@ impl Plugin for ScenePlugin {
             .add_systems(
                 OnEnter(SceneState::Active),
                 (
-                    tag_scene_elements,
+                    (tag_scene_elements, apply_deferred, tag_skeleton).chain(),
                     spawn_main_scene_light,
                     frame_scene_in_view,
                 )
@@ -113,6 +114,9 @@ pub struct SceneElement;
 #[derive(Component)]
 pub struct SkeletalJoint(pub Color);
 
+#[derive(Component)]
+pub struct Skeleton;
+
 #[derive(Component)]
 pub struct SceneRoot;
 
@@ -128,7 +132,7 @@ pub struct BindPoseGlobal(pub Affine3A);
 #[derive(Component, Clone, Copy, Debug, Deref, DerefMut)]
 pub struct InvBindPose(pub Affine3A);
 
-fn initialize(mut cmd: Commands, r_assets: Res<AssetServer>) {
+fn initialize(mut cmd: Commands) {
     // Camera, graphics settings
     let focus = Vec3::ZERO;
     let camera_xform = Transform::from_xyz(-2.5, 4.5, 5.0).looking_at(focus, Vec3::Y);
@@ -323,6 +327,17 @@ fn tag_scene_elements(
     }
 }
 
+fn tag_skeleton(
+    mut cmd: Commands,
+    q_joints: Query<Entity, With<SkeletalJoint>>,
+    q_parents: Query<Option<&Parent>>,
+) {
+    let joints: Vec<_> = q_joints.iter().collect();
+    if let Some(root) = algo::nearest_common_ancestor(&joints, &q_parents) {
+        cmd.entity(root).insert(Skeleton);
+    }
+}
+
 fn spawn_main_scene_light(
     mut cmd: Commands,
     r_scene_spawner: Res<SceneSpawner>,