shdrlib 0.1.5

//! Memory module - Buffer and image allocation
//!
//! **WARNING**: CORE objects do NOT enforce lifetime dependencies.
//! Memory objects must not outlive the device they were created from.
//!
//! **NOTE**: This module provides allocation FUNCTIONS, not an allocator struct.
//! CORE tier does not track state - that's the EX tier's job.

use crate::core::Device;
use ash::vk;
use thiserror::Error;

/// Memory operation error
#[derive(Debug, Error)]
pub enum MemoryError {
    #[error("Memory allocation failed: {0}")]
    AllocationFailed(vk::Result),

    #[error("Memory mapping failed: {0}")]
    MapFailed(vk::Result),

    #[error("Memory binding failed: {0}")]
    BindFailed(vk::Result),

    #[error("Out of device memory")]
    OutOfDeviceMemory,

    #[error("Out of host memory")]
    OutOfHostMemory,

    #[error("Buffer creation failed: {0}")]
    BufferCreationFailed(vk::Result),

    #[error("Image creation failed: {0}")]
    ImageCreationFailed(vk::Result),

    #[error("Image view creation failed: {0}")]
    ImageViewCreationFailed(vk::Result),

    #[error("Memory type not found")]
    MemoryTypeNotFound,
}

/// Buffer wrapper with owned memory
///
/// # Safety
///
/// This object does NOT enforce that the device outlives it.
/// Using the buffer after dropping the device causes undefined behavior.
pub struct Buffer {
    buffer: vk::Buffer,
    memory: vk::DeviceMemory,
    size: vk::DeviceSize,
    usage: vk::BufferUsageFlags,
}

