vidya/crates/renderer_resources/src/buffers.rs

use std::{borrow::Cow, sync::Arc};

use ash::vk;
use renderer_device::queue::QueueFamilyIndices;

use crate::{
    Device,
    device::{
        Allocation, AllocationDesc, AllocationStrategy, DeviceObject, MemoryLocation, QueueFlags,
    },
};

#[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_location: MemoryLocation,
    pub alloc_scheme: AllocationStrategy,
}

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_location.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_location", &self.mem_location)
            .field("alloc_scheme", &self.alloc_scheme)
            .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.size == other.size
            && self.usage == other.usage
            && self.queue_families == other.queue_families
            && self.mem_location == other.mem_location
    }
}

impl Default for BufferDesc {
    fn default() -> Self {
        Self {
            flags: Default::default(),
            name: Default::default(),
            size: Default::default(),
            usage: Default::default(),
            queue_families: QueueFlags::empty(),
            mem_location: MemoryLocation::Unknown,
            alloc_scheme: AllocationStrategy::default(),
        }
    }
}

#[derive(Debug)]
pub struct Buffer {
    buffer: BufferHandle,
    pub desc: BufferDesc,
    alloc: Allocation,
}

#[derive(Debug)]
enum BufferHandle {
    Owned(DeviceObject<vk::Buffer>),
    #[cfg(feature = "testing")]
    Testing(vk::Buffer),
}

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

impl Buffer {
    fn owned(&self) -> Option<&DeviceObject<vk::Buffer>> {
        match &self.buffer {
            BufferHandle::Owned(buffer) => Some(buffer),
            #[cfg(feature = "testing")]
            BufferHandle::Testing(_) => None,
        }
    }

    pub fn new(device: Device, desc: BufferDesc) -> crate::Result<Self> {
        let buffer = Self::new_raw(device.clone(), &desc)?;
        let alloc = device.allocate_for_buffer(
            *buffer,
            AllocationDesc {
                name: desc.name.as_deref().unwrap_or_default(),
                location: desc.mem_location,
                strategy: desc.alloc_scheme,
            },
        )?;

        let buffer = Self {
            buffer: BufferHandle::Owned(buffer),
            desc,
            alloc: Allocation::Owned(DeviceObject::new(device.clone(), alloc)),
        };
        device.bind_buffer_memory(
            buffer.alloc.allocation().expect("owned buffer allocation"),
            buffer.raw(),
        )?;

        Ok(buffer)
    }

    fn new_raw(device: Device, desc: &BufferDesc) -> crate::Result<DeviceObject<vk::Buffer>> {
        let queue_families = device.queues.family_indices(desc.queue_families);

        let sharing_mode = if queue_families.len() > 1 {
            vk::SharingMode::CONCURRENT
        } else {
            vk::SharingMode::EXCLUSIVE
        };
        let create_info = vk::BufferCreateInfo::default()
            .size(desc.size)
            .usage(desc.usage)
            .queue_family_indices(queue_families.as_u32_slice())
            .sharing_mode(sharing_mode);

        let buffer = unsafe { device.dev().create_buffer(&create_info, None)? };
        Ok(DeviceObject::new_debug_named(
            device,
            buffer,
            desc.name.clone(),
        ))
    }

    pub fn map(&mut self) -> Option<&[u8]> {
        if let Some(alloc) = self.alloc.allocation() {
            alloc
                .mapped_slice()
                .map(|slice| &slice[..self.desc.size as usize])
        } else {
            None
        }
    }

    pub fn map_mut(&mut self) -> Option<&mut [u8]> {
        if let Some(alloc) = self.alloc.allocation_mut() {
            alloc
                .mapped_slice_mut()
                .map(|slice| &mut slice[..self.desc.size as usize])
        } else {
            None
        }
    }

    pub fn raw(&self) -> vk::Buffer {
        match &self.buffer {
            BufferHandle::Owned(buffer) => **buffer,
            #[cfg(feature = "testing")]
            BufferHandle::Testing(buffer) => *buffer,
        }
    }
    pub fn len(&self) -> u64 {
        self.desc.size
    }

    /// Creates a non-owning buffer fixture for device-free integration tests.
    ///
    /// This constructor is only available with the crate's `testing` feature. The raw handle is
    /// never destroyed and must not be used for Vulkan calls.
    #[cfg(feature = "testing")]
    #[doc(hidden)]
    pub fn for_testing(raw: vk::Buffer, size: u64) -> Self {
        Self {
            buffer: BufferHandle::Testing(raw),
            desc: BufferDesc {
                size,
                ..Default::default()
            },
            alloc: Allocation::Unmanaged,
        }
    }
}

