vulkan_lib/vulkan-rs/src/buffer.rs

use crate::prelude::*;

use anyhow::Result;

use std;
use std::mem;
use std::sync::Arc;

pub struct BufferBuilder<'a, T> {
    flags: VkBufferCreateFlagBits,
    usage: VkBufferUsageFlagBits,
    memory_usage: Option<MemoryUsage>,
    sharing_mode: VkSharingMode,
    data: Option<&'a [T]>,
    size: VkDeviceSize,

    alignment: Option<VkDeviceSize>,
}

impl<'a, T> BufferBuilder<'a, T> {
    pub fn set_memory_usage(mut self, usage: MemoryUsage) -> Self {
        self.memory_usage = Some(usage);

        self
    }

    pub fn set_usage(mut self, usage: impl Into<VkBufferUsageFlagBits>) -> Self {
        self.usage = usage.into();

        self
    }

    pub fn set_data(mut self, data: &'a [T]) -> Self {
        self.data = Some(data);

        self
    }

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

        self
    }

    pub fn set_sharing_mode(mut self, sharing_mode: VkSharingMode) -> Self {
        self.sharing_mode = sharing_mode;

        self
    }

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

        self
    }

    pub fn force_alignment(mut self, alignment: VkDeviceSize) -> Self {
        self.alignment = Some(alignment);

        self
    }
}

impl<'a, T: Clone + Send + Sync + 'static> BufferBuilder<'a, T> {
    pub fn build(self, device: Arc<Device>) -> Result<Arc<Buffer<T>>> {
        let size = match self.data {
            Some(data) => data.len() as VkDeviceSize,
            None => self.size,
        };

        if size == 0 {
            panic!("Vulkan buffer size must not be zero");
        }

        // create buffer
        let buffer_ci = VkBufferCreateInfo::new(
            self.flags,
            size * mem::size_of::<T>() as VkDeviceSize,
            self.usage,
            self.sharing_mode,
            &[],
        );

        let buffer = device.create_buffer(&buffer_ci)?;

        // create memory
        let memory = match self.alignment {
            Some(alignment) => {
                let mut memory_requirements = device.buffer_memory_requirements(buffer);
                memory_requirements.alignment = alignment;

                Memory::forced_requirements(
                    &device,
                    memory_requirements,
                    buffer,
                    MemoryUsage::into_vma(self.memory_usage),
                )?
            }
            None => {
                Memory::buffer_memory(&device, buffer, MemoryUsage::into_vma(self.memory_usage))?
            }
        };

        let buffer = Arc::new(Buffer {
            device,
            buffer,
            memory,

            _usage: self.usage,

            _sharing_mode: self.sharing_mode,

            size,
        });

        if let Some(data) = self.data {
            buffer.fill(data)?;
        }

        Ok(buffer)
    }
}

#[derive(Debug)]
pub struct Buffer<T> {
    device: Arc<Device>,
    buffer: VkBuffer,

    memory: Arc<Memory<T>>,

    _usage: VkBufferUsageFlagBits,

    _sharing_mode: VkSharingMode,
    size: VkDeviceSize,
}

impl<T: Clone + Send + Sync + 'static> Buffer<T> {
    pub fn fill(&self, data: &[T]) -> Result<()> {
        let mut buffer_map = self.map(data.len() as VkDeviceSize)?;

        buffer_map.copy(data);

        Ok(())
    }

    pub fn map(&self, length: VkDeviceSize) -> Result<VkMappedMemory<'_, T>> {
        self.memory.map(length)
    }

    pub fn map_complete(&self) -> Result<VkMappedMemory<'_, T>> {
        self.memory.map(self.size)
    }

    pub fn into_device_local(
        self: &Arc<Buffer<T>>,
        buffer_recorder: &mut CommandBufferRecorder<'_>,
        access_mask: impl Into<VkAccessFlagBits>,
        stage: impl Into<VkPipelineStageFlagBits>,
        usage: impl Into<VkBufferUsageFlagBits>,
    ) -> Result<Arc<Buffer<T>>> {
        let new_usage = usage.into() | VK_BUFFER_USAGE_TRANSFER_DST_BIT;

        let device_local_buffer = Buffer::builder()
            .set_memory_usage(MemoryUsage::GpuOnly)
            .set_usage(new_usage)
            .set_size(self.size)
            .build(self.device.clone())?;

        // copy complete buffer
        buffer_recorder.copy_buffer(
            self,
            &device_local_buffer,
            &[VkBufferCopy {
                srcOffset: 0,
                dstOffset: 0,
                size: self.byte_size(),
            }],
        );

        // make sure buffer is copied before using it
        buffer_recorder.buffer_barrier(
            &device_local_buffer,
            VK_ACCESS_TRANSFER_WRITE_BIT,
            VK_PIPELINE_STAGE_TRANSFER_BIT,
            access_mask,
            stage,
        );

        Ok(device_local_buffer)
    }
}

impl<T> Buffer<T> {
    pub fn builder<'a>() -> BufferBuilder<'a, T> {
        BufferBuilder {
            flags: 0u32.into(),
            usage: 0u32.into(),
            memory_usage: None,
            sharing_mode: VK_SHARING_MODE_EXCLUSIVE,
            data: None,
            size: 0,

            alignment: None,
        }
    }

    pub fn byte_size(&self) -> VkDeviceSize {
        self.size * mem::size_of::<T>() as VkDeviceSize
    }

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

    pub fn device_address(&self) -> Address {
        self.device.get_buffer_device_address(self.buffer)
    }
}

impl<T> VulkanDevice for Buffer<T> {
    fn device(&self) -> &Arc<Device> {
        &self.device
    }
}

impl_vk_handle_t!(Buffer, VkBuffer, buffer);

impl<T> VkHandle<VkDeviceMemory> for Buffer<T> {
    fn vk_handle(&self) -> VkDeviceMemory {
        self.memory.vk_handle()
    }
}

impl<'a, T> VkHandle<VkDeviceMemory> for &'a Buffer<T> {
    fn vk_handle(&self) -> VkDeviceMemory {
        self.memory.vk_handle()
    }
}

impl<T> VkHandle<VkDeviceMemory> for Arc<Buffer<T>> {
    fn vk_handle(&self) -> VkDeviceMemory {
        self.memory.vk_handle()
    }
}

impl<'a, T> VkHandle<VkDeviceMemory> for &'a Arc<Buffer<T>> {
    fn vk_handle(&self) -> VkDeviceMemory {
        self.memory.vk_handle()
    }
}

impl<T> Drop for Buffer<T> {
    fn drop(&mut self) {
        self.device.destroy_buffer(self.buffer);
    }
}

// use crate::{ffi::*, handle_ffi_result};

impl<T> FFIBufferTrait for Buffer<T> {
    fn byte_size(&self) -> VkDeviceSize {
        self.byte_size()
    }
}

pub trait FFIBufferTrait {
    fn byte_size(&self) -> VkDeviceSize;
}

pub struct FFIBuffer {
    trait_obj: Box<dyn FFIBufferTrait>,
}

impl FFIBuffer {
    fn byte_size(&self) -> VkDeviceSize {
        self.trait_obj.byte_size()
    }
}

#[no_mangle]
pub extern "C" fn create_buffer(_device: *const Device) -> *const FFIBuffer {
    todo!()
}

#[no_mangle]
pub extern "C" fn byte_size(buffer: *const FFIBuffer) -> VkDeviceSize {
    unsafe { &*buffer }.byte_size()
}