From 7d5b7ecf382327ee93e872b4088d70b9cec4eb87 Mon Sep 17 00:00:00 2001
From: Matthew Pomes <the10thwiz@gmail.com>
Date: Mon, 12 Sep 2022 15:09:49 -0500
Subject: [PATCH] Add  type to conviently lookup entities by name

---
 crates/bevy_animation/src/lib.rs |  45 +++-------
 crates/bevy_core/src/name.rs     | 148 +++++++++++++++++++++++++++++++
 crates/bevy_gltf/src/loader.rs   |   2 +-
 3 files changed, 162 insertions(+), 33 deletions(-)

diff --git a/crates/bevy_animation/src/lib.rs b/crates/bevy_animation/src/lib.rs
index 0ab68da7c8788..4f8d4c24aba0e 100644
--- a/crates/bevy_animation/src/lib.rs
+++ b/crates/bevy_animation/src/lib.rs
@@ -6,14 +6,17 @@ use std::ops::Deref;
 
 use bevy_app::{App, CoreStage, Plugin};
 use bevy_asset::{AddAsset, Assets, Handle};
-use bevy_core::Name;
+pub use bevy_core::EntityPath;
+use bevy_core::{Name, NameLookup};
 use bevy_ecs::{
     change_detection::DetectChanges,
     entity::Entity,
     prelude::Component,
+    query::QueryEntityError,
     reflect::ReflectComponent,
     schedule::IntoSystemDescriptor,
-    system::{Query, Res},
+    schedule::ParallelSystemDescriptorCoercion,
+    system::{Query, Res, SystemParam},
 };
 use bevy_hierarchy::Children;
 use bevy_math::{Quat, Vec3};
@@ -52,13 +55,6 @@ pub struct VariableCurve {
     pub keyframes: Keyframes,
 }
 
-/// Path to an entity, with [`Name`]s. Each entity in a path must have a name.
-#[derive(Reflect, FromReflect, Clone, Debug, Hash, PartialEq, Eq, Default)]
-pub struct EntityPath {
-    /// Parts of the path
-    pub parts: Vec<Name>,
-}
-
 /// A list of [`VariableCurve`], and the [`EntityPath`] to which they apply.
 #[derive(Reflect, FromReflect, Clone, TypeUuid, Debug, Default)]
 #[uuid = "d81b7179-0448-4eb0-89fe-c067222725bf"]
