vulkan_lib/vma-rs/src/allocation.rs

use crate::vma_bindings::*;
use anyhow::Result;
use vulkan_sys::prelude::*;

use crate::allocator_pool::AllocatorPool;

use std::mem::MaybeUninit;
use std::ptr;
use std::slice;

#[derive(Debug, Clone)]
pub struct AllocationBuilder {
    allocator: VmaAllocator,

    flags: VmaAllocationCreateFlags,
    usage: VmaMemoryUsage,
    required_flags: VkMemoryPropertyFlagBits,
    preferred_flags: VkMemoryPropertyFlagBits,
    memory_type_bits: u32,
    priority: f32,

    pool: Option<AllocatorPool>,
}

pub trait Build<T> {
    fn build(self, argument: T) -> Result<Allocation>;
}

impl AllocationBuilder {
    pub fn set_flags(mut self, flags: impl Into<VmaAllocationCreateFlags>) -> Self {
        self.flags = flags.into();

        self
    }

    pub fn set_usage(mut self, usage: VmaMemoryUsage) -> Self {
        self.usage = usage;

        self
    }

    pub fn priority(mut self, priority: f32) -> Self {
        self.priority = priority;

        self
    }

    pub fn set_required_flags(
        mut self,
        required_flags: impl Into<VkMemoryPropertyFlagBits>,
    ) -> Self {
        self.required_flags = required_flags.into();

        self
    }

    pub fn set_preferred_flags(
        mut self,
        preferred_flags: impl Into<VkMemoryPropertyFlagBits>,
    ) -> Self {
        self.preferred_flags = preferred_flags.into();

        self
    }

    pub fn set_memory_type_bits(mut self, memory_type_bits: u32) -> Self {
        self.memory_type_bits = memory_type_bits;

        self
    }

    pub fn set_pool(mut self, pool: AllocatorPool) -> Self {
        self.pool = Some(pool);

        self
    }

    fn vma_allocation_create_info(&self) -> VmaAllocationCreateInfo {
        let mut create_info = VmaAllocationCreateInfo::new(
            self.flags,
            self.usage,
            self.required_flags,
            self.preferred_flags,
            self.memory_type_bits,
            self.priority,
        );

        if let Some(pool) = &self.pool {
            create_info.set_pool(pool.pool());
        }

        create_info
    }
}

impl Build<&VkMemoryRequirements> for AllocationBuilder {
    fn build(self, memory_requirements: &VkMemoryRequirements) -> Result<Allocation> {
        let create_info = self.vma_allocation_create_info();

        let mut allocation = MaybeUninit::uninit();
        let mut allocation_info = MaybeUninit::uninit();

        let result = unsafe {
            vmaAllocateMemory(
                self.allocator,
                memory_requirements,
                &create_info,
                allocation.as_mut_ptr(),
                allocation_info.as_mut_ptr(),
            )
        };

        if result == VK_SUCCESS {
            Ok(Allocation::new(
                self.allocator,
                unsafe { allocation.assume_init() },
                AllocationType::MemoryOnly,
                unsafe { allocation_info.assume_init() },
            ))
        } else {
            Err(anyhow::Error::msg(format!(
                "Failed allocating memory for Buffer: {:?}",
                result
            )))
        }
    }
}

impl Build<VkBuffer> for AllocationBuilder {
    fn build(self, buffer: VkBuffer) -> Result<Allocation> {
        let create_info = self.vma_allocation_create_info();

        let mut allocation = MaybeUninit::uninit();
        let mut allocation_info = MaybeUninit::uninit();

        let result = unsafe {
            vmaAllocateMemoryForBuffer(
                self.allocator,
                buffer,
                &create_info,
                allocation.as_mut_ptr(),
                allocation_info.as_mut_ptr(),
            )
        };

        if result == VK_SUCCESS {
            let mut allocation = Allocation::new(
                self.allocator,
                unsafe { allocation.assume_init() },
                AllocationType::MemoryOnly,
                unsafe { allocation_info.assume_init() },
            );

            allocation.bind_buffer_memory(buffer)?;

            Ok(allocation)
        } else {
            Err(anyhow::Error::msg(format!(
                "Failed allocating memory for Buffer: {:?}",
                result
            )))
        }
    }
}

