jay-compositor 1.15.0

use crate::gfx_apis::vulkan::VulkanError;
use crate::gfx_apis::vulkan::allocator::VulkanAllocator;
use crate::gfx_apis::vulkan::buffer_cache::VulkanBufferUncached;
use crate::gfx_apis::vulkan::device::DescriptorHeapDevice;
use crate::gfx_apis::vulkan::device::VulkanDevice;
use crate::utils::page_alloc::PAGE_ALLOC_PAGE_SIZE;
use crate::utils::page_alloc::PageAlloc;
use crate::utils::page_alloc::PageAllocEntry;
use ash::vk::BindHeapInfoEXT;
use ash::vk::BufferUsageFlags;
use ash::vk::CommandBuffer;
use ash::vk::DeviceAddressRangeEXT;
use ash::vk::DeviceSize;
use linearize::Linearize;
use std::rc::Rc;

pub struct DescriptorHeap {
    ty: DescriptorHeapType,
    page_alloc: Rc<PageAlloc>,
    buffer: VulkanBufferUncached,
    size: usize,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, Linearize)]
pub enum DescriptorHeapType {
    Sampler,
    Resource,
}

impl VulkanDevice {
    pub fn allocate_descriptor_heap(
        self: &Rc<Self>,
        dh: &DescriptorHeapDevice,
        allocator: &Rc<VulkanAllocator>,
        size: usize,
        ty: DescriptorHeapType,
    ) -> Result<Rc<DescriptorHeap>, VulkanError> {
        let size = size.next_multiple_of(PAGE_ALLOC_PAGE_SIZE as usize);
        let full_size = size as DeviceSize + dh.min_reserved_range[ty];
        if full_size > dh.max_size[ty] {
            return Err(VulkanError::MaximumHeapSize);
        }
        let buffer = self.allocate_uncached_buffer(
            full_size,
            (PAGE_ALLOC_PAGE_SIZE as DeviceSize).max(dh.alignment[ty]),
            allocator,
            BufferUsageFlags::DESCRIPTOR_HEAP_EXT | BufferUsageFlags::SHADER_DEVICE_ADDRESS,
        )?;
        let pages = size / PAGE_ALLOC_PAGE_SIZE as usize;
        Ok(Rc::new(DescriptorHeap {
            ty,
            size,
            page_alloc: dh.alloc.create_alloc(pages),
            buffer,
        }))
    }
}

impl DescriptorHeap {
    pub fn size(&self) -> usize {
        self.size
    }

    pub fn not_contains(&self, entry: &PageAllocEntry) -> bool {
        entry.is_not_in_alloc(&self.page_alloc)
    }

    pub fn bind(&self, dh: &DescriptorHeapDevice, buf: CommandBuffer) {
        let info = BindHeapInfoEXT::default()
            .heap_range(DeviceAddressRangeEXT {
                address: self.buffer.address,
                size: self.buffer.size,
            })
            .reserved_range_offset(self.size as _)
            .reserved_range_size(dh.min_reserved_range[self.ty]);
        unsafe {
            match self.ty {
                DescriptorHeapType::Sampler => dh.device.cmd_bind_sampler_heap(buf, &info),
                DescriptorHeapType::Resource => dh.device.cmd_bind_resource_heap(buf, &info),
            }
        }
    }

    pub fn allocate(
        self: &Rc<Self>,
        data: &[u8],
    ) -> Result<Option<Rc<PageAllocEntry>>, VulkanError> {
        let Some(res) = self.page_alloc.allocate(data.len()) else {
            return Ok(None);
        };
        unsafe {
            self.buffer
                .allocation
                .upload_range(res.offset() as usize, data.len(), |p| {
                    std::ptr::copy_nonoverlapping(data.as_ptr(), p, data.len());
                })?;
        }
        Ok(Some(res))
    }
}