lyra-engine/lyra-ecs/src/system/mod.rs

use std::{ptr::NonNull, marker::PhantomData, cell::{Ref, RefMut}};

use crate::{world::World, Access, ResourceObject, query::{Query, View}};

pub mod graph;
pub use graph::*;

pub mod criteria;
pub use criteria::*;

pub mod batched;
pub use batched::*;

/// A system that does not mutate the world
pub trait System {
    /// A method that indicates the type of access of the world the system requires.
    fn world_access(&self) -> Access;

    /// The actual execution of the system.
    fn execute(&mut self, world: NonNull<World>) -> anyhow::Result<()>;

    /// A setup step of the System, called before `execute` ever runs.
    fn setup(&self, world: NonNull<World>) -> anyhow::Result<()> {
        Ok(())
    }
}

pub trait IntoSystem<T> {
    type System: System;

    fn into_system(self) -> Self::System;
}

pub trait FnArgFetcher {
    type Arg<'a>: FnArg<Fetcher = Self>;

    fn new() -> Self;

    /// Return the appropriate world access if this fetcher gets the world directly.
    /// Return [`Access::None`] if you're only fetching components, or resources.
    fn world_access(&self) -> Access {
        Access::None
    }

    /// Get the arg from the world
    unsafe fn get<'a>(&mut self, world: NonNull<World>) -> Self::Arg<'a>;
}

pub trait FnArg {
    type Fetcher: FnArgFetcher;
}

pub struct FnSystem<F, Args> {
    inner: F,
    args: Args,
}

macro_rules! impl_fn_system_tuple {
    ( $($name: ident),+ ) => (
        #[allow(non_snake_case)]
        impl<F, $($name: FnArgFetcher,)+> System for FnSystem<F, ($($name,)+)>
        where
            F: for<'a> FnMut($($name::Arg<'a>,)+) -> anyhow::Result<()>,
        {
            fn world_access(&self) -> Access {
                todo!()
            }

            fn execute(&mut self, world: NonNull<World>) -> anyhow::Result<()> {
                $(let $name = unsafe { $name::new().get(world) };)+

                (self.inner)($($name,)+)?;
                
                Ok(())
            }
        }

        /* impl<F, $($name: FnArg,)+> IntoSystem for F
        where
            F: FnMut($($name,)+) -> anyhow::Result<()>,
            F: for<'a> FnMut($(<$name::Fetcher as FnArgFetcher>::Arg<'a>,)+) -> anyhow::Result<()>,
        {
            type System = FnSystem<F, ($($name::Fetcher,)+)>;

            fn into_system(self) -> Self::System {
                FnSystem {
                    args: ($($name::Fetcher::new(),)+),
                    inner: self
                }
            }
        } */

        impl<F, $($name: FnArg,)+> IntoSystem<($($name,)+)> for F
        where
            F: FnMut($($name,)+) -> anyhow::Result<()>,
            F: for<'a> FnMut($(<$name::Fetcher as FnArgFetcher>::Arg<'a>,)+) -> anyhow::Result<()>,
        {
            type System = FnSystem<F, ($($name::Fetcher,)+)>;

            fn into_system(self) -> Self::System {
                FnSystem {
                    args: ($($name::Fetcher::new(),)+),
                    inner: self
                }
            }
        }
    );
}

// Hopefully up to 16 arguments in a system is good enough.
impl_fn_system_tuple!{ A }
impl_fn_system_tuple!{ A, B }
impl_fn_system_tuple!{ A, B, C }
impl_fn_system_tuple!{ A, B, C, D }
impl_fn_system_tuple!{ A, B, C, D, E }
impl_fn_system_tuple!{ A, B, C, D, E, F2 }
impl_fn_system_tuple!{ A, B, C, D, E, F2, G }
impl_fn_system_tuple!{ A, B, C, D, E, F2, G, H }
impl_fn_system_tuple!{ A, B, C, D, E, F2, G, H, I }
impl_fn_system_tuple!{ A, B, C, D, E, F2, G, H, I, J }
impl_fn_system_tuple!{ A, B, C, D, E, F2, G, H, I, J, K, L }
impl_fn_system_tuple!{ A, B, C, D, E, F2, G, H, I, J, K, L, M }
impl_fn_system_tuple!{ A, B, C, D, E, F2, G, H, I, J, K, L, M, N }
impl_fn_system_tuple!{ A, B, C, D, E, F2, G, H, I, J, K, L, M, N, O }
impl_fn_system_tuple!{ A, B, C, D, E, F2, G, H, I, J, K, L, M, N, O, P }

