vulkan_lib/vma-rs/src/allocator.rs

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

use crate::allocation::AllocationBuilder;

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

pub struct AllocatorBuilder {
    flags: VmaAllocatorCreateFlags,
    preferred_large_heap_block_size: VkDeviceSize,
    allocation_callbacks: Option<VkAllocationCallbacks>,
    device_memory_callbacks: Option<VmaDeviceMemoryCallbacks>,
    frame_in_use_count: u32,
    heap_size_limits: Vec<VkDeviceSize>,
    vulkan_functions: Option<VmaVulkanFunctions>,
    record_settings: Option<VmaRecordSettings>,
}

impl AllocatorBuilder {
    pub fn set_flags(mut self, flags: VmaAllocatorCreateFlags) -> Self {
        self.flags = flags.into();

        self
    }

    pub fn set_preferred_large_heap_block_size(mut self, size: VkDeviceSize) -> Self {
        self.preferred_large_heap_block_size = size;

        self
    }

    pub fn set_allocation_callbacks(mut self, allocation_callbacks: VkAllocationCallbacks) -> Self {
        self.allocation_callbacks = Some(allocation_callbacks);

        self
    }

    pub fn set_device_memory_callbacks(
        mut self,
        device_memory_callbacks: VmaDeviceMemoryCallbacks,
    ) -> Self {
        self.device_memory_callbacks = Some(device_memory_callbacks);

        self
    }

    pub fn set_frame_in_use_count(mut self, use_count: u32) -> Self {
        self.frame_in_use_count = use_count;

        self
    }

    pub fn set_heap_size_limits(mut self, heap_size_limits: Vec<VkDeviceSize>) -> Self {
        self.heap_size_limits = heap_size_limits;

        self
    }

    pub fn set_vulkan_functions(mut self, vulkan_functions: VmaVulkanFunctions) -> Self {
        self.vulkan_functions = Some(vulkan_functions);

        self
    }

    pub fn set_record_settings(mut self, record_settings: VmaRecordSettings) -> Self {
        self.record_settings = Some(record_settings);

        self
    }

    pub fn build(
        self,
        instance: VkInstance,
        device: VkDevice,
        physical_device: VkPhysicalDevice,
        vulkan_api_version: u32,
    ) -> Result<Allocator> {
        let allocator_create_info = VmaAllocatorCreateInfo {
            flags: self.flags,
            physicalDevice: physical_device,
            instance,
            device,
            vulkanApiVersion: vulkan_api_version,
            preferredLargeHeapBlockSize: self.preferred_large_heap_block_size,
            pAllocationCallbacks: match &self.allocation_callbacks {
                Some(callbacks) => callbacks as *const _,
                None => ptr::null(),
            },
            pDeviceMemoryCallbacks: match &self.device_memory_callbacks {
                Some(callbacks) => callbacks as *const _,
                None => ptr::null(),
            },
            frameInUseCount: self.frame_in_use_count,
            pHeapSizeLimit: if self.heap_size_limits.is_empty() {
                ptr::null()
            } else {
                self.heap_size_limits.as_ptr()
            },
            pVulkanFunctions: match &self.vulkan_functions {
                Some(functions) => functions as *const _,
                None => ptr::null(),
            },
            pRecordSettings: match &self.record_settings {
                Some(settings) => settings as *const _,
                None => ptr::null(),
            },
        };

        let mut allocator = MaybeUninit::uninit();

        let result = unsafe { vmaCreateAllocator(&allocator_create_info, allocator.as_mut_ptr()) };

        if result == VK_SUCCESS {
            Ok(Allocator {
                allocator: unsafe { allocator.assume_init() },
            })
        } else {
            Err(anyhow::Error::msg(format!(
                "Failed creating memory allocator: {:?}",
                result
            )))
        }
    }
}

impl Default for AllocatorBuilder {
    fn default() -> Self {
        AllocatorBuilder {
            flags: 0,
            preferred_large_heap_block_size: 0,
            allocation_callbacks: None,
            device_memory_callbacks: None,
            frame_in_use_count: 0,
            heap_size_limits: Vec::new(),
            vulkan_functions: None,
            record_settings: None,
        }
    }
}

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

impl Allocator {
    pub fn allocate(&self) -> AllocationBuilder {
        AllocationBuilder::new(self.allocator)
    }

    pub fn statistics(&self) -> VmaStats {
        let mut stats = MaybeUninit::uninit();

        unsafe {
            vmaCalculateStats(self.allocator, stats.as_mut_ptr());

            stats.assume_init()
        }
    }
}

unsafe impl Send for Allocator {}
unsafe impl Sync for Allocator {}

impl Allocator {
    pub fn builder() -> AllocatorBuilder {
        AllocatorBuilder::default()
    }
}