impl Buffer {
    /// Create a new buffer with allocated memory
    ///
    /// # Arguments
    ///
    /// * `device` - The device to create the buffer on
    /// * `size` - Size of the buffer in bytes
    /// * `usage` - Buffer usage flags
    /// * `properties` - Memory property flags
    ///
    /// # Errors
    ///
    /// Returns `MemoryError` if buffer creation or memory allocation fails.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use shdrlib::core::{Buffer, Device};
    /// use ash::vk;
    ///
    /// # fn example(device: &Device) -> Result<(), shdrlib::core::MemoryError> {
    /// let buffer = Buffer::new(
    ///     device,
    ///     1024,
    ///     vk::BufferUsageFlags::VERTEX_BUFFER,
    ///     vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
    /// )?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn new(
        device: &Device,
        size: vk::DeviceSize,
        usage: vk::BufferUsageFlags,
        properties: vk::MemoryPropertyFlags,
    ) -> Result<Self, MemoryError> {
        // Create buffer
        let buffer_info = vk::BufferCreateInfo {
            size,
            usage,
            sharing_mode: vk::SharingMode::EXCLUSIVE,
            ..Default::default()
        };

        // SAFETY: device is valid, buffer_info is properly initialized
        let buffer = unsafe {
            device
                .handle()
                .create_buffer(&buffer_info, None)
                .map_err(MemoryError::BufferCreationFailed)?
        };

        // Get memory requirements
        // SAFETY: device and buffer are valid
        let mem_requirements = unsafe { device.handle().get_buffer_memory_requirements(buffer) };

        // Find suitable memory type
        let memory_type_index = device
            .find_memory_type(mem_requirements.memory_type_bits, properties)
            .ok_or(MemoryError::MemoryTypeNotFound)?;

        // Allocate memory
        let alloc_info = vk::MemoryAllocateInfo {
            allocation_size: mem_requirements.size,
            memory_type_index,
            ..Default::default()
        };

        // SAFETY: device is valid, alloc_info is properly initialized
        let memory = unsafe {
            device
                .handle()
                .allocate_memory(&alloc_info, None)
                .map_err(MemoryError::AllocationFailed)?
        };

        // Bind memory to buffer
        // SAFETY: device, buffer, and memory are valid
        unsafe {
            device
                .handle()
                .bind_buffer_memory(buffer, memory, 0)
                .map_err(MemoryError::BindFailed)?;
        }

        Ok(Self {
            buffer,
            memory,
            size,
            usage,
        })
    }

    /// Create a Buffer from existing Vulkan handles
    ///
    /// # Safety
    ///
    /// Caller must ensure the buffer and memory are valid and properly bound.
    /// This is primarily for internal use by EX tier.
    #[inline]
    pub(crate) fn from_raw(
        buffer: vk::Buffer,
        memory: vk::DeviceMemory,
        size: vk::DeviceSize,
    ) -> Self {
        Self {
            buffer,
            memory,
            size,
            usage: vk::BufferUsageFlags::empty(), // Unknown at this point
        }
    }

    /// Get the raw Vulkan buffer handle
    #[inline]
    pub fn handle(&self) -> vk::Buffer {
        self.buffer
    }

    /// Get the device memory handle
    #[inline]
    pub fn memory(&self) -> vk::DeviceMemory {
        self.memory
    }

    /// Get the buffer size
    #[inline]
    pub fn size(&self) -> vk::DeviceSize {
        self.size
    }

    /// Get the buffer usage flags
    #[inline]
    pub fn usage(&self) -> vk::BufferUsageFlags {
        self.usage
    }

    /// Map buffer memory to host-visible pointer
    ///
    /// # Safety
    ///
    /// The memory must be host-visible. Caller must ensure proper synchronization.
    ///
    /// # Errors
    ///
    /// Returns `MemoryError::MapFailed` if mapping fails.
    pub unsafe fn map(&self, device: &Device) -> Result<*mut u8, MemoryError> {
        // SAFETY: Caller guarantees memory is host-visible
        let ptr = unsafe {
            device
                .handle()
                .map_memory(self.memory, 0, self.size, vk::MemoryMapFlags::empty())
                .map_err(MemoryError::MapFailed)? as *mut u8
        };
        Ok(ptr)
    }

    /// Unmap buffer memory
    pub fn unmap(&self, device: &Device) {
        // SAFETY: device and memory are valid
        unsafe {
            device.handle().unmap_memory(self.memory);
        }
    }

    /// Copy data to buffer (convenience method for host-visible buffers)
    ///
    /// # Safety
    ///
    /// The buffer must be host-visible. The data must fit within the buffer size.
    pub unsafe fn copy_from_slice<T: Copy>(
        &self,
        device: &Device,
        data: &[T],
    ) -> Result<(), MemoryError> {
        let size = std::mem::size_of_val(data) as vk::DeviceSize;
        if size > self.size {
            return Err(MemoryError::AllocationFailed(
                vk::Result::ERROR_OUT_OF_HOST_MEMORY,
            ));
        }

        // SAFETY: Caller guarantees buffer is host-visible and data fits
        unsafe {
            let ptr = self.map(device)?;
            std::ptr::copy_nonoverlapping(data.as_ptr() as *const u8, ptr, size as usize);
            self.unmap(device);
        }

        Ok(())
    }

    /// Destroy the buffer manually
    ///
    /// # Safety
    ///
    /// The device must be the same device used to create this buffer.
    /// After calling destroy(), the buffer handle is invalidated.
    pub fn destroy(&self, device: &Device) {
        // SAFETY: device, buffer, and memory are valid
        unsafe {
            device.handle().destroy_buffer(self.buffer, None);
            device.handle().free_memory(self.memory, None);
        }
    }
}

impl Drop for Buffer {
    fn drop(&mut self) {
        if self.buffer != vk::Buffer::null() {
            eprintln!("WARNING: Buffer dropped without calling .destroy() - potential memory leak");
        }
    }
}

/// Image wrapper with owned memory and view
///
/// **NOTE**: Image ALWAYS owns an ImageView - images need views to be used.
///
/// # Safety
///
/// This object does NOT enforce that the device outlives it.
/// Using the image after dropping the device causes undefined behavior.
pub struct Image {
    image: vk::Image,
    memory: vk::DeviceMemory,
    view: vk::ImageView,
    extent: vk::Extent3D,
    format: vk::Format,
    usage: vk::ImageUsageFlags,
}