/// An ArgFetcher implementation for query [`View`]s
pub struct ViewArgFetcher<Q> {
    query: Q
}

impl<'a, Q: Query> FnArg for View<'a, Q> {
    type Fetcher = ViewArgFetcher<Q>;
}

impl<Q: Query> FnArgFetcher for ViewArgFetcher<Q> {
    type Arg<'a> = View<'a, Q>;

    fn new() -> Self {
        ViewArgFetcher {
            query: Q::new(),
        }
    }

    fn world_access(&self) -> Access {
        todo!()
    }

    unsafe fn get<'a>(&mut self, world: NonNull<World>) -> Self::Arg<'a> {
        let world = &*world.as_ptr();
        let arch = world.archetypes.values().collect();
        let v = View::new(world, self.query, arch);

        v
    }
}

/// An ArgFetcher implementation for borrowing the [`World`].
pub struct WorldArgFetcher;

impl<'a> FnArg for &'a World {
    type Fetcher = WorldArgFetcher;
}

impl FnArgFetcher for WorldArgFetcher {
    type Arg<'a> = &'a World;

    fn new() -> Self {
        WorldArgFetcher
    }

    fn world_access(&self) -> Access {
        Access::Read
    }

    unsafe fn get<'a>(&mut self, world: NonNull<World>) -> Self::Arg<'a> {
        &*world.as_ptr()
    }

}

/// An ArgFetcher implementation for mutably borrowing the [`World`].
pub struct WorldMutArgFetcher;

impl<'a> FnArg for &'a mut World {
    type Fetcher = WorldMutArgFetcher;
}

impl FnArgFetcher for WorldMutArgFetcher {
    type Arg<'a> = &'a mut World;

    fn new() -> Self {
        WorldMutArgFetcher
    }

    fn world_access(&self) -> Access {
        Access::Write
    }

    unsafe fn get<'a>(&mut self, world: NonNull<World>) -> Self::Arg<'a> {
        &mut *world.as_ptr()
    }
}

pub struct ResourceArgFetcher<R: ResourceObject> {
    phantom: PhantomData<fn() -> R>
}

impl<'a, R: ResourceObject> FnArg for Ref<'a, R> {
    type Fetcher = ResourceArgFetcher<R>;
}

impl<R: ResourceObject> FnArgFetcher for ResourceArgFetcher<R> {
    type Arg<'a> = Ref<'a, R>;

    fn new() -> Self {
        ResourceArgFetcher {
            phantom: PhantomData
        }
    }

    unsafe fn get<'a>(&mut self, world: NonNull<World>) -> Self::Arg<'a> {
        let world = world.as_ref();
        world.get_resource::<R>()
    }
}

pub struct ResourceMutArgFetcher<R: ResourceObject> {
    phantom: PhantomData<fn() -> R>
}

impl<'a, R: ResourceObject> FnArg for RefMut<'a, R> {
    type Fetcher = ResourceMutArgFetcher<R>;
}

impl<R: ResourceObject> FnArgFetcher for ResourceMutArgFetcher<R> {
    type Arg<'a> = RefMut<'a, R>;

    fn new() -> Self {
        ResourceMutArgFetcher {
            phantom: PhantomData
        }
    }

    unsafe fn get<'a>(&mut self, world: NonNull<World>) -> Self::Arg<'a> {
        let world = world.as_ref();
        world.get_resource_mut::<R>()
    }
}

#[cfg(test)]
mod tests {
    use std::{ptr::NonNull, cell::RefMut};

    use crate::{tests::{Vec2, Vec3}, world::World, query::{QueryBorrow, View}};
    use super::{System, IntoSystem};

    struct SomeCounter(u32);

