vulkan-rust 0.10.0

use std::sync::Arc;

use crate::loader::Loader;
use crate::vk;

/// Wrapper around a `VkDevice` handle and its loaded command table.
///
/// Owns a `Box<DeviceCommands>` containing all device-level function
/// pointers, loaded at construction via `vkGetDeviceProcAddr`. Using the
/// real device handle gives the ICD's direct function pointers, bypassing
/// the loader trampoline, this is the fastest dispatch path in Vulkan.
///
/// Holds an optional reference to the Vulkan shared library so that
/// function pointers remain valid even if the originating `Entry` is
/// dropped. When created via `from_raw_parts`, the caller manages the
/// library lifetime and this field is `None`.
///
/// Does **not** implement `Drop`, the caller must explicitly call
/// `destroy_device` when done. This avoids double-destroy bugs when
/// wrapping externally managed handles via `from_raw_parts`.
///
/// **Guide:** Most device-level operations are covered across
/// [Memory Management](https://hiddentale.github.io/vulkan_rust/concepts/memory.html),
/// [Command Buffers](https://hiddentale.github.io/vulkan_rust/concepts/command-buffers.html),
/// and [Synchronization](https://hiddentale.github.io/vulkan_rust/concepts/synchronization.html).
///
/// # Examples
///
/// ```no_run
/// use vulkan_rust::vk::*;
///
/// # let (entry, instance, device) = vulkan_rust::test_helpers::create_test_device().unwrap();
/// // Use the device to create Vulkan objects.
/// let fence_info = FenceCreateInfo::builder();
/// let fence = unsafe { device.create_fence(&fence_info, None) }
///     .expect("create_fence failed");
///
/// // Clean up (reverse creation order).
/// unsafe {
///     device.destroy_fence(fence, None);
///     device.destroy_device(None);
///     instance.destroy_instance(None);
/// };
/// ```
pub struct Device {
    handle: vk::Device,
    commands: Box<vk::commands::DeviceCommands>,
    _loader: Option<Arc<dyn Loader>>,
}

impl Device {
    /// Internal construction path. Called by `Instance::create_device`.
    ///
    /// # Safety
    ///
    /// - `handle` must be a valid `VkDevice`.
    /// - `get_device_proc_addr` must resolve device-level commands for
    ///   this handle.
    pub(crate) unsafe fn load(
        handle: vk::Device,
        get_device_proc_addr: vk::commands::PFN_vkGetDeviceProcAddr,
        loader: Option<Arc<dyn Loader>>,
    ) -> Self {
        let get_device_proc_addr_fn = get_device_proc_addr.expect("vkGetDeviceProcAddr not loaded");
        // SAFETY: handle is valid per caller contract; transmute converts raw fn ptrs.
        let commands = Box::new(unsafe {
            vk::commands::DeviceCommands::load(|name| {
                std::mem::transmute(get_device_proc_addr_fn(handle, name.as_ptr()))
            })
        });
        Self {
            handle,
            commands,
            _loader: loader,
        }
    }

    /// Wrap a raw handle created externally (OpenXR, middleware, testing).
    ///
    /// # Safety
    ///
    /// - `handle` must be a valid `VkDevice`.
    /// - `get_device_proc_addr` must resolve commands for this device.
    /// - The caller owns the device lifetime.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use vulkan_rust::Device;
    /// # use vulkan_rust::vk::Handle;
    /// # let entry = vulkan_rust::test_helpers::create_test_entry().unwrap();
    ///
    /// // Given a raw device handle and proc addr from an external source:
    /// # let raw_device = vulkan_rust::vk::Device::null();
    /// let device = unsafe {
    ///     Device::from_raw_parts(raw_device, entry.get_device_proc_addr())
    /// };
    /// ```
    pub unsafe fn from_raw_parts(
        handle: vk::Device,
        get_device_proc_addr: vk::commands::PFN_vkGetDeviceProcAddr,
    ) -> Self {
        // SAFETY: forwards caller's safety guarantees to `load`.
        unsafe { Self::load(handle, get_device_proc_addr, None) }
    }

    /// Returns the raw `VkDevice` handle.
    pub fn handle(&self) -> vk::Device {
        self.handle
    }

    /// Returns the loaded device-level command table.
    ///
    /// Use this to call any of the ~200 device-level commands directly,
    /// including those without hand-written ergonomic wrappers.
    pub fn commands(&self) -> &vk::commands::DeviceCommands {
        &self.commands
    }

