vidya/crates/renderer/src/device.rs

use std::{borrow::Cow, collections::BTreeMap, ops::Deref, sync::Arc};

use ash::{
    khr,
    prelude::VkResult,
    vk::{self, Handle},
};
use tinyvec::{array_vec, ArrayVec};

use crate::{sync, Instance, PhysicalDevice, Queue};

#[derive(Debug, Default)]
pub struct DeviceQueueFamilies {
    pub(crate) families: Vec<(u32, u32)>,
    pub(crate) graphics: (u32, u32),
    pub(crate) present: (u32, u32),
    pub(crate) async_compute: (u32, u32),
    pub(crate) transfer: (u32, u32),
}

impl DeviceQueueFamilies {
    pub fn swapchain_family_indices(&self) -> ArrayVec<[u32; 2]> {
        let mut indices = array_vec!([u32; 2] => self.graphics.0);

        if self.present.0 != self.graphics.0 {
            indices.push(self.present.0);
        }

        indices
    }
    pub fn graphics_familty(&self) -> u32 {
        self.graphics.0
    }
    pub fn present_familty(&self) -> u32 {
        self.present.0
    }
    pub fn async_compute_familty(&self) -> u32 {
        self.async_compute.0
    }
    pub fn transfer_familty(&self) -> u32 {
        self.transfer.0
    }
}

#[repr(transparent)]
struct DeviceWrapper(ash::Device);

impl Deref for DeviceWrapper {
    type Target = ash::Device;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Drop for DeviceWrapper {
    fn drop(&mut self) {
        unsafe {
            _ = self.0.device_wait_idle();
            self.0.destroy_device(None);
        }
    }
}

#[allow(unused)]
pub struct DeviceInner {
    alloc: vk_mem::Allocator,
    device: DeviceWrapper,
    physical: PhysicalDevice,
    instance: Arc<Instance>,
    swapchain: khr::swapchain::Device,
    debug_utils: ash::ext::debug_utils::Device,
    allocated_queues: BTreeMap<(u32, u32), Queue>,
    // these are resident in allocated_queues, and may in fact be clones of each
    // other, for ease of access
    main_queue: Queue,
    compute_queue: Queue,
    transfer_queue: Queue,
    present_queue: Queue,
    sync_threadpool: sync::SyncThreadpool,
    features: crate::PhysicalDeviceFeatures,
}

impl core::fmt::Debug for DeviceInner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DeviceInner")
            .field("device", &self.device.handle())
            .finish()
    }
}

#[derive(Clone, Debug)]
pub struct Device(Arc<DeviceInner>);
pub type WeakDevice = std::sync::Weak<DeviceInner>;

