shdrlib 0.1.5

//! Descriptor module - Descriptor sets and layouts
//!
//! **WARNING**: CORE objects do NOT enforce lifetime dependencies.
//! Descriptor sets must not outlive their pool, pools must not outlive the device.

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

/// Descriptor pool operation error
#[derive(Debug, Error)]
pub enum DescriptorPoolError {
    #[error("Descriptor pool creation failed: {0}")]
    CreationFailed(vk::Result),

    #[error("Descriptor set allocation failed: {0}")]
    AllocationFailed(vk::Result),

    #[error("Descriptor pool reset failed: {0}")]
    ResetFailed(vk::Result),
}

/// Descriptor operation error
#[derive(Debug, Error)]
pub enum DescriptorError {
    #[error("Descriptor layout creation failed: {0}")]
    LayoutCreationFailed(vk::Result),

    #[error("Descriptor update failed: {0}")]
    UpdateFailed(vk::Result),
}

/// Descriptor pool wrapper
///
/// Allocates descriptor sets.
///
/// # Safety
///
/// This object does NOT enforce that the device outlives it.
/// Descriptor sets allocated from this pool are invalidated when the pool is destroyed.
pub struct DescriptorPool {
    pool: vk::DescriptorPool,
}

impl DescriptorPool {
    /// Create a new descriptor pool
    ///
    /// # Arguments
    ///
    /// * `device` - The device to create the pool on
    /// * `max_sets` - Maximum number of sets that can be allocated
    /// * `pool_sizes` - Pool size descriptors
    /// * `flags` - Pool creation flags
    ///
    /// # Errors
    ///
    /// Returns `DescriptorPoolError::CreationFailed` if creation fails.
    pub fn new(
        device: &Device,
        max_sets: u32,
        pool_sizes: &[vk::DescriptorPoolSize],
        flags: vk::DescriptorPoolCreateFlags,
    ) -> Result<Self, DescriptorPoolError> {
        let pool_info = vk::DescriptorPoolCreateInfo {
            flags,
            max_sets,
            pool_size_count: pool_sizes.len() as u32,
            p_pool_sizes: pool_sizes.as_ptr(),
            ..Default::default()
        };

        // SAFETY: device is valid, pool_info is properly initialized
        let pool = unsafe {
            device
                .handle()
                .create_descriptor_pool(&pool_info, None)
                .map_err(DescriptorPoolError::CreationFailed)?
        };

        Ok(Self { pool })
    }

    /// Create a DescriptorPool from an existing Vulkan handle
    ///
    /// # Safety
    ///
    /// Caller must ensure the pool is valid.
    /// This is primarily for internal use by EX tier.
    #[inline]
    pub(crate) fn from_raw(pool: vk::DescriptorPool) -> Self {
        Self { pool }
    }

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

    /// Allocate descriptor sets from the pool
    ///
    /// # Arguments
    ///
    /// * `device` - The device handle
    /// * `layouts` - Descriptor set layouts to allocate
    ///
    /// # Errors
    ///
    /// Returns `DescriptorPoolError::AllocationFailed` if allocation fails.
    pub fn allocate(
        &self,
        device: &Device,
        layouts: &[vk::DescriptorSetLayout],
    ) -> Result<Vec<DescriptorSet>, DescriptorPoolError> {
        let alloc_info = vk::DescriptorSetAllocateInfo {
            descriptor_pool: self.pool,
            descriptor_set_count: layouts.len() as u32,
            p_set_layouts: layouts.as_ptr(),
            ..Default::default()
        };

        // SAFETY: device and pool are valid, alloc_info is properly initialized
        let sets = unsafe {
            device
                .handle()
                .allocate_descriptor_sets(&alloc_info)
                .map_err(DescriptorPoolError::AllocationFailed)?
        };

        Ok(sets.into_iter().map(|set| DescriptorSet { set }).collect())
    }

    /// Reset the descriptor pool
    ///
    /// This frees all descriptor sets allocated from the pool.
    ///
    /// # Errors
    ///
    /// Returns `DescriptorPoolError::ResetFailed` if reset fails.
    pub fn reset(
        &self,
        device: &Device,
        flags: vk::DescriptorPoolResetFlags,
    ) -> Result<(), DescriptorPoolError> {
        // SAFETY: device and pool are valid
        unsafe {
            device
                .handle()
                .reset_descriptor_pool(self.pool, flags)
                .map_err(DescriptorPoolError::ResetFailed)
        }
    }

    /// Destroy the pool manually
    pub fn destroy(&mut self, device: &Device) {
        // SAFETY: device and pool are valid
        unsafe {
            device.handle().destroy_descriptor_pool(self.pool, None);
        }
        self.pool = vk::DescriptorPool::null();
    }
}

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

/// Descriptor set wrapper
///
/// Represents an allocated descriptor set.
///
/// **NOTE**: DescriptorSet has NO Drop impl - pool owns and frees them.
pub struct DescriptorSet {
    set: vk::DescriptorSet,
}