    /// Create a single graphics pipeline.
    ///
    /// Convenience wrapper around [`create_graphics_pipelines`](Self::create_graphics_pipelines)
    /// for the common single-pipeline case.
    ///
    /// # Safety
    ///
    /// Same requirements as `create_graphics_pipelines`.
    pub unsafe fn create_graphics_pipeline(
        &self,
        pipeline_cache: vk::PipelineCache,
        create_info: &vk::GraphicsPipelineCreateInfo,
        allocator: Option<&vk::AllocationCallbacks>,
    ) -> crate::VkResult<vk::Pipeline> {
        unsafe { self.create_graphics_pipelines(pipeline_cache, &[*create_info], allocator) }
            .map(|v| v[0])
    }

    /// Create a single compute pipeline.
    ///
    /// Convenience wrapper around [`create_compute_pipelines`](Self::create_compute_pipelines)
    /// for the common single-pipeline case.
    ///
    /// # Safety
    ///
    /// Same requirements as `create_compute_pipelines`.
    pub unsafe fn create_compute_pipeline(
        &self,
        pipeline_cache: vk::PipelineCache,
        create_info: &vk::ComputePipelineCreateInfo,
        allocator: Option<&vk::AllocationCallbacks>,
    ) -> crate::VkResult<vk::Pipeline> {
        unsafe { self.create_compute_pipelines(pipeline_cache, &[*create_info], allocator) }
            .map(|v| v[0])
    }

    /// Map a device memory object into host address space.
    ///
    /// Returns a pointer to the mapped region. Wraps
    /// [`vkMapMemory`](https://registry.khronos.org/vulkan/specs/latest/man/html/vkMapMemory.html).
    ///
    /// # Safety
    ///
    /// - `memory` must be a valid, non-mapped `DeviceMemory`.
    /// - `offset + size` must not exceed the allocation size (`WHOLE_SIZE` is valid for `size`).
    /// - The returned pointer is only valid until `unmap_memory` is called.
    pub unsafe fn map_memory(
        &self,
        memory: vk::DeviceMemory,
        offset: u64,
        size: u64,
        flags: vk::MemoryMapFlags,
    ) -> crate::VkResult<*mut core::ffi::c_void> {
        let fp = self.commands().map_memory.expect("vkMapMemory not loaded");
        let mut data: *mut core::ffi::c_void = core::ptr::null_mut();
        crate::error::check(unsafe { fp(self.handle(), memory, offset, size, flags, &mut data) })?;
        Ok(data)
    }

