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

use std::{any::Any, cell::RefMut, mem::{self, MaybeUninit}, 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: *mut (), world: Option<&mut World>) -> usize;

#[repr(C, packed)]
struct PackedCommand<T: Command> {
    run: RunCommand,
    cmd: T,
}

/// 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 {
    data: Vec<MaybeUninit<u8>>,
}

impl CommandQueue {
    /// Execute the commands in the queue.
    /// 
    /// If `world` is `None`, the commands will just be dropped and the memory freed.
    fn execute(&mut self, mut world: Option<&mut World>) {
        let range = self.data.as_mut_ptr_range();
        let mut current = range.start;
        let end = range.end;

        while current < end {
            // Retrieve the runner for the command.
            // Safety: current pointer will either be the start of the buffer, or at the start of a new PackedCommand
            let run_fn = unsafe { current.cast::<RunCommand>().read_unaligned() };
            
            // Retrieves the pointer to the command which is just after RunCommand due to PackedCommand.
            // Safety: PackedCommand is repr C and packed, so it will be right after the RunCommand.
            current = unsafe { current.add(mem::size_of::<RunCommand>()) };

            // Now run the command, providing the type erased pointer to the command.
            let read_size = unsafe { run_fn(current.cast(), world.as_deref_mut()) };

            // The pointer is added to so that it is just after the command that was ran.
            // Safety: the RunCommand returns the size of the command
            current = unsafe { current.add(read_size) };
        }

        // Safety: all of the commands were just read from the pointers.
        unsafe { self.data.set_len(0) };
    }
}

impl Drop for CommandQueue {
    fn drop(&mut self) {
        if !self.data.is_empty() {
            println!("CommandQueue has commands but is being dropped");
        }

        self.execute(None);
    }
}

/// 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 [`View`](crate::query::View).
/// 
/// ```nobuild
/// fn particle_spawner_system(
///     commands: Commands,
///     view: View<(&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 run_fn = |cmd_ptr: *mut (), world: Option<&mut World>| {
            // Safety: the pointer is a type-erased pointer to the command. The pointer is read
            // then dropped out of scope, this closure will not be ran again so no use-after-free
            // will occur.
            let cmd: C = unsafe { ptr::read_unaligned(cmd_ptr.cast::<C>()) };
            match world {
                Some(world) => cmd.run(world),
                None => {} // cmd just gets dropped
            }

            // the size of the command must be returned to increment the pointer when applying
            // the command queue.
            mem::size_of::<C>()
        };

        let data = &mut self.queue.data;
        
        // Reserve enough bytes from the vec to store the packed command and its run fn.
        let old_len = data.len();
        data.reserve(mem::size_of::<PackedCommand<C>>());
        
        // Get a pointer to the end of the packed data. Safe since we just reserved enough memory
        // to store this command.
        let end_ptr = unsafe { data.as_mut_ptr().add(old_len) };

        unsafe {
            // write the command and its runner into the buffer
            end_ptr.cast::<PackedCommand<C>>()
                // written unaligned to keep everything packed
                .write_unaligned(PackedCommand {
                    run: run_fn,
                    cmd,
                });

            // we wrote to the vec's buffer without using its api, so we need manually
            // set the length of the vec.
            data.set_len(old_len + mem::size_of::<PackedCommand<C>>());
        }
    }

    /// 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) {
        self.queue.execute(Some(world));
    }
}

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);
    }
}

/// A system for executing deferred commands that are stored in a [`World`] as a Resource.
/// 
/// Commands are usually added inside a system from a [`Commands`] object created just for it
/// as an fn argument. However, there may be cases that commands cannot be added that way, so 
/// they can also be added as a resource and executed later in this system.
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, sync::{atomic::{AtomicU32, Ordering}, Arc}};

    use crate::{system::{GraphExecutor, IntoSystem}, tests::Vec2, CommandQueue, 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);
    }

    /// A test that ensures a command in a command queue will only ever run once.
    #[test]
    fn commands_only_one_exec() {
        let mut world = World::new();

        let counter = Arc::new(AtomicU32::new(0));
        
        let mut queue = CommandQueue::default();
        let mut commands = Commands::new(&mut queue, &mut world);

        let counter_cl = counter.clone();
        commands.add(move |_world: &mut World| {
            counter_cl.fetch_add(1, Ordering::AcqRel);
        });

        queue.execute(Some(&mut world));
        assert_eq!(1, counter.load(Ordering::Acquire));

        queue.execute(Some(&mut world));
        // If its not one, the command somehow was executed.
        // I would be surprised it wouldn't cause some segfault but still increment the counter
        assert_eq!(1, counter.load(Ordering::Acquire));
    }
}