impl Build<VkImage> for AllocationBuilder {
    fn build(self, image: VkImage) -> Result<Allocation> {
        let create_info = self.vma_allocation_create_info();

        let mut allocation = MaybeUninit::uninit();
        let mut allocation_info = MaybeUninit::uninit();

        let result = unsafe {
            vmaAllocateMemoryForImage(
                self.allocator,
                image,
                &create_info,
                allocation.as_mut_ptr(),
                allocation_info.as_mut_ptr(),
            )
        };

        if result == VK_SUCCESS {
            let mut allocation = Allocation::new(
                self.allocator,
                unsafe { allocation.assume_init() },
                AllocationType::MemoryOnly,
                unsafe { allocation_info.assume_init() },
            );

            allocation.bind_image_memory(image)?;

            Ok(allocation)
        } else {
            Err(anyhow::Error::msg(format!(
                "Failed allocating memory for Image: {:?}",
                result
            )))
        }
    }
}

impl Build<&VkImageCreateInfo> for AllocationBuilder {
    fn build(self, image_create_info: &VkImageCreateInfo) -> Result<Allocation> {
        let create_info = self.vma_allocation_create_info();

        let mut image = MaybeUninit::uninit();
        let mut allocation = MaybeUninit::uninit();
        let mut allocation_info = MaybeUninit::uninit();

        let result = unsafe {
            vmaCreateImage(
                self.allocator,
                image_create_info,
                &create_info,
                image.as_mut_ptr(),
                allocation.as_mut_ptr(),
                allocation_info.as_mut_ptr(),
            )
        };

        if result == VK_SUCCESS {
            Ok(Allocation {
                allocator: self.allocator,

                allocation: unsafe { allocation.assume_init() },
                allocation_type: AllocationType::Image(unsafe { image.assume_init() }),
                allocation_info: unsafe { allocation_info.assume_init() },
            })
        } else {
            Err(anyhow::Error::msg(format!(
                "Failed creating Image and allocating memory for Image: {:?}",
                result
            )))
        }
    }
}

impl Build<&VkBufferCreateInfo> for AllocationBuilder {
    fn build(self, buffer_create_info: &VkBufferCreateInfo) -> Result<Allocation> {
        let create_info = self.vma_allocation_create_info();

        let mut buffer = MaybeUninit::uninit();
        let mut allocation = MaybeUninit::uninit();
        let mut allocation_info = MaybeUninit::uninit();

        let result = unsafe {
            vmaCreateBuffer(
                self.allocator,
                buffer_create_info,
                &create_info,
                buffer.as_mut_ptr(),
                allocation.as_mut_ptr(),
                allocation_info.as_mut_ptr(),
            )
        };

        if result == VK_SUCCESS {
            Ok(Allocation {
                allocator: self.allocator,

                allocation: unsafe { allocation.assume_init() },
                allocation_type: AllocationType::Buffer(unsafe { buffer.assume_init() }),
                allocation_info: unsafe { allocation_info.assume_init() },
            })
        } else {
            Err(anyhow::Error::msg(format!(
                "Failed creating Buffer and allocating memory for Buffer: {:?}",
                result
            )))
        }
    }
}

impl AllocationBuilder {
    pub(crate) fn new(allocator: VmaAllocator) -> Self {
        AllocationBuilder {
            allocator,

            flags: VmaAllocationCreateFlags::default(),
            usage: VmaMemoryUsage::VMA_MEMORY_USAGE_UNKNOWN,
            required_flags: VkMemoryPropertyFlagBits::default(),
            preferred_flags: VkMemoryPropertyFlagBits::default(),
            memory_type_bits: 0,
            priority: 0.0,

            pool: None,
        }
    }
}

#[derive(Debug, Clone)]
enum AllocationType {
    Buffer(VkBuffer),
    Image(VkImage),
    MemoryOnly,
}

#[derive(Debug, Clone)]
pub struct Allocation {
    allocator: VmaAllocator,

    allocation: VmaAllocation,
    allocation_type: AllocationType,
    allocation_info: VmaAllocationInfo,
}

unsafe impl Send for Allocation {}
unsafe impl Sync for Allocation {}

impl Allocation {
    fn new(
        allocator: VmaAllocator,
        allocation: VmaAllocation,
        allocation_type: AllocationType,
        allocation_info: VmaAllocationInfo,
    ) -> Self {
        Self {
            allocator,
            allocation,
            allocation_type,
            allocation_info,
        }
    }

    pub fn memory_type_index(&self) -> u32 {
        self.allocation_info.memoryType
    }

    pub fn device_memory(&self) -> VkDeviceMemory {
        self.allocation_info.deviceMemory
    }