impl Device {
    pub fn new(
        instance: Arc<Instance>,
        physical: PhysicalDevice,
        mut features: crate::PhysicalDeviceFeatures,
    ) -> VkResult<Self> {
        // we have 4 queues at most: graphics, compute, transfer, present
        let priorities = [1.0f32; 4];

        let queue_infos = physical
            .queue_families
            .families
            .iter()
            .map(|&(family, queues)| {
                vk::DeviceQueueCreateInfo::default()
                    .queue_family_index(family)
                    .queue_priorities(&priorities[..queues as usize])
            })
            .collect::<Vec<_>>();

        let extensions = features
            .device_extensions
            .iter()
            .map(|ext| ext.extension_name.as_ptr())
            .collect::<Vec<_>>();

        let mut features2 = features.features2();
        let device_info = vk::DeviceCreateInfo::default()
            .queue_create_infos(&queue_infos)
            .enabled_extension_names(&extensions)
            .push_next(&mut features2);

        let device = unsafe {
            let device = instance
                .instance
                .create_device(physical.pdev, &device_info, None)?;

            let allocated_queues = queue_infos
                .iter()
                .flat_map(|info| {
                    (0..info.queue_count).map(|i| {
                        (
                            (info.queue_family_index, i),
                            Queue::new(&device, info.queue_family_index, i),
                        )
                    })
                })
                .collect::<BTreeMap<_, _>>();

            let get_queue =
                |(family, index)| allocated_queues.get(&(family, index)).cloned().unwrap();

            let main_queue = get_queue(physical.queue_families.graphics);
            let present_queue = get_queue(physical.queue_families.present);
            let compute_queue = get_queue(physical.queue_families.async_compute);
            let transfer_queue = get_queue(physical.queue_families.transfer);

            let alloc_info =
                vk_mem::AllocatorCreateInfo::new(&instance.instance, &device, physical.pdev);

            let alloc = vk_mem::Allocator::new(alloc_info)?;

            DeviceInner {
                device: DeviceWrapper(device.clone()),
                physical,
                swapchain: khr::swapchain::Device::new(&instance.instance, &device),
                debug_utils: ash::ext::debug_utils::Device::new(&instance.instance, &device),
                instance,
                alloc,
                allocated_queues,
                main_queue,
                present_queue,
                compute_queue,
                transfer_queue,
                features,
                sync_threadpool: sync::SyncThreadpool::new(),
            }
        };

        Ok(Self(Arc::new(device)))
    }
    pub fn sync_threadpool(&self) -> &sync::SyncThreadpool {
        &self.0.sync_threadpool
    }
    pub fn weak(&self) -> WeakDevice {
        Arc::downgrade(&self.0)
    }
    pub fn alloc(&self) -> &vk_mem::Allocator {
        &self.0.alloc
    }
    pub fn dev(&self) -> &ash::Device {
        &self.0.device
    }
    pub fn swapchain(&self) -> &khr::swapchain::Device {
        &self.0.swapchain
    }
    pub fn debug_utils(&self) -> &ash::ext::debug_utils::Device {
        &self.0.debug_utils
    }
    pub fn queue_families(&self) -> &DeviceQueueFamilies {
        &self.0.physical.queue_families
    }
    pub fn phy(&self) -> vk::PhysicalDevice {
        self.0.physical.pdev
    }
    pub fn features(&self) -> &crate::PhysicalDeviceFeatures {
        &self.0.features
    }
    pub fn physical_device(&self) -> &PhysicalDevice {
        &self.0.physical
    }
    pub fn graphics_queue(&self) -> &Queue {
        &self.0.main_queue
    }
    pub fn present_queue(&self) -> &Queue {
        &self.0.present_queue
    }

    pub unsafe fn lock_queues(&self) {
        // this is obviously awful, allocating for this
        self.0
            .allocated_queues
            .values()
            .for_each(|q| core::mem::forget(q.lock()));
    }

    pub unsafe fn unlock_queues(&self) {
        self.0
            .allocated_queues
            .values()
            .for_each(|q| unsafe { q.0.force_unlock() });
    }

    pub fn wait_queue_idle(&self, queue: &Queue) -> VkResult<()> {
        tracing::warn!("locking queue {queue:?} and waiting for idle");

        queue.with_locked(|q| unsafe { self.dev().queue_wait_idle(q) })?;

        tracing::warn!("finished waiting: unlocking queue {queue:?}.");
        Ok(())
    }

    pub fn wait_idle(&self) -> VkResult<()> {
        tracing::warn!("locking all queues and waiting for device to idle");
        unsafe {
            self.lock_queues();
            self.dev().device_wait_idle()?;
            self.unlock_queues();
        }
        tracing::warn!("finished waiting: unlocking all queues.");
        Ok(())
    }

    pub fn debug_name_object<T: vk::Handle>(&self, handle: T, name: &str) -> VkResult<()> {
        let name = std::ffi::CString::new(name.as_bytes()).unwrap_or(c"invalid name".to_owned());
        unsafe {
            self.debug_utils().set_debug_utils_object_name(
                &vk::DebugUtilsObjectNameInfoEXT::default()
                    .object_handle(handle)
                    .object_name(&name),
            )?;
        }
        Ok(())
    }
}

