lyra-engine/lyra-ecs/src/command.rs

use std::{any::Any, cell::RefMut, collections::VecDeque, ptr::{self, NonNull}};

use crate::{system::FnArgFetcher, Access, Bundle, Entities, Entity, World};

/// A Command be used to delay mutation of the world until after this system is ran.
pub trait Command: Any {
    fn as_any_boxed(self: Box<Self>) -> Box<dyn Any>;

    /// Executes the command
    fn run(self, world: &mut World);
}

impl<F> Command for F
where
    F: FnOnce(&mut World) + 'static
{
    fn as_any_boxed(self: Box<Self>) -> Box<dyn Any> {
        self
    }

    fn run(self, world: &mut World) {
        self(world)
    }
}

type RunCommand = unsafe fn(cmd: Box<dyn Command>, world: &mut World);

/// Stores a queue of commands that will get executed after the system is ran.
/// 
/// This struct can be inserted as a resource into the world, and the commands will be
/// executed by the [`GraphExecutor`](crate::system::GraphExecutor) after the system is executed.
#[derive(Default)]
pub struct CommandQueue(VecDeque<(RunCommand, Box<dyn Command>)>);

/// Used in a system to queue up commands that will run right after this system.
/// 
/// This can be used to delay the mutation of the world until after the system is ran. These
/// must be used if you're mutating the world inside a [`ViewState`](crate::query::ViewState).
/// 
/// ```nobuild
/// fn particle_spawner_system(
///     commands: Commands,
///     view: ViewState<(&Campfire, &Transform)>
/// ) -> anyhow::Result<()> {
///     for (campfire, pos) in view.iter() {
///         // If you do not use commands to spawn this, the next iteration
///         // of the view will cause a segfault.
///         commands.spawn((pos, Particle::new(/* ... */)));
///     }
/// 
///     Ok(())
/// }
/// ```
pub struct Commands<'a, 'b> {
    queue: &'b mut CommandQueue,
    entities: &'a mut Entities,
}

impl<'a, 'b> Commands<'a, 'b> {
    pub fn new(queue: &'b mut CommandQueue, world: &'a mut World) -> Self {
        Self {
            queue,
            entities: &mut world.entities,
        }
    }

    /// Add a command to the end of the command queue
    pub fn add<C: Command>(&mut self, cmd: C) {
        let cmd = Box::new(cmd);

        let run_fn = |cmd: Box<dyn Command>, world: &mut World| {
            let cmd = cmd.as_any_boxed()
                .downcast::<C>()
                .unwrap();
            
            cmd.run(world);
        };

        self.queue.0.push_back((run_fn, cmd));
    }

    /// Spawn an entity into the World. See [`World::spawn`]
    pub fn spawn<B: Bundle + 'static>(&mut self, bundle: B) -> Entity {
        let e = self.entities.reserve();

        self.add(move |world: &mut World| {
            world.spawn_into(e, bundle);
        });

        e
    }

    /// Execute all commands in the queue, in order of insertion
    pub fn execute(&mut self, world: &mut World) -> anyhow::Result<()> {
        while let Some((cmd_fn, cmd_ptr)) = self.queue.0.pop_front() {
            unsafe {
                cmd_fn(cmd_ptr, world);
            }
        }

        Ok(())
    }
}

impl FnArgFetcher for Commands<'_, '_> {
    type State = CommandQueue;
    type Arg<'a, 'state> = Commands<'a, 'state>;

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

    unsafe fn get<'a, 'state>(state: &'state mut Self::State, mut world_ptr: ptr::NonNull<World>) -> Self::Arg<'a, 'state> {
        let world = world_ptr.as_mut();
        Commands::new(state, world)
    }

    fn create_state(_: NonNull<World>) -> Self::State {
        CommandQueue::default()
    }

    fn apply_deferred<'a>(mut state: Self::State, mut world_ptr: NonNull<World>) {
        let world = unsafe { world_ptr.as_mut() };

        let mut cmds = Commands::new(&mut state, world);
        // safety: Commands has a mut borrow only to entities in the world
        let world = unsafe { world_ptr.as_mut() };
        cmds.execute(world).unwrap()
    }
}

pub fn execute_deferred_commands(world: &mut World, mut commands: RefMut<Commands>) -> anyhow::Result<()> {
    commands.execute(world)?;

    Ok(())
}

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

    use crate::{system::{GraphExecutor, IntoSystem}, tests::Vec2, Commands, DynTypeId, World};

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

        let spawned_vec = Vec2::rand();
        
        let spawned_vec_cl = spawned_vec.clone();
        let test_sys = move |mut commands: Commands| -> anyhow::Result<()> {
            commands.spawn((spawned_vec_cl.clone(),));

            Ok(())
        };

        let mut graph_exec = GraphExecutor::new();
        graph_exec.insert_system("test", test_sys.into_system(), &[]);
        graph_exec.execute(NonNull::from(&world), true).unwrap();

        assert_eq!(world.entities.len(), 4);
        
        // there's only one archetype
        let arch = world.archetypes.values().next().unwrap();
        let col = arch.get_column(DynTypeId::of::<Vec2>()).unwrap();
        let vec2: Ref<Vec2> = unsafe { col.get(3) };
        assert_eq!(vec2.clone(), spawned_vec);
    }
}