@@ -187,9 +183,8 @@ pub fn animation_player(
     time: Res<Time>,
     animations: Res<Assets<AnimationClip>>,
     mut animation_players: Query<(Entity, &mut AnimationPlayer)>,
-    names: Query<&Name>,
     mut transforms: Query<&mut Transform>,
-    children: Query<&Children>,
+    lookup: NameLookup,
 ) {
     for (entity, mut player) in &mut animation_players {
         if let Some(animation_clip) = animations.get(&player.animation_clip) {
@@ -208,29 +203,15 @@ pub fn animation_player(
             if elapsed < 0.0 {
                 elapsed += animation_clip.duration;
             }
-            'entity: for (path, curves) in &animation_clip.curves {
+            for (path, curves) in &animation_clip.curves {
                 // PERF: finding the target entity can be optimised
-                let mut current_entity = entity;
-                // Ignore the first name, it is the root node which we already have
-                for part in path.parts.iter().skip(1) {
-                    let mut found = false;
-                    if let Ok(children) = children.get(current_entity) {
-                        for child in children.deref() {
-                            if let Ok(name) = names.get(*child) {
-                                if name == part {
-                                    // Found a children with the right name, continue to the next part
-                                    current_entity = *child;
-                                    found = true;
-                                    break;
-                                }
-                            }
-                        }
-                    }
-                    if !found {
-                        warn!("Entity not found for path {:?} on part {:?}", path, part);
-                        continue 'entity;
+                let current_entity = match lookup.lookup(entity, path) {
+                    Ok(e) => e,
+                    Err(e) => {
+                        warn!("Entity for path {path:?} was not found");
+                        continue;
                     }
-                }
+                };
                 if let Ok(mut transform) = transforms.get_mut(current_entity) {
                     for curve in curves {
                         // Some curves have only one keyframe used to set a transform
diff --git a/crates/bevy_core/src/name.rs b/crates/bevy_core/src/name.rs
index f7fad0ca18e26..eacddacac2255 100644
--- a/crates/bevy_core/src/name.rs
+++ b/crates/bevy_core/src/name.rs
@@ -1,4 +1,13 @@
+use bevy_ecs::{
+    component::Component,
+    entity::Entity,
+    query::QueryEntityError,
+    reflect::ReflectComponent,
+    system::{Query, SystemParam},
+};
 use bevy_ecs::{component::Component, reflect::ReflectComponent};
+use bevy_hierarchy::Children;
+use bevy_reflect::std_traits::ReflectDefault;
 use bevy_reflect::Reflect;
 use bevy_reflect::{std_traits::ReflectDefault, FromReflect};
 use bevy_utils::AHasher;
@@ -150,3 +159,142 @@ impl Deref for Name {
         self.name.as_ref()
     }
 }
+
+/// Path to an entity, with [`Name`]s. Each entity in a path must have a name.
+#[derive(Reflect, FromReflect, Clone, Debug, Hash, PartialEq, Eq, Default)]
+pub struct EntityPath {
+    /// Parts of the path
+    pub parts: Vec<Name>,
+}
+
+/// System param to enable entity lookup of an entity via EntityPath
+#[derive(SystemParam)]
+pub struct NameLookup<'w, 's> {
+    named: Query<'w, 's, (Entity, &'static Name)>,
+    children: Query<'w, 's, &'static Children>,
+}
+
+/// Errors when looking up an entity by name
+#[derive(Debug)]
+pub enum LookupError {
+    /// An entity could not be found, this either means the entity has been
+    /// despawned, or the entity doesn't have the required components
+    Query(QueryEntityError),
+    /// The root node does not have the corrent name
+    // TODO: add expected / found name
+    RootNotFound,
+    /// A child was not found
+    // TODO: add expected name
+    ChildNotFound,
+    /// The name does not uniquely identify an entity
+    // TODO: add name
+    NameNotUnique,
+}
+
+impl From<QueryEntityError> for LookupError {
+    fn from(q: QueryEntityError) -> Self {
+        Self::Query(q)
+    }
+}
+
+impl<'w, 's> NameLookup<'w, 's> {
+    /// Find an entity by entity path, may return an error if the root name isn't unique
+    pub fn lookup_any(&self, path: &EntityPath) -> Result<Entity, LookupError> {
+        let mut path = path.parts.iter();
+        let root_name = path.next().unwrap();
+        let mut root = None;
+        for (entity, name) in self.named.iter() {
+            if root_name == name {
+                if root.is_some() {
+                    return Err(LookupError::NameNotUnique);
+                }
+                root = Some(entity);
+            }
+        }
+        let mut current_node = root.ok_or(LookupError::RootNotFound)?;
+        for part in path {
+            current_node = self.find_child(current_node, part)?;
+        }
+        Ok(current_node)
+    }
+
+    /// Find an entity by the root & entity path
+    pub fn lookup(&self, root: Entity, path: &EntityPath) -> Result<Entity, LookupError> {
+        let mut path = path.parts.iter();
+        let (_, root_name) = self.named.get(root)?;
+        if root_name != path.next().unwrap() {
+            return Err(LookupError::RootNotFound);
+        }
+        let mut current_node = root;
+        for part in path {
+            current_node = self.find_child(current_node, part)?;
+        }
+        Ok(current_node)
+    }
+
+    /// Internal function to get the child of `current_node` that has the name `part`
+    fn find_child(&self, current_node: Entity, part: &Name) -> Result<Entity, LookupError> {
+        let children = self.children.get(current_node)?;
+        let mut ret = Err(LookupError::ChildNotFound);
+        for child in children {
+            if let Ok((_, name)) = self.named.get(*child) {
+                if name == part {
+                    if ret.is_ok() {
+                        return Err(LookupError::NameNotUnique);
+                    }
+                    ret = Ok(*child);
+                }
+            }
+        }
+        ret
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use bevy_app::App;
+    use bevy_ecs::{
+        prelude::Bundle,
+        query::With,
+        schedule::{ParallelSystemDescriptorCoercion, Stage},
+        system::Commands,
+        world::World,
+    };
+    use bevy_hierarchy::BuildChildren;
+
+    use super::*;
+
+    #[derive(Component)]
+    struct Root;
+
+    fn create_heirachy(mut cmds: Commands) {
+        cmds.spawn()
+            .insert(Name::new("root"))
+            .insert(Root)
+            .with_children(|cmds| {
+                cmds.spawn().insert(Name::new("child a"));
+                cmds.spawn().insert(Name::new("child b"));
+                cmds.spawn().insert(Name::new("child c"));
+            });
+    }
+
+    #[test]
+    fn test_lookup() {
+        fn validate(root: Query<Entity, With<Root>>, lookup: NameLookup) {
+            let root = root.single();
+            let a = lookup
+                .lookup(
+                    root,
+                    &EntityPath {
+                        parts: vec![Name::new("root"), Name::new("child a")],
+                    },
+                )
+                .unwrap();
+        }
+
+        let mut app = App::empty();
+        // app.add_startup_stage_after("startup", "", )
+        app.add_startup_system(create_heirachy);
+        app.add_startup_system(validate.after(create_heirachy));
+    }
+}
diff --git a/crates/bevy_gltf/src/loader.rs b/crates/bevy_gltf/src/loader.rs
index 81d793993ac7c..8e51be15bc9b8 100644
--- a/crates/bevy_gltf/src/loader.rs
+++ b/crates/bevy_gltf/src/loader.rs
@@ -204,7 +204,7 @@ async fn load_gltf<'a, 'b>(
                 if let Some((root_index, path)) = paths.get(&node.index()) {
                     animation_roots.insert(root_index);
                     animation_clip.add_curve_to_path(
-                        bevy_animation::EntityPath {
+                        bevy_core::EntityPath {
                             parts: path.clone(),
                         },
                         bevy_animation::VariableCurve {