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

    /// Creates a Vulkan buffer without allocating or binding memory.
    pub fn new_unallocated(device: Device, desc: BufferDesc) -> crate::Result<Self> {
        let buffer = Self::new_raw(device, &desc)?;

        Ok(Self {
            buffer: BufferHandle::Owned(buffer),
            desc,
            alloc: Allocation::Unmanaged {
                requirements: std::sync::OnceLock::new(),
            },
        })
    }

    /// Returns the Vulkan memory requirements, including dedicated-allocation hints.
    pub fn memory_requirements(&self) -> crate::memory::Result<crate::memory::MemoryRequirements> {
        let Some(buffer) = self.owned() else {
            return Err(crate::memory::Error::RequirementQuery(
                crate::Error::Unspecified,
            ));
        };

        let requirements_cache = match &self.alloc {
            Allocation::Unmanaged { requirements } => Some(requirements),
            Allocation::Owned(_) | Allocation::Shared(_) => None,
        };

        if let Some(requirements) = requirements_cache.and_then(std::sync::OnceLock::get) {
            return Ok(*requirements);
        }

        let requirements = buffer
            .device()
            .buffer_memory_requirements(**buffer)
            .map_err(crate::memory::Error::RequirementQuery)?;

        if let Some(cache) = requirements_cache {
            let _ = cache.set(requirements);
        }

        Ok(requirements)
    }

    /// Binds this buffer into an aliasable allocation at an allocation-local offset.
    ///
    /// Requirements are queried once and cached in the resource. The Vulkan bind remains the
    /// authoritative validation.
    pub fn bind_allocation(
        &mut self,
        allocation: Arc<DeviceObject<crate::device::GpuAllocation>>,
        offset: u64,
    ) -> crate::memory::Result<()> {
        if !matches!(self.alloc, Allocation::Unmanaged { .. }) {
            return Err(crate::memory::Error::AlreadyBound);
        }

        let requirements = self.memory_requirements()?;
        crate::memory::validate_binding(&requirements, allocation.as_ref(), offset)?;
        let device = self.owned().expect("owned buffer").device();

        self.alloc = device
            .bind_buffer_at_offset(self.raw(), &requirements, allocation, offset)
            .map_err(crate::memory::Error::Bind)?;

        Ok(())
    }

    pub(crate) fn validate_unallocated(&self) -> crate::memory::Result<()> {
        if matches!(self.alloc, Allocation::Unmanaged { .. }) {
            Ok(())
        } else {
            Err(crate::memory::Error::AlreadyBound)
        }
    }

    pub(crate) fn validate_device(&self, device: &Device) -> crate::memory::Result<()> {
        if self.owned().expect("unallocated buffer is owned").device() == device {
            Ok(())
        } else {
            Err(crate::memory::Error::Bind(
                crate::Error::AllocationDeviceMismatch,
            ))
        }
    }

    pub(crate) fn set_shared_allocation(
        &mut self,
        allocation: Arc<DeviceObject<crate::device::GpuAllocation>>,
    ) {
        debug_assert!(matches!(self.alloc, Allocation::Unmanaged { .. }));
        self.alloc = Allocation::Shared(allocation);
    }

    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 {
                requirements: std::sync::OnceLock::new(),
            },
        }
    }
}

#[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();
    }
}

#[cfg(all(test, feature = "testing"))]
mod unallocated_memory_tests {
    use ash::vk;
    use renderer_device::testing::TestVulkan;

    use super::{Buffer, BufferDesc};
    use crate::{
        device::MemoryLocation,
        images::{Image, ImageDesc},
        memory::{Error, MemoryGroup, MemoryMember},
    };

    fn buffer_desc() -> BufferDesc {
        BufferDesc {
            size: 256,
            usage: vk::BufferUsageFlags::TRANSFER_SRC | vk::BufferUsageFlags::TRANSFER_DST,
            mem_location: MemoryLocation::GpuOnly,
            ..Default::default()
        }
    }