    /// Map a device memory object into host address space (Vulkan 1.4+).
    ///
    /// Returns a pointer to the mapped region. Wraps
    /// [`vkMapMemory2`](https://registry.khronos.org/vulkan/specs/latest/man/html/vkMapMemory2.html).
    ///
    /// # Safety
    ///
    /// - `p_memory_map_info` must describe a valid, non-mapped memory range.
    /// - The returned pointer is only valid until the memory is unmapped.
    pub unsafe fn map_memory2(
        &self,
        p_memory_map_info: &vk::MemoryMapInfo,
    ) -> crate::VkResult<*mut core::ffi::c_void> {
        let fp = self
            .commands()
            .map_memory2
            .expect("vkMapMemory2 not loaded");
        let mut data: *mut core::ffi::c_void = core::ptr::null_mut();
        crate::error::check(unsafe { fp(self.handle(), p_memory_map_info, &mut data) })?;
        Ok(data)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::ffi::c_char;
    use vk::Handle;

    fn fake_handle() -> vk::Device {
        vk::Device::from_raw(0xBEEF)
    }

    /// Stub `vkGetDeviceProcAddr` that returns null for all lookups.
    unsafe extern "system" fn mock_get_device_proc_addr(
        _device: vk::Device,
        _name: *const c_char,
    ) -> vk::PFN_vkVoidFunction {
        None
    }

    #[test]
    fn from_raw_parts_stores_handle() {
        let device =
            unsafe { Device::from_raw_parts(fake_handle(), Some(mock_get_device_proc_addr)) };
        assert_eq!(device.handle().as_raw(), fake_handle().as_raw());
    }

    #[test]
    fn handle_returns_value_from_construction() {
        let device = unsafe { Device::load(fake_handle(), Some(mock_get_device_proc_addr), None) };
        assert_eq!(device.handle().as_raw(), fake_handle().as_raw());
    }

    #[test]
    fn commands_returns_reference() {
        let device = unsafe { Device::load(fake_handle(), Some(mock_get_device_proc_addr), None) };
        // Commands were loaded with a null-returning proc addr, so all
        // function pointers are None,but the reference is valid.
        let _ = device.commands();
    }

    #[test]
    fn load_with_loader_reference() {
        use std::ffi::{CStr, c_void};
        struct DummyLoader;
        unsafe impl Loader for DummyLoader {
            unsafe fn load(&self, _name: &CStr) -> *const c_void {
                std::ptr::null()
            }
        }
        let loader: Arc<dyn Loader> = Arc::new(DummyLoader);
        let device = unsafe {
            Device::load(
                fake_handle(),
                Some(mock_get_device_proc_addr),
                Some(loader.clone()),
            )
        };
        assert_eq!(Arc::strong_count(&loader), 2);
        assert_eq!(device.handle().as_raw(), fake_handle().as_raw());
    }

    #[test]
    fn load_without_loader() {
        let device = unsafe { Device::load(fake_handle(), Some(mock_get_device_proc_addr), None) };
        assert_eq!(device.handle().as_raw(), fake_handle().as_raw());
        // All commands should be None since mock returns null.
        assert!(device.commands().device_wait_idle.is_none());
    }

    #[test]
    fn commands_all_none_from_null_mock() {
        let device =
            unsafe { Device::from_raw_parts(fake_handle(), Some(mock_get_device_proc_addr)) };
        assert!(device.commands().create_buffer.is_none());
        assert!(device.commands().destroy_device.is_none());
        assert!(device.commands().get_device_queue.is_none());
    }

    // -- Rich mock that provides some real function pointers ------------------

    unsafe extern "system" fn mock_device_wait_idle(_device: vk::Device) -> vk::Result {
        vk::Result::SUCCESS
    }

    unsafe extern "system" fn mock_destroy_device(
        _device: vk::Device,
        _p_allocator: *const vk::AllocationCallbacks,
    ) {
    }

    /// `vkGetDeviceProcAddr` that resolves a few commands for richer testing.
    unsafe extern "system" fn rich_get_device_proc_addr(
        _device: vk::Device,
        name: *const c_char,
    ) -> vk::PFN_vkVoidFunction {
        let name = unsafe { std::ffi::CStr::from_ptr(name) };
        match name.to_bytes() {
            b"vkDeviceWaitIdle" => Some(unsafe {
                std::mem::transmute::<
                    unsafe extern "system" fn(vk::Device) -> vk::Result,
                    unsafe extern "system" fn(),
                >(mock_device_wait_idle)
            }),
            b"vkDestroyDevice" => Some(unsafe {
                std::mem::transmute::<
                    unsafe extern "system" fn(vk::Device, *const vk::AllocationCallbacks),
                    unsafe extern "system" fn(),
                >(mock_destroy_device)
            }),
            _ => None,
        }
    }

    #[test]
    fn load_with_rich_mock_populates_some_commands() {
        let device = unsafe { Device::load(fake_handle(), Some(rich_get_device_proc_addr), None) };
        assert!(
            device.commands().device_wait_idle.is_some(),
            "device_wait_idle should be loaded"
        );
        assert!(
            device.commands().destroy_device.is_some(),
            "destroy_device should be loaded"
        );
        // Commands not returned by the mock should still be None.
        assert!(device.commands().create_buffer.is_none());
    }

    #[test]
    fn from_raw_parts_with_rich_mock_populates_commands() {
        let device =
            unsafe { Device::from_raw_parts(fake_handle(), Some(rich_get_device_proc_addr)) };
        assert!(device.commands().device_wait_idle.is_some());
        assert!(device.commands().destroy_device.is_some());
    }

    #[test]
    fn device_wait_idle_succeeds_with_mock() {
        let device =
            unsafe { Device::from_raw_parts(fake_handle(), Some(rich_get_device_proc_addr)) };
        let result = unsafe { device.device_wait_idle() };
        assert!(result.is_ok());
    }

    #[test]
    fn destroy_device_succeeds_with_mock() {
        let device =
            unsafe { Device::from_raw_parts(fake_handle(), Some(rich_get_device_proc_addr)) };
        // Should not panic; the mock destroy is a no-op.
        unsafe { device.destroy_device(None) };
    }

    #[test]
    fn from_raw_parts_stores_no_loader() {
        let device =
            unsafe { Device::from_raw_parts(fake_handle(), Some(mock_get_device_proc_addr)) };
        // from_raw_parts passes None for loader, verify handle is correct.
        assert_eq!(device.handle().as_raw(), fake_handle().as_raw());
    }

    #[test]
    fn load_with_loader_keeps_arc_alive() {
        use std::ffi::{CStr, c_void};
        struct DummyLoader;
        unsafe impl Loader for DummyLoader {
            unsafe fn load(&self, _name: &CStr) -> *const c_void {
                std::ptr::null()
            }
        }
        let loader: Arc<dyn Loader> = Arc::new(DummyLoader);
        let weak = Arc::downgrade(&loader);
        let device =
            unsafe { Device::load(fake_handle(), Some(mock_get_device_proc_addr), Some(loader)) };
        // The Arc should be kept alive by the device.
        assert!(weak.upgrade().is_some(), "loader should still be alive");
        drop(device);
        // After device is dropped, the Arc should be released.
        assert!(weak.upgrade().is_none(), "loader should be dropped");
    }

    #[test]
    #[ignore] // requires Vulkan runtime
    fn device_wait_idle_succeeds() {
        let _vk = crate::VK_TEST_MUTEX.lock().expect("VK_TEST_MUTEX poisoned");
        let (instance, device) = create_real_device();
        unsafe { device.device_wait_idle() }.expect("device_wait_idle failed");
        unsafe { device.destroy_device(None) };
        unsafe { instance.destroy_instance(None) };
    }

    #[test]
    #[ignore] // requires Vulkan runtime
    fn get_device_queue_returns_non_null_queue() {
        let _vk = crate::VK_TEST_MUTEX.lock().expect("VK_TEST_MUTEX poisoned");
        let (instance, device) = create_real_device();
        let queue = unsafe { device.get_device_queue(0, 0) };
        assert!(!queue.is_null(), "expected non-null queue handle");
        unsafe { device.destroy_device(None) };
        unsafe { instance.destroy_instance(None) };
    }

    #[test]
    #[ignore] // requires Vulkan runtime
    fn queue_wait_idle_succeeds() {
        let _vk = crate::VK_TEST_MUTEX.lock().expect("VK_TEST_MUTEX poisoned");
        let (instance, device) = create_real_device();
        let queue = unsafe { device.get_device_queue(0, 0) };
        unsafe { device.queue_wait_idle(queue) }.expect("queue_wait_idle failed");
        unsafe { device.destroy_device(None) };
        unsafe { instance.destroy_instance(None) };
    }

    fn create_real_device() -> (crate::instance::Instance, Device) {
        use crate::entry::Entry;
        use crate::loader::LibloadingLoader;

        let loader = LibloadingLoader::new().expect("failed to load Vulkan");
        let entry = unsafe { Entry::new(loader) }.expect("failed to create Entry");

        let api_version_1_0 = crate::Version::new(1, 0, 0).to_raw();
        let app_info = vk::ApplicationInfo {
            s_type: vk::StructureType::APPLICATION_INFO,
            p_next: std::ptr::null(),
            p_application_name: std::ptr::null(),
            application_version: 0,
            p_engine_name: std::ptr::null(),
            engine_version: 0,
            api_version: api_version_1_0,
        };
        let instance_create_info = vk::InstanceCreateInfo {
            s_type: vk::StructureType::INSTANCE_CREATE_INFO,
            p_next: std::ptr::null(),
            flags: vk::InstanceCreateFlagBits::empty(),
            p_application_info: &app_info,
            enabled_layer_count: 0,
            pp_enabled_layer_names: std::ptr::null(),
            enabled_extension_count: 0,
            pp_enabled_extension_names: std::ptr::null(),
        };
        let instance = unsafe { entry.create_instance(&instance_create_info, None) }
            .expect("failed to create instance");

        let physical_devices = unsafe { instance.enumerate_physical_devices() }
            .expect("failed to enumerate physical devices");
        let physical_device = physical_devices[0];

        let queue_priority = 1.0f32;
        let queue_create_info = vk::DeviceQueueCreateInfo {
            s_type: vk::StructureType::DEVICE_QUEUE_CREATE_INFO,
            p_next: std::ptr::null(),
            flags: vk::DeviceQueueCreateFlagBits::empty(),
            queue_family_index: 0,
            queue_count: 1,
            p_queue_priorities: &queue_priority,
        };
        let device_create_info = vk::DeviceCreateInfo {
            s_type: vk::StructureType::DEVICE_CREATE_INFO,
            p_next: std::ptr::null(),
            flags: 0,
            queue_create_info_count: 1,
            p_queue_create_infos: &queue_create_info,
            enabled_layer_count: 0,
            pp_enabled_layer_names: std::ptr::null(),
            enabled_extension_count: 0,
            pp_enabled_extension_names: std::ptr::null(),
            p_enabled_features: std::ptr::null(),
        };
        let device = unsafe { instance.create_device(physical_device, &device_create_info, None) }
            .expect("failed to create device");

        (instance, device)
    }
}