impl Image {
    /// Create a new image with allocated memory and view
    ///
    /// # Arguments
    ///
    /// * `device` - The device to create the image on
    /// * `extent` - Image dimensions
    /// * `format` - Image format
    /// * `usage` - Image usage flags
    /// * `properties` - Memory property flags
    ///
    /// # Errors
    ///
    /// Returns `MemoryError` if image creation, memory allocation, or view creation fails.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use shdrlib::core::{Image, Device};
    /// use ash::vk;
    ///
    /// # fn example(device: &Device) -> Result<(), shdrlib::core::MemoryError> {
    /// let image = Image::new(
    ///     device,
    ///     vk::Extent3D { width: 512, height: 512, depth: 1 },
    ///     vk::Format::R8G8B8A8_UNORM,
    ///     vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::TRANSFER_DST,
    ///     vk::MemoryPropertyFlags::DEVICE_LOCAL,
    /// )?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn new(
        device: &Device,
        extent: vk::Extent3D,
        format: vk::Format,
        usage: vk::ImageUsageFlags,
        properties: vk::MemoryPropertyFlags,
    ) -> Result<Self, MemoryError> {
        // Create image
        let image_info = vk::ImageCreateInfo {
            image_type: vk::ImageType::TYPE_2D,
            format,
            extent,
            mip_levels: 1,
            array_layers: 1,
            samples: vk::SampleCountFlags::TYPE_1,
            tiling: vk::ImageTiling::OPTIMAL,
            usage,
            sharing_mode: vk::SharingMode::EXCLUSIVE,
            initial_layout: vk::ImageLayout::UNDEFINED,
            ..Default::default()
        };

        // SAFETY: device is valid, image_info is properly initialized
        let image = unsafe {
            device
                .handle()
                .create_image(&image_info, None)
                .map_err(MemoryError::ImageCreationFailed)?
        };

        // Get memory requirements
        // SAFETY: device and image are valid
        let mem_requirements = unsafe { device.handle().get_image_memory_requirements(image) };

        // Find suitable memory type
        let memory_type_index = device
            .find_memory_type(mem_requirements.memory_type_bits, properties)
            .ok_or(MemoryError::MemoryTypeNotFound)?;

        // Allocate memory
        let alloc_info = vk::MemoryAllocateInfo {
            allocation_size: mem_requirements.size,
            memory_type_index,
            ..Default::default()
        };

        // SAFETY: device is valid, alloc_info is properly initialized
        let memory = unsafe {
            device
                .handle()
                .allocate_memory(&alloc_info, None)
                .map_err(MemoryError::AllocationFailed)?
        };

        // Bind memory to image
        // SAFETY: device, image, and memory are valid
        unsafe {
            device
                .handle()
                .bind_image_memory(image, memory, 0)
                .map_err(MemoryError::BindFailed)?;
        }

        // Create image view
        let view_info = vk::ImageViewCreateInfo {
            image,
            view_type: vk::ImageViewType::TYPE_2D,
            format,
            subresource_range: vk::ImageSubresourceRange {
                aspect_mask: vk::ImageAspectFlags::COLOR,
                base_mip_level: 0,
                level_count: 1,
                base_array_layer: 0,
                layer_count: 1,
            },
            ..Default::default()
        };

        // SAFETY: device is valid, view_info is properly initialized
        let view = unsafe {
            device
                .handle()
                .create_image_view(&view_info, None)
                .map_err(MemoryError::ImageViewCreationFailed)?
        };

        Ok(Self {
            image,
            memory,
            view,
            extent,
            format,
            usage,
        })
    }

    /// Create an Image from existing Vulkan handles
    ///
    /// # Safety
    ///
    /// Caller must ensure the image, memory, and view are valid and properly bound.
    /// This is primarily for internal use by EX tier.
    #[inline]
    pub(crate) fn from_raw(
        image: vk::Image,
        memory: vk::DeviceMemory,
        view: vk::ImageView,
        format: vk::Format,
        extent: vk::Extent3D,
    ) -> Self {
        Self {
            image,
            memory,
            view,
            extent,
            format,
            usage: vk::ImageUsageFlags::empty(), // Unknown at this point
        }
    }

    /// Get the raw Vulkan image handle
    #[inline]
    pub fn handle(&self) -> vk::Image {
        self.image
    }

    /// Get the device memory handle
    #[inline]
    pub fn memory(&self) -> vk::DeviceMemory {
        self.memory
    }

    /// Get the image view handle
    #[inline]
    pub fn view(&self) -> vk::ImageView {
        self.view
    }

    /// Get the image extent
    #[inline]
    pub fn extent(&self) -> vk::Extent3D {
        self.extent
    }

    /// Get the image format
    #[inline]
    pub fn format(&self) -> vk::Format {
        self.format
    }

    /// Get the image usage flags
    #[inline]
    pub fn usage(&self) -> vk::ImageUsageFlags {
        self.usage
    }

    /// Destroy the image manually
    ///
    /// # Safety
    ///
    /// The device must be the same device used to create this image.
    /// After calling destroy(), the image handle is invalidated.
    pub fn destroy(&self, device: &Device) {
        // SAFETY: device, view, image, and memory are valid
        unsafe {
            device.handle().destroy_image_view(self.view, None);
            device.handle().destroy_image(self.image, None);
            device.handle().free_memory(self.memory, None);
        }
    }
}

