lyra-engine/lyra-scene/src/lib.rs

use std::{collections::VecDeque, ops::{Deref, DerefMut}};

use lyra_ecs::{query::Entities, relation::ChildOf, Bundle, Component, Entity, World};

// So we can use lyra_ecs::Component derive macro
pub(crate) mod lyra_engine {
    pub(crate) mod ecs {
        pub use lyra_ecs::*;
    }
}

mod node;
use lyra_math::{Transform, Vec3};
pub use node::*;

/// A flag spawned on all scene node entities
#[derive(Component)]
pub struct SceneNodeFlag;

/// A flag spawned on only the scene root node
#[derive(Component)]
pub struct SceneNodeRoot;

enum MutCow<'a, T> {
    Mut(&'a mut T),
    Owned(T),
}

impl<'a, T> Deref for MutCow<'a, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        match self {
            MutCow::Mut(t) => t,
            MutCow::Owned(t) => t,
        }
    }
}

impl<'a, T> DerefMut for MutCow<'a, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            MutCow::Mut(t) => t,
            MutCow::Owned(t) => t,
        }
    }
}

/// A SceneGraph is a Graph of nodes that represents the hierarchy of a scene.
/// 
/// This SceneGraph is special in the sense that it is literally just an ECS world with methods
/// implemented for it that make it easier to use for a SceneGraph.
//#[derive(Default)]
pub struct SceneGraph<'a>  {
    pub(crate) world: MutCow<'a, World>,
    root_node: SceneNode,
}

impl<'a> SceneGraph<'a> {
    /// Create a new SceneGraph with its own ECS World.
    pub fn new() -> Self {
        let mut world = World::new();
        
        let e = world.spawn((Transform::from_translation(Vec3::new(0.0, 0.0, 0.0)), SceneNodeRoot));
        let root = SceneNode::new(None, e);

        Self {
            world: MutCow::Owned(world),
            root_node: root
        }
    }

    /// Retrieve a SceneGraph from an ECS World.
    /// 
    /// Returns `None` if the `root_entity` was not created from a `SceneGraph` that was later
    /// inserted into another world with [`SceneGraph::into_world`].
    pub fn from_world(world: &'a mut World, root_entity: Entity) -> Option<Self> {
        if world.view_one::<(&SceneNodeRoot, &Transform)>(root_entity).get().is_none() {
            None
        } else {
            Some(Self {
                world: MutCow::Mut(world),
                root_node: SceneNode::new(None, root_entity),
            })
        }
    }

    /// Create a new SceneGraph inside an existing ECS World.
    pub fn new_from_world(world: &'a mut World) -> Self {
        let root_en = world.spawn((Transform::from_translation(Vec3::new(0.0, 0.0, 0.0)), SceneNodeRoot));
        let root = SceneNode::new(None, root_en);

        Self {
            world: MutCow::Mut(world),
            root_node: root,
        }
    }

    /// Consume the SceneGraph, inserting its internal world into the provided world.
    /// 
    /// The `into` World will contain all nodes of the SceneGraph. You could iterate through the
    /// SceneGraph manually if you'd like. The root entity has a [`SceneNodeRoot`] flag component,
    /// and a Transform component. The children nodes have a [`SceneNodeFlag`] component and a 
    /// Transform component. The transforms of the children will be local, you must find its world
    /// transform by traversing up the hierarchy of the scene manually with the `ChildOf` relation.
    /// 
    /// Returns: the Entity of the root node in the provided world.
    pub fn into_world(mut self, into: &mut World) -> Entity {
        // first insert the root entity into the World.
        let v = self.world.view_one::<&Transform>(self.root_node.entity());
        let pos = *v.get().unwrap();
        let new_root = into.spawn((pos, SceneNodeRoot));

        // now process the children of the root node.
        Self::traverse_inserting_into(&mut self.world, self.root_node.entity(), new_root, into);
        new_root
    }

