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

mod node;
pub use node::*;

mod world_transform;
pub use world_transform::*;

use lyra_ecs::{query::{Entities, ViewOne}, relation::ChildOf, Bundle, Component, World};
use lyra_math::Transform;

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

/// 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;

/// 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  {
    pub(crate) world: World,
    root_node: SceneNode,
}

impl SceneGraph {
    /// Create a new SceneGraph with its own ECS World.
    pub fn new() -> Self {
        let world = World::new();
        
        Self::from_world(world)
    }

    /// Create a new SceneGraph inside an existing ECS World.
    pub fn from_world(mut world: World) -> Self {
        let root_en = world.spawn((WorldTransform::default(), Transform::default(), SceneNodeRoot));
        let root = SceneNode::new(None, root_en);

        Self {
            world,
            root_node: root,
        }
    }

    /// 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, bundle: B) -> SceneNode {
        let node = self.root_node.clone();
        self.add_node_under(&node, 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, bundle: B) -> SceneNode {
        world_add_child_node(&mut self.world, parent, bundle)
    }

    /// Insert a component bundle to a SceneNode.
    /// 
    /// See [`lyra_ecs::World::insert`].
    pub fn insert<B: Bundle>(&mut self, node: &SceneNode, bundle: B) {
        self.world.insert(node.entity(), 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, Q>(&self, mut callback: F)
    where
        F: FnMut(&World, &SceneNode, ViewOne<Q::Query>),
        Q: lyra_ecs::query::AsQuery,
    {
        self.traverse_down_from::<F, Q>(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, Q>(&self, start: SceneNode, callback: &mut F)
    where
        F: FnMut(&World, &SceneNode, ViewOne<Q::Query>),
        Q: lyra_ecs::query::AsQuery,
    {
        let v = self.world
            .view::<Entities>()
            .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);
            let v = self.world.view_one::<Q>(e);
            callback(&self.world, &node, v);

            self.traverse_down_from::<F, Q>(node, callback);
        }
    }

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

    /// Retrieve a borrow of the world that backs the Scene
    pub fn world(&self) -> &World {
        &self.world
    }
}

/// 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, bundle: B) -> SceneNode {
    let e = world.spawn(bundle);
    world.insert(e, (SceneNodeFlag,));
    world.add_relation(e, ChildOf, parent.entity());

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

#[cfg(test)]
pub mod tests {
    use lyra_ecs::{query::{filter::{Has, Not}, Entities}, relation::{ChildOf, RelationOriginComponent}, Component};
    use lyra_math::{Transform, Vec3};

    use crate::{lyra_engine, WorldTransform, 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((WorldTransform::default(), Transform::from_translation(v2s[0]), FakeMesh));
        assert!(a.parent(&scene).unwrap() == scene.root_node);
        let b = a.add_node(&mut scene, (WorldTransform::default(), Transform::from_translation(v2s[1]), FakeMesh));
        assert!(b.parent(&scene).unwrap() == a);

        let view = scene.world.view::<(Entities, &mut WorldTransform, &Transform, Not<Has<RelationOriginComponent<ChildOf>>>)>();
        crate::system_update_world_transforms(&scene.world, view).unwrap();

        let mut idx = 0;
        scene.traverse_down::<_, &WorldTransform>(|_, _, v| {
            let pos = v.get().unwrap();
            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;
        });
    }
}