    #[test]
    fn simple_system() {
        let mut world = World::new();
        let vecs = &[Vec2::rand(), Vec2::rand(), Vec2::rand(), Vec2::rand()];
        world.spawn((vecs[0],));
        world.spawn((vecs[1],));
        world.spawn((vecs[2],));
        world.spawn((vecs[3],));

        let mut count = 0;

        let test_system = |view: View<QueryBorrow<Vec2>>| -> anyhow::Result<()> {
            let mut vecs = vecs.to_vec();
            for v in view.into_iter() {
                let pos = vecs.iter().position(|vec| *vec == *v)
                    .expect("Failure to find vec inside list");
                vecs.remove(pos);

                println!("Got v at: {:?}", v);
                count += 1;
            }

            Ok(())
        };

        test_system.into_system().execute(NonNull::from(&world)).unwrap();

        assert_eq!(count, 4);
    }

    #[test]
    fn multi_view_system() {
        let mut world = World::new();
        world.spawn((Vec2::rand(), Vec3::rand()));
        world.spawn((Vec2::rand(), Vec3::rand()));
        world.spawn((Vec2::rand(),));
        world.spawn((Vec2::rand(),));

        let mut count = 0;

        let test_system = |view: View<(QueryBorrow<Vec2>, QueryBorrow<Vec3>)>| -> anyhow::Result<()> {
            for (v2, v3) in view.into_iter() {
                println!("Got v2 at '{:?}' and v3 at: '{:?}'", v2, v3);
                count += 1;
            }

            Ok(())
        };

        test_system.into_system().execute(NonNull::from(&world)).unwrap();

        assert_eq!(count, 2);
    }

    #[test]
    fn world_system() {
        let mut world = World::new();
        world.spawn((Vec2::rand(), Vec3::rand()));
        world.spawn((Vec2::rand(), Vec3::rand()));
        world.add_resource(SomeCounter(0));

        let test_system = |world: &World| -> anyhow::Result<()> {
            let mut counter = world.get_resource_mut::<SomeCounter>();
            counter.0 += 10;

            Ok(())
        };

        test_system.into_system().execute(NonNull::from(&world)).unwrap();

        let test_system = |world: &mut World| -> anyhow::Result<()> {
            let mut counter = world.get_resource_mut::<SomeCounter>();
            counter.0 += 10;

            Ok(())
        };

        test_system.into_system().execute(NonNull::from(&world)).unwrap();

        let counter = world.get_resource::<SomeCounter>();
        assert_eq!(counter.0, 20);
    }

    #[test]
    fn function_system() {
        let mut world = World::new();
        world.spawn((Vec2::rand(), Vec3::rand()));
        world.spawn((Vec2::rand(), Vec3::rand()));
        world.add_resource(SomeCounter(0));

        let test_system = |world: &World| -> anyhow::Result<()> {
            let mut counter = world.get_resource_mut::<SomeCounter>();
            counter.0 += 10;

            Ok(())
        };

        test_system.into_system().execute(NonNull::from(&world)).unwrap();

        fn test_system(world: &mut World) -> anyhow::Result<()> {
            let mut counter = world.get_resource_mut::<SomeCounter>();
            counter.0 += 10;

            Ok(())
        }

        test_system.into_system().execute(NonNull::from(&world)).unwrap();

        let counter = world.get_resource::<SomeCounter>();
        assert_eq!(counter.0, 20);
    }

    #[test]
    fn resource_system() {
        let mut world = World::new();
        world.spawn((Vec2::rand(), Vec3::rand()));
        world.spawn((Vec2::rand(), Vec3::rand()));
        world.add_resource(SomeCounter(0));

        let test_system = |mut counter: RefMut<SomeCounter>| -> anyhow::Result<()> {
            counter.0 += 10;

            Ok(())
        };

        test_system.into_system().execute(NonNull::from(&world)).unwrap();

        let counter = world.get_resource::<SomeCounter>();
        assert_eq!(counter.0, 10);
    }

    #[test]
    fn multi_arg_system() {
        let mut world = World::new();
        world.spawn((Vec2::rand(), ));
        world.spawn((Vec2::rand(), ));
        world.add_resource(SomeCounter(0));

        let test_system = |mut counter: RefMut<SomeCounter>, view: View<QueryBorrow<Vec2>>| -> anyhow::Result<()> {
            for v2 in view.into_iter() {
                println!("Got v2 at '{:?}'", v2);
                counter.0 += 1;
            }

            Ok(())
        };

        test_system.into_system().execute(NonNull::from(&world)).unwrap();

        let counter = world.get_resource::<SomeCounter>();
        assert_eq!(counter.0, 2);
    }
}