impl AsRef<khr::swapchain::Device> for Device {
    fn as_ref(&self) -> &khr::swapchain::Device {
        &self.0.swapchain
    }
}

impl AsRef<ash::Device> for Device {
    fn as_ref(&self) -> &ash::Device {
        &self.0.device
    }
}

#[allow(dead_code)]
pub struct DeviceAndQueues {
    pub(crate) device: Device,
    pub(crate) main_queue: Queue,
    pub(crate) compute_queue: Queue,
    pub(crate) transfer_queue: Queue,
    pub(crate) present_queue: Queue,
}

impl AsRef<ash::Device> for DeviceAndQueues {
    fn as_ref(&self) -> &ash::Device {
        &self.device.as_ref()
    }
}
impl AsRef<ash::khr::swapchain::Device> for DeviceAndQueues {
    fn as_ref(&self) -> &ash::khr::swapchain::Device {
        &self.device.as_ref()
    }
}

#[derive(Clone)]
pub struct DeviceOwnedDebugObject<T> {
    device: Device,
    object: T,
    #[cfg(debug_assertions)]
    name: Option<Cow<'static, str>>,
}

impl<T: std::fmt::Debug + vk::Handle + Copy> std::fmt::Debug for DeviceOwnedDebugObject<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut fmt = f.debug_struct(core::any::type_name::<T>());

        fmt.field_with("device", |f| {
            write!(f, "0x{:x}", self.device.0.device.handle().as_raw())
        })
        .field_with("handle", |f| write!(f, "0x{:x}", &self.object.as_raw()));

        #[cfg(debug_assertions)]
        {
            fmt.field("name", &self.name);
        }

        fmt.finish()
    }
}

impl<T> DeviceOwnedDebugObject<T> {
    pub fn new(
        device: crate::Device,
        object: T,
        name: Option<Cow<'static, str>>,
    ) -> ash::prelude::VkResult<Self>
    where
        T: vk::Handle + Copy,
    {
        if let Some(name) = name.as_ref() {
            device.debug_name_object(object, name)?;
        }

        Ok(Self {
            device,
            object,
            #[cfg(debug_assertions)]
            name,
        })
    }

    pub fn dev(&self) -> &crate::Device {
        &self.device
    }
    pub fn handle(&self) -> T
    where
        T: Copy,
    {
        self.object
    }
}

pub trait DeviceOwned<T> {
    fn device(&self) -> &Device;
    fn handle(&self) -> T;
}

#[macro_export]
macro_rules! define_device_owned_handle {
    ($(#[$attr:meta])*
        $ty_vis:vis $ty:ident($handle:ty) {
        $($field_vis:vis $field_name:ident : $field_ty:ty),*
        $(,)?
    } $(=> |$this:ident| $dtor:stmt)?) => {
        $(#[$attr])*
            $ty_vis struct $ty {
                inner: crate::device::DeviceOwnedDebugObject<$handle>,
                $(
                    $field_vis $field_name: $field_ty,
                )*
            }

        impl crate::device::DeviceOwned<$handle> for $ty {
            fn device(&self) -> &Device {
                self.inner.dev()
            }
            fn handle(&self) -> $handle {
                self.inner.handle()
            }
        }

        impl $ty {
            fn construct(
                device: crate::device::Device,
                handle: $handle,
                name: Option<::std::borrow::Cow<'static, str>>,
                $($field_name: $field_ty,)*
            ) -> ::ash::prelude::VkResult<Self> {
                Ok(Self {
                    inner: crate::device::DeviceOwnedDebugObject::new(
                        device,
                        handle,
                        name,
                    )?,
                    $($field_name,)*
                })
            }
        }

        $(
            impl Drop for $ty {
                fn drop(&mut self) {
                    #[allow(unused_mut)]
                    let mut $this = self;
                    $dtor
                }
            }
        )?
    };
}