    /// Recursively traverse the `from` world, starting at `node_en`, inserting the
    /// children entity into the world.
    /// 
    /// Parameters:
    /// * `scene_world` - The world that the SceneGraph exists in
    /// * `scene_en` - The entity of the parent entity inside of the `scene_world`.
    /// * `parent_en` - The entity of the parent entity inside of the `into` world.
    /// * `into` - The world to insert the SceneGraph world into.
    fn traverse_inserting_into(scene_world: &mut World, scene_en: Entity, parent_en: Entity, into: &mut World) {
        let v = scene_world.view::<(Entities, &Transform)>()
            .relates_to::<ChildOf>(scene_en);

        // unfortunately, the borrow checker exists, and wasn't happy that `scene_world` was
        // immutably borrowed with the view (v), and being mutably borrowed in the
        // recursive call.
        let mut child_entities = VecDeque::new();

        for ((child_scene_en, pos), _rel) in v.iter() {
            let into_en = into.spawn((*pos, SceneNodeFlag));
            into.add_relation(into_en, ChildOf, parent_en);

            child_entities.push_back((child_scene_en, into_en));
        }

        while let Some((child_scene_en, into_en)) = child_entities.pop_front() {
            Self::traverse_inserting_into(scene_world, child_scene_en, into_en, into);
        }
    }

    /// Adds a node to the root node of the SceneGraph.
    /// 
    /// The spawned entity will have a `ChildOf` relation targeting the root node, the
    /// `SceneNodeFlag` component is also added to the entity.
    pub fn add_node<B: Bundle>(&mut self, local_transform: Transform, bundle: B) -> SceneNode {
        let node = self.root_node.clone();
        self.add_node_under(&node, local_transform, bundle)
    }

    /// Add a node under a parent node.
    /// 
    /// The spawned entity will have a `ChildOf` relation targeting the provided parent node,
    /// the `SceneNodeFlag` component is also added to the entity.
    pub fn add_node_under<B: Bundle>(&mut self, parent: &SceneNode, local_transform: Transform, bundle: B) -> SceneNode {
        world_add_child_node(&mut self.world, parent, local_transform, bundle)
    }

    pub fn add_empty_node_under(&mut self, parent: &SceneNode, local_transform: Transform) -> SceneNode {
        let e = self.world.spawn((SceneNodeFlag, local_transform));
        self.world.add_relation(e, ChildOf, parent.entity());

        SceneNode::new(Some(parent.entity()), e)
    }

    /// Traverses down the SceneGraph, calling `callback` with each SceneNode and its world transform.
    /// 
    /// The traversal does not include the root scene node.
    pub fn traverse_down<F>(&self, mut callback: F)
    where
        F: FnMut(&SceneNode, Transform),
    {
        self.traverse_down_from(self.root_node.clone(), &mut callback);
    }

    /// Recursively Traverses down the SceneGraph from a starting node, calling `callback` with each
    /// SceneNode and its world transform.
    fn traverse_down_from<F>(&self, start: SceneNode, callback: &mut F)
    where
        F: FnMut(&SceneNode, Transform),
    {
        let v = self.world.view::<(Entities, &Transform)>()
            .relates_to::<ChildOf>(start.entity());

        for ((e, _), _rel) in v.iter() {
            let node = SceneNode::new(Some(start.entity()), e);
            let world_pos = node.world_transform(self);
            callback(&node, world_pos);

            self.traverse_down_from(node, callback);
        }
    }

    pub fn root_node(&self) -> SceneNode {
        self.root_node.clone()
    }
}

/// Add a node under a parent node.
/// 
/// The spawned entity will have a `ChildOf` relation targeting the provided parent node,
/// the `SceneNodeFlag` component is also added to the entity.
pub(crate) fn world_add_child_node<B: Bundle>(world: &mut World, parent: &SceneNode, local_transform: Transform, bundle: B) -> SceneNode {
    let e = world.spawn(bundle);
    world.insert(e, (SceneNodeFlag, local_transform));
    world.add_relation(e, ChildOf, parent.entity());

    SceneNode::new(Some(parent.entity()), e)
}