impl Drop for Image {
    fn drop(&mut self) {
        if self.image != vk::Image::null() {
            eprintln!("WARNING: Image dropped without calling .destroy() - potential memory leak");
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::{Device, DeviceCreateInfo, Instance, InstanceCreateInfo, QueueCreateInfo};

    fn create_test_device() -> (Instance, Device) {
        let instance = Instance::new(InstanceCreateInfo {
            enable_validation: false,
            ..Default::default()
        })
        .unwrap();
        let physical_devices = instance.enumerate_physical_devices().unwrap();
        let physical_device = physical_devices[0];

        let graphics_family = unsafe {
            instance
                .get_physical_device_queue_family_properties(physical_device)
                .iter()
                .enumerate()
                .find(|(_, qf)| qf.queue_flags.contains(vk::QueueFlags::GRAPHICS))
                .map(|(i, _)| i as u32)
                .unwrap()
        };

        let device = Device::new(
            &instance,
            physical_device,
            DeviceCreateInfo {
                queue_create_infos: vec![QueueCreateInfo {
                    queue_family_index: graphics_family,
                    queue_count: 1,
                    queue_priorities: vec![1.0],
                }],
                ..Default::default()
            },
        )
        .unwrap();

        (instance, device)
    }

    #[test]
    fn test_buffer_creation() {
        let (_instance, device) = create_test_device();
        let buffer = Buffer::new(
            &device,
            1024,
            vk::BufferUsageFlags::VERTEX_BUFFER,
            vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
        )
        .unwrap();

        assert_ne!(buffer.handle(), vk::Buffer::null());
        assert_eq!(buffer.size(), 1024);

        buffer.destroy(&device);
    }

    #[test]
    fn test_buffer_mapping() {
        let (_instance, device) = create_test_device();
        let buffer = Buffer::new(
            &device,
            1024,
            vk::BufferUsageFlags::VERTEX_BUFFER,
            vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
        )
        .unwrap();

        // Map and unmap
        unsafe {
            let ptr = buffer.map(&device).unwrap();
            assert!(!ptr.is_null());
            buffer.unmap(&device);
        }

        buffer.destroy(&device);
    }

    #[test]
    fn test_buffer_copy() {
        let (_instance, device) = create_test_device();
        let buffer = Buffer::new(
            &device,
            1024,
            vk::BufferUsageFlags::VERTEX_BUFFER,
            vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
        )
        .unwrap();

        let data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0];
        unsafe {
            buffer.copy_from_slice(&device, &data).unwrap();
        }

        buffer.destroy(&device);
    }

    #[test]
    fn test_image_creation() {
        let (_instance, device) = create_test_device();
        let image = Image::new(
            &device,
            vk::Extent3D {
                width: 512,
                height: 512,
                depth: 1,
            },
            vk::Format::R8G8B8A8_UNORM,
            vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::TRANSFER_DST,
            vk::MemoryPropertyFlags::DEVICE_LOCAL,
        )
        .unwrap();

        assert_ne!(image.handle(), vk::Image::null());
        assert_ne!(image.view(), vk::ImageView::null());
        assert_eq!(image.extent().width, 512);
        assert_eq!(image.extent().height, 512);

        image.destroy(&device);
    }
}