    pub fn offset(&self) -> VkDeviceSize {
        self.allocation_info.offset
    }

    pub fn size(&self) -> VkDeviceSize {
        self.allocation_info.size
    }

    pub fn map<T: Clone>(&self, length: VkDeviceSize) -> Result<VkMappedMemory<'_, T>> {
        let mut data = MaybeUninit::uninit();

        let result = unsafe { vmaMapMemory(self.allocator, self.allocation, data.as_mut_ptr()) };

        let slice =
            unsafe { slice::from_raw_parts_mut(data.assume_init() as *mut T, length as usize) };

        if result == VK_SUCCESS {
            let mut mapped_memory = VkMappedMemory::new(slice);

            let allocation = self.allocation;
            let allocator = self.allocator;

            mapped_memory.set_unmap(move || unsafe { vmaUnmapMemory(allocator, allocation) });

            Ok(mapped_memory)
        } else {
            Err(anyhow::Error::msg(format!(
                "Mapping VkMemory failed: {:?}",
                result
            )))
        }
    }

    pub fn buffer(&self) -> VkBuffer {
        match self.allocation_type {
            AllocationType::Buffer(buffer) => buffer,
            AllocationType::Image(_) => panic!("Allocation is a VkBuffer"),
            AllocationType::MemoryOnly => panic!("Allocation is memory only"),
        }
    }

    pub fn image(&self) -> VkImage {
        match self.allocation_type {
            AllocationType::Buffer(_) => panic!("Allocation is a VkImage"),
            AllocationType::Image(image) => image,
            AllocationType::MemoryOnly => panic!("Allocation is memory only"),
        }
    }

    pub fn bind_buffer_memory(&mut self, buffer: VkBuffer) -> Result<()> {
        #[cfg(debug_assertions)]
        {
            match self.allocation_type {
                AllocationType::Buffer(_) => panic!("allocation already bound to buffer"),
                AllocationType::Image(_) => panic!("allocation already bound to image"),
                AllocationType::MemoryOnly => (),
            }
        }

        let result = unsafe { vmaBindBufferMemory(self.allocator, self.allocation, buffer) };

        if result == VK_SUCCESS {
            Ok(())
        } else {
            Err(anyhow::Error::msg(format!(
                "Failed binding Buffer to memory: {:?}",
                result
            )))
        }
    }

    pub fn bind_image_memory(&mut self, image: VkImage) -> Result<()> {
        #[cfg(debug_assertions)]
        {
            match self.allocation_type {
                AllocationType::Buffer(_) => panic!("allocation already bound to buffer"),
                AllocationType::Image(_) => panic!("allocation already bound to image"),
                AllocationType::MemoryOnly => (),
            }
        }

        let result = unsafe { vmaBindImageMemory(self.allocator, self.allocation, image) };

        if result == VK_SUCCESS {
            Ok(())
        } else {
            Err(anyhow::Error::msg(format!(
                "Failed binding Image to memory: {:?}",
                result
            )))
        }
    }
}

impl Drop for Allocation {
    fn drop(&mut self) {
        match self.allocation_type {
            AllocationType::Buffer(buffer) => unsafe {
                vmaDestroyBuffer(self.allocator, buffer, self.allocation)
            },
            AllocationType::Image(image) => unsafe {
                vmaDestroyImage(self.allocator, image, self.allocation)
            },
            AllocationType::MemoryOnly => unsafe { vmaFreeMemory(self.allocator, self.allocation) },
        }
    }
}

impl VmaAllocationCreateInfo {
    pub fn new(
        flags: impl Into<VmaAllocationCreateFlags>,
        usage: VmaMemoryUsage,
        required_flags: VkMemoryPropertyFlagBits,
        preferred_flags: VkMemoryPropertyFlagBits,
        memory_type_bits: u32,
        priority: f32,
    ) -> Self {
        VmaAllocationCreateInfo {
            flags: flags.into(),
            usage,
            requiredFlags: required_flags,
            preferredFlags: preferred_flags,
            memoryTypeBits: memory_type_bits,
            pool: ptr::null_mut(),
            pUserData: ptr::null_mut(),
            priority,
        }
    }

    pub fn set_pool(&mut self, pool: VmaPool) {
        self.pool = pool;
    }

    // pub fn set_user_data<T>(&mut self, data: &mut T) {
    //     self.pUserData = data as *mut T as *mut _
    // }
}