vidya/crates/renderer/src/buffers.rs

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

use ash::{prelude::VkResult, vk};
use itertools::Itertools;
use vk_mem::Alloc;

use crate::{
    define_device_owned_handle,
    device::{DeviceOwned, QueueFlags},
    Device,
};

#[derive(Clone)]
pub struct BufferDesc {
    pub flags: vk::BufferCreateFlags,
    pub name: Option<Cow<'static, str>>,
    pub size: u64,
    pub usage: vk::BufferUsageFlags,
    pub queue_families: QueueFlags,

    pub mem_usage: vk_mem::MemoryUsage,
    pub alloc_flags: vk_mem::AllocationCreateFlags,
}

impl std::hash::Hash for BufferDesc {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.flags.hash(state);
        self.size.hash(state);
        self.usage.hash(state);
        self.queue_families.hash(state);
        self.mem_usage.hash(state);
        self.alloc_flags.bits().hash(state);
    }
}

impl std::fmt::Debug for BufferDesc {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BufferDesc")
            .field("flags", &self.flags)
            .field("name", &self.name)
            .field("size", &self.size)
            .field("usage", &self.usage)
            .field("queue_families", &self.queue_families)
            .field("mem_usage", &self.mem_usage)
            .field_with("alloc_flags", |f| {
                write!(
                    f,
                    "{}",
                    self.alloc_flags
                        .iter_names()
                        .map(|(name, _)| name)
                        .format(" | ")
                )
            })
            .finish()
    }
}

impl Eq for BufferDesc {}
impl PartialEq for BufferDesc {
    fn eq(&self, other: &Self) -> bool {
        self.flags == other.flags
            // for hashmaps, `Eq` may be more strict than `Hash`
            && self.name == other.name
            && self.size == other.size
            && self.usage == other.usage
            && self.queue_families == other.queue_families
            && self.mem_usage == other.mem_usage
            && self.alloc_flags.bits() == other.alloc_flags.bits()
    }
}

impl Default for BufferDesc {
    fn default() -> Self {
        Self {
            flags: Default::default(),
            name: Default::default(),
            size: Default::default(),
            usage: Default::default(),
            queue_families: QueueFlags::empty(),
            alloc_flags: vk_mem::AllocationCreateFlags::empty(),
            mem_usage: vk_mem::MemoryUsage::Auto,
        }
    }
}

define_device_owned_handle! {
    #[derive(Debug)]
    pub Buffer(vk::Buffer) {
        alloc: vk_mem::Allocation,
        size: u64,
    } => |this|  unsafe {
        this.device().clone().alloc().destroy_buffer(this.handle(), &mut this.alloc);
    }
}

impl Eq for Buffer {}
impl PartialEq for Buffer {
    fn eq(&self, other: &Self) -> bool {
        self.inner == other.inner
    }
}

impl Buffer {
    pub fn new(device: Device, desc: BufferDesc) -> VkResult<Self> {
        let queue_families = device.queue_families().family_indices(desc.queue_families);

        let sharing_mode = if queue_families.len() > 1 {
            vk::SharingMode::CONCURRENT
        } else {
            vk::SharingMode::EXCLUSIVE
        };

        let (buffer, allocation) = unsafe {
            device.alloc().create_buffer(
                &vk::BufferCreateInfo::default()
                    .size(desc.size)
                    .usage(desc.usage)
                    .queue_family_indices(&queue_families)
                    .sharing_mode(sharing_mode),
                &vk_mem::AllocationCreateInfo {
                    flags: desc.alloc_flags,
                    usage: desc.mem_usage,
                    ..Default::default()
                },
            )?
        };

        Ok(Self::construct(
            device, buffer, desc.name, allocation, desc.size,
        )?)
    }

    #[allow(dead_code)]
    pub fn map_arc(self: &mut Arc<Self>) -> VkResult<MappedBuffer<'_>> {
        Arc::get_mut(self).map(Self::map).unwrap()
    }

    pub fn map(&mut self) -> VkResult<MappedBuffer<'_>> {
        let bytes = unsafe {
            let data = self.inner.dev().alloc().map_memory(&mut self.alloc)?;
            let slice = core::slice::from_raw_parts_mut(data, self.size as usize);

            slice
        };

        Ok(MappedBuffer { inner: self, bytes })
    }
    pub fn buffer(&self) -> vk::Buffer {
        self.handle()
    }
    pub fn len(&self) -> u64 {
        self.size
    }
}

pub struct MappedBuffer<'a> {
    bytes: &'a mut [u8],
    inner: &'a mut Buffer,
}

impl Drop for MappedBuffer<'_> {
    fn drop(&mut self) {
        unsafe {
            self.inner
                .inner
                .dev()
                .alloc()
                .unmap_memory(&mut self.inner.alloc);
        }
    }
}

impl Deref for MappedBuffer<'_> {
    type Target = [u8];

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

impl DerefMut for MappedBuffer<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.bytes
    }
}