impl DescriptorSet {
    /// Get the raw Vulkan descriptor set handle
    #[inline]
    pub fn handle(&self) -> vk::DescriptorSet {
        self.set
    }
}

// No Drop implementation - descriptor sets are freed when pool is destroyed/reset

/// Descriptor set layout wrapper
///
/// Describes the structure of descriptor sets.
pub struct DescriptorSetLayout {
    layout: vk::DescriptorSetLayout,
}

impl DescriptorSetLayout {
    /// Create a new descriptor set layout
    ///
    /// # Arguments
    ///
    /// * `device` - The device to create the layout on
    /// * `bindings` - Layout binding descriptors
    ///
    /// # Errors
    ///
    /// Returns `DescriptorError::LayoutCreationFailed` if creation fails.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use shdrlib::core::{DescriptorSetLayout, Device};
    /// use ash::vk;
    ///
    /// # fn example(device: &Device) -> Result<(), shdrlib::core::DescriptorError> {
    /// let bindings = [
    ///     vk::DescriptorSetLayoutBinding {
    ///         binding: 0,
    ///         descriptor_type: vk::DescriptorType::UNIFORM_BUFFER,
    ///         descriptor_count: 1,
    ///         stage_flags: vk::ShaderStageFlags::VERTEX,
    ///         ..Default::default()
    ///     },
    /// ];
    ///
    /// let layout = DescriptorSetLayout::new(device, &bindings)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn new(
        device: &Device,
        bindings: &[vk::DescriptorSetLayoutBinding],
    ) -> Result<Self, DescriptorError> {
        let layout_info = vk::DescriptorSetLayoutCreateInfo {
            binding_count: bindings.len() as u32,
            p_bindings: bindings.as_ptr(),
            ..Default::default()
        };

        // SAFETY: device is valid, layout_info is properly initialized
        let layout = unsafe {
            device
                .handle()
                .create_descriptor_set_layout(&layout_info, None)
                .map_err(DescriptorError::LayoutCreationFailed)?
        };

        Ok(Self { layout })
    }

    /// Create a DescriptorSetLayout from an existing Vulkan handle
    ///
    /// # Safety
    ///
    /// Caller must ensure the layout is valid.
    /// This is primarily for internal use by EX tier.
    #[inline]
    pub(crate) fn from_raw(layout: vk::DescriptorSetLayout) -> Self {
        Self { layout }
    }

    /// Get the raw Vulkan descriptor set layout handle
    #[inline]
    pub fn handle(&self) -> vk::DescriptorSetLayout {
        self.layout
    }

    /// Destroy the layout manually
    pub fn destroy(&mut self, device: &Device) {
        // SAFETY: device and layout are valid
        unsafe {
            device
                .handle()
                .destroy_descriptor_set_layout(self.layout, None);
        }
        self.layout = vk::DescriptorSetLayout::null();
    }
}

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

/// Descriptor write operation
///
/// Describes an update to a descriptor set.
pub struct DescriptorWrite {
    pub set: vk::DescriptorSet,
    pub binding: u32,
    pub array_element: u32,
    pub descriptor_type: vk::DescriptorType,
    pub resource: DescriptorResource,
}

/// Descriptor resource types
#[derive(Clone)]
pub enum DescriptorResource {
    /// Buffer descriptor
    Buffer {
        buffer: vk::Buffer,
        offset: u64,
        range: u64,
    },
    /// Image descriptor
    Image {
        view: vk::ImageView,
        layout: vk::ImageLayout,
        sampler: Option<vk::Sampler>,
    },
}