#[cfg(test)]
mod allocation_validation_tests {
    use ash::vk;
    use renderer_device::testing::TestVulkan;

    use super::{Buffer, BufferDesc};
    use crate::{
        device::MemoryLocation,
        images::{Image, ImageDesc},
    };
    use renderer_device::device::AllocationStrategy;

    #[test]
    fn vma_mapping_and_dedicated_resources_under_validation() {
        let test = TestVulkan::expect();
        let device = test.device().clone();

        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<Buffer>();
        assert_send_sync::<Image>();

        let mut mapped_buffer = Buffer::new(
            device.clone(),
            BufferDesc {
                size: 64,
                usage: vk::BufferUsageFlags::TRANSFER_SRC,
                mem_location: MemoryLocation::CpuToGpu,
                ..Default::default()
            },
        )
        .expect("create CpuToGpu buffer");
        let allocation = mapped_buffer
            .alloc
            .allocation()
            .expect("CpuToGpu buffer allocation");
        let memory_type = test
            .physical_device()
            .properties
            .memory_properties
            .memory_types[allocation.memory_type_index() as usize];
        assert!(memory_type.property_flags.contains(
            vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT
        ));

        let expected = (0..64).map(|byte| byte as u8).collect::<Vec<_>>();
        mapped_buffer
            .map_mut()
            .expect("CpuToGpu buffer is mapped")
            .copy_from_slice(&expected);
        assert_eq!(
            mapped_buffer.map().expect("CpuToGpu buffer remains mapped"),
            expected
        );

        let dedicated_buffer = Buffer::new(
            device.clone(),
            BufferDesc {
                size: 256,
                usage: vk::BufferUsageFlags::TRANSFER_DST,
                alloc_scheme: AllocationStrategy::Dedicated,
                ..Default::default()
            },
        )
        .expect("create dedicated buffer");
        assert!(
            dedicated_buffer
                .alloc
                .allocation()
                .expect("dedicated buffer allocation")
                .is_dedicated(),
            "VMA should allocate the dedicated buffer from its own memory block"
        );

        let dedicated_image = Image::new(
            device,
            ImageDesc {
                format: vk::Format::R8G8B8A8_UNORM,
                extent: vk::Extent3D {
                    width: 16,
                    height: 16,
                    depth: 1,
                },
                usage: vk::ImageUsageFlags::SAMPLED,
                alloc_scheme: AllocationStrategy::Dedicated,
                ..Default::default()
            },
        )
        .expect("create dedicated image");
        assert!(
            dedicated_image
                .allocation()
                .expect("dedicated image allocation")
                .allocation()
                .expect("dedicated image VMA allocation")
                .is_dedicated(),
            "VMA should allocate the dedicated image from its own memory block"
        );

        drop(dedicated_image);
        drop(dedicated_buffer);
        drop(mapped_buffer);
        test.device().wait_idle().expect("wait for device idle");
        test.assert_no_validation_errors();
    }
}

/// Byte range selected by a buffer view.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BufferRange {
    pub offset: u64,
    pub size: u64,
}

/// An owned Vulkan buffer view, retaining the buffer it refers to.
#[derive(Debug)]
pub struct BufferView {
    raw: vk::BufferView,
    buffer: Arc<Buffer>,
    range: BufferRange,
}

impl BufferView {
    pub fn new(buffer: Arc<Buffer>, format: vk::Format, range: BufferRange) -> crate::Result<Self> {
        let info = vk::BufferViewCreateInfo::default()
            .buffer(buffer.raw())
            .format(format)
            .offset(range.offset)
            .range(range.size);
        let Some(buffer_handle) = buffer.owned() else {
            return Err(crate::Error::Unspecified);
        };
        let raw = unsafe { buffer_handle.device().raw.create_buffer_view(&info, None)? };
        Ok(Self { raw, buffer, range })
    }

    pub fn raw(&self) -> vk::BufferView {
        self.raw
    }
    pub fn buffer(&self) -> &Arc<Buffer> {
        &self.buffer
    }
    pub fn range(&self) -> BufferRange {
        self.range
    }
}

impl Drop for BufferView {
    fn drop(&mut self) {
        if let Some(buffer) = self.buffer.owned() {
            unsafe { buffer.device().raw.destroy_buffer_view(self.raw, None) };
        }
    }
}