    #[test]
    fn unallocated_resources_query_requirements_and_bind() {
        let test = TestVulkan::expect();
        let device = test.device().clone();
        let mut buffer = Buffer::new_unallocated(device.clone(), buffer_desc())
            .expect("create unallocated buffer");
        let buffer_requirements = buffer.memory_requirements().expect("buffer requirements");
        assert!(buffer_requirements.size >= buffer.len());
        assert!(matches!(
            &buffer.alloc,
            crate::device::Allocation::Unmanaged { .. }
        ));

        let mut image = Image::new_unallocated(
            device.clone(),
            ImageDesc {
                format: vk::Format::R8G8B8A8_UNORM,
                extent: vk::Extent3D {
                    width: 16,
                    height: 16,
                    depth: 1,
                },
                usage: vk::ImageUsageFlags::TRANSFER_DST,
                ..Default::default()
            },
        )
        .expect("create unallocated image");
        let image_requirements = image.memory_requirements().expect("image requirements");
        assert!(image_requirements.size > 0);
        assert!(
            image
                .allocation()
                .is_some_and(|allocation| allocation.allocation().is_none())
        );

        let buffer_memory = device
            .allocate_memory(&buffer_requirements, vk::MemoryPropertyFlags::empty())
            .expect("allocate buffer memory");
        buffer
            .bind_allocation(buffer_memory, 0)
            .expect("bind buffer at offset zero");
        assert!(matches!(
            &buffer.alloc,
            crate::device::Allocation::Shared(_)
        ));

        let image_memory = device
            .allocate_memory(&image_requirements, vk::MemoryPropertyFlags::empty())
            .expect("allocate image memory");
        image
            .bind_allocation(image_memory, 0)
            .expect("bind image at offset zero");
        assert!(matches!(
            image.allocation(),
            Some(crate::device::Allocation::Shared(_))
        ));
        drop(image);
        drop(buffer);
        test.device().wait_idle().expect("wait for device idle");
        test.assert_no_validation_errors();
    }

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

        let mut buffer = Buffer::new_unallocated(device.clone(), buffer_desc())
            .expect("create buffer for offset checks");
        let requirements = buffer.memory_requirements().expect("buffer requirements");
        let allocation = device
            .allocate_memory(&requirements, vk::MemoryPropertyFlags::empty())
            .expect("allocate offset-check memory");
        if requirements.alignment > 1 {
            assert!(matches!(
                buffer.bind_allocation(allocation.clone(), 1),
                Err(Error::MisalignedOffset { .. })
            ));
        }
        assert!(matches!(
            buffer.bind_allocation(allocation.clone(), requirements.size),
            Err(Error::OutOfBounds { .. })
        ));
        buffer
            .bind_allocation(allocation.clone(), 0)
            .expect("first buffer bind");
        assert!(matches!(
            buffer.bind_allocation(allocation, 0),
            Err(Error::AlreadyBound)
        ));

        let allocated = device
            .allocate_memory(&requirements, vk::MemoryPropertyFlags::empty())
            .expect("allocate for synthetic type-bits validation");
        let impossible = crate::memory::MemoryRequirements {
            memory_type_bits: 0,
            ..requirements
        };
        assert!(matches!(
            crate::memory::validate_binding(&impossible, allocated.as_ref(), 0),
            Err(Error::IncompatibleMemoryType { .. })
        ));

        drop(buffer);
        test.device().wait_idle().expect("wait for device idle");
        test.assert_no_validation_errors();
    }

    #[test]
    fn memory_group_packs_and_binds_two_buffers() {
        let test = TestVulkan::expect();
        let device = test.device().clone();
        let mut first =
            Buffer::new_unallocated(device.clone(), buffer_desc()).expect("create first buffer");
        let mut second =
            Buffer::new_unallocated(device.clone(), buffer_desc()).expect("create second buffer");
        let first_requirements = first.memory_requirements().expect("first requirements");
        let second_requirements = second.memory_requirements().expect("second requirements");
        let common_bits =
            first_requirements.memory_type_bits & second_requirements.memory_type_bits;
        assert_ne!(common_bits, 0, "the test adapter needs a common type");
        let memory_type_index = common_bits.trailing_zeros();
        let alignment = first_requirements
            .alignment
            .max(second_requirements.alignment);
        let first_end = first_requirements.size;
        let expected_second_offset = first_end
            .checked_add(alignment - 1)
            .expect("offset arithmetic")
            / alignment
            * alignment;
        let size = expected_second_offset + second_requirements.size;
        let mut group = MemoryGroup::new(
            memory_type_index,
            vk::MemoryPropertyFlags::empty(),
            size,
            alignment,
        )
        .expect("create memory group");
        let first_offset = group
            .add(MemoryMember::Buffer(&mut first))
            .expect("pack first buffer");
        let second_offset = group
            .add(MemoryMember::Buffer(&mut second))
            .expect("pack second buffer");
        assert_eq!(first_offset, 0);
        assert_eq!(second_offset, expected_second_offset);
        assert_eq!(second_offset % alignment, 0);
        group.allocate(&device).expect("allocate shared memory");
        group.bind_all().expect("bind group members");
        assert!(matches!(&first.alloc, crate::device::Allocation::Shared(_)));
        assert!(matches!(
            &second.alloc,
            crate::device::Allocation::Shared(_)
        ));
        assert_eq!(
            first
                .alloc
                .allocation()
                .map(crate::device::GpuAllocation::memory),
            second
                .alloc
                .allocation()
                .map(crate::device::GpuAllocation::memory),
            "both buffers retain the shared Vulkan memory block"
        );
        drop(first);
        drop(second);
        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) };
        }
    }
}