/// Update descriptor sets
///
/// # Arguments
///
/// * `device` - The device handle
/// * `writes` - Descriptor write operations to perform
///
/// # Example
///
/// ```rust,no_run
/// use shdrlib::core::{update_descriptor_sets, DescriptorWrite, DescriptorResource, Device};
/// use ash::vk;
///
/// # fn example(
/// #     device: &Device,
/// #     descriptor_set: vk::DescriptorSet,
/// #     buffer: vk::Buffer,
/// # ) {
/// let writes = vec![DescriptorWrite {
///     set: descriptor_set,
///     binding: 0,
///     array_element: 0,
///     descriptor_type: vk::DescriptorType::UNIFORM_BUFFER,
///     resource: DescriptorResource::Buffer {
///         buffer,
///         offset: 0,
///         range: 256,
///     },
/// }];
///
/// update_descriptor_sets(device, &writes);
/// # }
/// ```
pub fn update_descriptor_sets(device: &Device, writes: &[DescriptorWrite]) {
    let mut buffer_infos = Vec::new();
    let mut image_infos = Vec::new();
    let mut write_sets = Vec::new();

    for write in writes {
        match &write.resource {
            DescriptorResource::Buffer {
                buffer,
                offset,
                range,
            } => {
                buffer_infos.push(vk::DescriptorBufferInfo {
                    buffer: *buffer,
                    offset: *offset,
                    range: *range,
                });
                write_sets.push(vk::WriteDescriptorSet {
                    dst_set: write.set,
                    dst_binding: write.binding,
                    dst_array_element: write.array_element,
                    descriptor_count: 1,
                    descriptor_type: write.descriptor_type,
                    p_buffer_info: &buffer_infos[buffer_infos.len() - 1],
                    ..Default::default()
                });
            }
            DescriptorResource::Image {
                view,
                layout,
                sampler,
            } => {
                image_infos.push(vk::DescriptorImageInfo {
                    image_view: *view,
                    image_layout: *layout,
                    sampler: sampler.unwrap_or(vk::Sampler::null()),
                });
                write_sets.push(vk::WriteDescriptorSet {
                    dst_set: write.set,
                    dst_binding: write.binding,
                    dst_array_element: write.array_element,
                    descriptor_count: 1,
                    descriptor_type: write.descriptor_type,
                    p_image_info: &image_infos[image_infos.len() - 1],
                    ..Default::default()
                });
            }
        }
    }

    // SAFETY: device is valid, write_sets and infos are properly initialized
    unsafe {
        device.handle().update_descriptor_sets(&write_sets, &[]);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::{
        Buffer, 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_descriptor_set_layout_creation() {
        let (_instance, device) = create_test_device();

        let bindings = [vk::DescriptorSetLayoutBinding {
            binding: 0,
            descriptor_type: vk::DescriptorType::UNIFORM_BUFFER,
            descriptor_count: 1,
            stage_flags: vk::ShaderStageFlags::VERTEX,
            ..Default::default()
        }];

        let mut layout = DescriptorSetLayout::new(&device, &bindings).unwrap();
        assert_ne!(layout.handle(), vk::DescriptorSetLayout::null());
        layout.destroy(&device);
    }

    #[test]
    fn test_descriptor_pool_creation() {
        let (_instance, device) = create_test_device();

        let pool_sizes = [vk::DescriptorPoolSize {
            ty: vk::DescriptorType::UNIFORM_BUFFER,
            descriptor_count: 10,
        }];

        let mut pool = DescriptorPool::new(
            &device,
            10,
            &pool_sizes,
            vk::DescriptorPoolCreateFlags::empty(),
        )
        .unwrap();

        assert_ne!(pool.handle(), vk::DescriptorPool::null());
        pool.destroy(&device);
    }

    #[test]
    fn test_descriptor_set_allocation() {
        let (_instance, device) = create_test_device();

        let bindings = [vk::DescriptorSetLayoutBinding {
            binding: 0,
            descriptor_type: vk::DescriptorType::UNIFORM_BUFFER,
            descriptor_count: 1,
            stage_flags: vk::ShaderStageFlags::VERTEX,
            ..Default::default()
        }];

        let mut layout = DescriptorSetLayout::new(&device, &bindings).unwrap();

        let pool_sizes = [vk::DescriptorPoolSize {
            ty: vk::DescriptorType::UNIFORM_BUFFER,
            descriptor_count: 10,
        }];

        let mut pool = DescriptorPool::new(
            &device,
            10,
            &pool_sizes,
            vk::DescriptorPoolCreateFlags::empty(),
        )
        .unwrap();

        let sets = pool.allocate(&device, &[layout.handle()]).unwrap();
        assert_eq!(sets.len(), 1);
        assert_ne!(sets[0].handle(), vk::DescriptorSet::null());

        pool.destroy(&device);
        layout.destroy(&device);
    }

    #[test]
    fn test_descriptor_update() {
        let (_instance, device) = create_test_device();

        let bindings = [vk::DescriptorSetLayoutBinding {
            binding: 0,
            descriptor_type: vk::DescriptorType::UNIFORM_BUFFER,
            descriptor_count: 1,
            stage_flags: vk::ShaderStageFlags::VERTEX,
            ..Default::default()
        }];

        let mut layout = DescriptorSetLayout::new(&device, &bindings).unwrap();

        let pool_sizes = [vk::DescriptorPoolSize {
            ty: vk::DescriptorType::UNIFORM_BUFFER,
            descriptor_count: 10,
        }];

        let mut pool = DescriptorPool::new(
            &device,
            10,
            &pool_sizes,
            vk::DescriptorPoolCreateFlags::empty(),
        )
        .unwrap();

        let sets = pool.allocate(&device, &[layout.handle()]).unwrap();

        // Create a buffer to bind
        let buffer = Buffer::new(
            &device,
            256,
            vk::BufferUsageFlags::UNIFORM_BUFFER,
            vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
        )
        .unwrap();

        // Update descriptor set
        let writes = vec![DescriptorWrite {
            set: sets[0].handle(),
            binding: 0,
            array_element: 0,
            descriptor_type: vk::DescriptorType::UNIFORM_BUFFER,
            resource: DescriptorResource::Buffer {
                buffer: buffer.handle(),
                offset: 0,
                range: 256,
            },
        }];

        update_descriptor_sets(&device, &writes);

        buffer.destroy(&device);
        pool.destroy(&device);
        layout.destroy(&device);
    }
}