vidya/crates/renderer/src/commands.rs

use std::{future::Future, marker::PhantomData, sync::Arc};

use crate::sync::{self, FenceFuture};

use super::{Device, Queue};
use ash::{prelude::*, vk};

pub struct SingleUseCommandPool {
    device: Device,
    pool: vk::CommandPool,
    queue: Queue,
}

impl Drop for SingleUseCommandPool {
    fn drop(&mut self) {
        unsafe {
            self.device.dev().destroy_command_pool(self.pool, None);
        }
    }
}

impl SingleUseCommandPool {
    pub fn new(device: Device, queue: Queue) -> VkResult<Arc<Self>> {
        let pool_info = vk::CommandPoolCreateInfo::default()
            .queue_family_index(queue.family())
            .flags(vk::CommandPoolCreateFlags::TRANSIENT);

        let pool =
            unsafe { device.dev().create_command_pool(&pool_info, None)? };

        Ok(Arc::new(Self {
            device,
            pool,
            queue,
        }))
    }

    pub fn alloc(self: &Arc<Self>) -> VkResult<SingleUseCommand> {
        SingleUseCommand::new(self.device.clone(), self.clone())
    }

    pub fn queue(&self) -> &Queue {
        &self.queue
    }

    pub fn pool(&self) -> vk::CommandPool {
        self.pool
    }
}

pub struct SingleUseCommand {
    device: Device,
    pool: Arc<SingleUseCommandPool>,
    buffer: vk::CommandBuffer,
}

impl Drop for SingleUseCommand {
    fn drop(&mut self) {
        unsafe {
            self.device
                .dev()
                .free_command_buffers(self.pool.pool(), &[self.buffer])
        };
    }
}

impl SingleUseCommand {
    pub fn new(
        device: Device,
        pool: Arc<SingleUseCommandPool>,
    ) -> VkResult<Self> {
        let buffer = unsafe {
            let alloc_info = vk::CommandBufferAllocateInfo::default()
                .command_buffer_count(1)
                .command_pool(pool.pool())
                .level(vk::CommandBufferLevel::PRIMARY);
            let buffer = device.dev().allocate_command_buffers(&alloc_info)?[0];

            device.dev().begin_command_buffer(
                buffer,
                &vk::CommandBufferBeginInfo::default()
                    .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT),
            )?;

            buffer
        };
        Ok(Self {
            device,
            pool,
            buffer,
        })
    }

    pub fn buffer(&self) -> vk::CommandBuffer {
        self.buffer
    }

    pub fn submit_fence(
        &self,
        wait: Option<(vk::Semaphore, vk::PipelineStageFlags)>,
        signal: Option<vk::Semaphore>,
        fence: Option<vk::Fence>,
    ) -> VkResult<()> {
        unsafe { self.device.dev().end_command_buffer(self.buffer)? };

        let buffers = [self.buffer];
        let mut submit_info =
            vk::SubmitInfo::default().command_buffers(&buffers);

        if let Some(semaphore) = signal.as_ref() {
            // SAFETY: T and [T;1] have the same layout
            submit_info = submit_info.signal_semaphores(unsafe {
                core::mem::transmute::<&vk::Semaphore, &[vk::Semaphore; 1]>(
                    semaphore,
                )
            });
        }
        if let Some((semaphore, stage)) = wait.as_ref() {
            submit_info = submit_info
                .wait_semaphores(core::slice::from_ref(semaphore))
                .wait_dst_stage_mask(core::slice::from_ref(stage));
        }

        let fence = fence.unwrap_or(vk::Fence::null());
        self.pool.queue().with_locked(|queue| unsafe {
            self.device.dev().queue_submit(queue, &[submit_info], fence)
        })?;
        tracing::info!(
            "submitted queue {:?} and fence {:?}",
            self.pool.queue(),
            fence
        );

        Ok(())
    }

    pub fn submit_async<'a>(
        &'a self,
        wait: Option<(vk::Semaphore, vk::PipelineStageFlags)>,
        signal: Option<vk::Semaphore>,
        fence: Arc<sync::Fence>,
    ) -> VkResult<FenceFuture<'a>> {
        let device = self.device.clone();
        self.submit_fence(wait, signal, Some(fence.fence()))?;

        Ok(unsafe { FenceFuture::new(fence) })
    }

    pub fn submit_blocking(
        self,
        wait: Option<(vk::Semaphore, vk::PipelineStageFlags)>,
        signal: Option<vk::Semaphore>,
    ) -> VkResult<()> {
        let fence = Arc::new(sync::Fence::create(self.device.clone())?);
        let future = self.submit_async(wait, signal, fence)?;
        future.block();
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    async fn async_submit(cmd: SingleUseCommand, queue: Queue) {
        cmd.submit_async(
            None,
            None,
            Arc::new(sync::Fence::create(cmd.device.clone()).unwrap()),
        )
        .unwrap()
        .await;
    }
}