#[cfg(test)]
pub mod tests {
    use lyra_ecs::Component;
    use lyra_math::{Transform, Vec3};

    use crate::{lyra_engine, SceneGraph};

    #[derive(Component)]
    pub struct FakeMesh;

    #[test]
    fn single_node_hierarchy() {
        let mut scene = SceneGraph::new();

        let a = scene.add_node(Transform::from_translation(Vec3::new(10.0, 10.0, 10.0)), FakeMesh);
        assert!(a.parent(&scene).unwrap() == scene.root_node);
    }

    #[test]
    fn double_node_hierarchy() {
        let mut scene = SceneGraph::new();

        let a = scene.add_node(Transform::from_translation(Vec3::new(10.0, 10.0, 10.0)), FakeMesh);
        assert!(a.parent(&scene).unwrap() == scene.root_node);

        let b = a.add_node(&mut scene, Transform::from_translation(Vec3::new(50.0, 50.0, 50.0)), FakeMesh);
        assert!(b.parent(&scene).unwrap() == a);
    }

    #[test]
    fn traverse_down() {
        let v2s = vec![Vec3::new(10.0, 10.0, 10.0), Vec3::new(50.0, 50.0, 50.0)];

        let mut scene = SceneGraph::new();

        let a = scene.add_node(Transform::from_translation(v2s[0]), FakeMesh);
        assert!(a.parent(&scene).unwrap() == scene.root_node);
        let b = a.add_node(&mut scene, Transform::from_translation(v2s[1]), FakeMesh);
        assert!(b.parent(&scene).unwrap() == a);

        let mut idx = 0;
        scene.traverse_down(|_e, pos| {
            if idx == 0 {
                assert_eq!(pos, Transform::from_translation(v2s[idx]));
            } else if idx == 1 {
                let t = v2s.iter().sum();
                assert_eq!(pos, Transform::from_translation(t));
            }

            idx += 1;
        });
    }

    /* #[test]
    fn inserting_and_from_world() {
        let v2s = vec![Vec3::new(10.0, 10.0, 10.0), Vec3::new(50.0, 50.0, 50.0)];

        let mut scene = SceneGraph::new();

        let a = scene.add_node(Transform::from_translation(v2s[0]), FakeMesh);
        assert!(a.parent(&scene).unwrap() == scene.root_node);
        let b = a.add_node(&mut scene, Transform::from_translation(v2s[1]), FakeMesh);
        assert!(b.parent(&scene).unwrap() == a);

        let mut other_world = World::new();
        let root = scene.into_world(&mut other_world);
        
        // check all of the entities inside of the World
        let (root_pos, _) = other_world.view_one::<(&Transform, &SceneNodeRoot)>(root).get().unwrap();
        assert_eq!(*root_pos, Transform::from_xyz(0.0, 0.0, 0.0));

        let ((child_en, child_pos), _)  = other_world.view::<(Entities, &Transform)>()
            .relates_to::<ChildOf>(root).iter().next().unwrap();
        assert_eq!(*child_pos, Transform::from_translation(v2s[0]));

        let ((_, childchild_pos), _)  = other_world.view::<(Entities, &Transform)>()
            .relates_to::<ChildOf>(child_en).iter().next().unwrap();
        assert_eq!(*childchild_pos, Transform::from_translation(v2s[1]));
        
        drop(root_pos);
        drop(child_pos);
        drop(childchild_pos);

        // Now get the SceneGraph inside the World and use the nice utility tools to traverse it.
        let scene = SceneGraph::from_world(&mut other_world, root)
            .expect("Failed to get SceneGraph from World!");

        let mut idx = 0;
        scene.traverse_down(|_e, pos| {
            if idx == 0 {
                assert_eq!(pos, Transform::from_translation(v2s[idx]));
            } else if idx == 1 {
                let t = v2s.iter().sum();
                assert_eq!(pos, Transform::from_translation(t));
            }

            idx += 1;
        });
    } */
}