punktfunk/crates/pf-vkdecode/src/device.rs

//! Borrowed-device wrap: the presenter's live Vulkan handles loaded into ash
//! function tables, plus the queue-lock contract every queue submission runs under.
//!
//! Ownership: everything in [`DeviceHandles`] is BORROWED. This crate never creates
//! and never destroys the instance/device — [`DecodeDevice`]'s ash wrappers are
//! function tables over foreign handles, and dropping them destroys nothing. The
//! objects this crate does create (sessions, images, buffers, pools) are destroyed
//! by their owning structs' `Drop` impls, all of which must run before the borrowed
//! device dies — the same liveness contract FFmpeg's decoder had over the identical
//! handle bundle (`pf-client-core`'s `VulkanDecodeDevice`), now written down.

use ash::vk;
use ash::vk::Handle;

/// The borrowed handles of the presenter's decode-capable device, as raw integers so
/// the type stays FFI-plain (mirrors `pf-client-core`'s `VulkanDecodeDevice`, which
/// adapts into this in WP-C — pf-vkdecode deliberately does not depend on it).
///
/// Caller contract (checked where cheap, otherwise trusted):
/// - All four handles are live, and stay live for the lifetime of every object this
///   crate builds from them (the presenter outlives every session pump).
/// - The instance/device were created with the Vulkan Video decode stack enabled:
///   `VK_KHR_video_queue`, `VK_KHR_video_decode_queue`, and the per-codec extension
///   of every decoder that will be built on the bundle
///   (`VK_KHR_video_decode_h264` for [`crate::VkH264Decoder`],
///   `VK_KHR_video_decode_h265` for [`crate::VkH265Decoder`]), plus the
///   `synchronization2` and `timelineSemaphore` features (the presenter's device
///   meets all of this when it advertises `video_decode`).
/// - `decode_qf`/`decode_queue_index` name a queue with `VIDEO_DECODE_KHR` ops.
///   Which CODEC operations that family advertises is not trusted but READ
///   ([`DecodeDevice::decode_codec_ops`]) and each decoder refuses up front when
///   its own is missing — a physical-device caps query answers for hardware
///   regardless of which extensions the device was created with, so this is the
///   only thing standing between a wrong bundle and `vkCreateVideoSessionKHR` on
///   an unenabled codec. `graphics_qf` is the family the presenter samples on
///   (image sharing crosses the two when they differ).
#[derive(Debug, Clone)]
pub struct DeviceHandles {
    /// `PFN_vkGetInstanceProcAddr` from the loader; everything else is resolved
    /// through it.
    pub get_instance_proc_addr: usize,
    pub instance: usize,
    pub physical_device: usize,
    pub device: usize,
    /// The video-decode queue family.
    pub decode_qf: u32,
    /// Queue index within `decode_qf` this decoder submits on.
    pub decode_queue_index: u32,
    /// The presenter's graphics+present family (the other side of image sharing).
    pub graphics_qf: u32,
}

/// External synchronization for `vkQueueSubmit`: the caller supplies the lock that
/// serializes EVERY submit on the shared device — in WP-C that is pf-client-core's
/// `QueueLock`, the same object the presenter holds around its own submits/presents
/// (the 2026-07-09 `VK_ERROR_DEVICE_LOST` race is why this is a first-class contract
/// and not an afterthought). Tests use [`NoopQueueLock`].
///
/// `lock` blocks until the queue is free and takes it; `unlock` releases it. Use
/// [`QueueSubmitGuard`] rather than calling the pair by hand.
pub trait QueueLock {
    fn lock(&self);
    fn unlock(&self);
}

/// A [`QueueLock`] that guards nothing — for tests and for callers whose decode
/// queue is provably not shared with any other submitter.
#[derive(Debug, Default)]
pub struct NoopQueueLock;

impl QueueLock for NoopQueueLock {
    fn lock(&self) {}
    fn unlock(&self) {}
}

/// RAII scope over a [`QueueLock`]: acquired for exactly the duration of a queue
/// submission, released on drop (including unwinds — though this crate's own paths
/// never panic while holding it).
pub struct QueueSubmitGuard<'a> {
    lock: &'a dyn QueueLock,
}

impl<'a> QueueSubmitGuard<'a> {
    /// Take the queue (blocking until free).
    pub fn acquire(lock: &'a dyn QueueLock) -> Self {
        lock.lock();
        Self { lock }
    }
}

impl Drop for QueueSubmitGuard<'_> {
    fn drop(&mut self) {
        self.lock.unlock();
    }
}

/// A [`DeviceHandles`] bundle that cannot host the decoder being built. Caller
/// bugs (a half-filled bundle) and device gaps (a decode family that does not run
/// this codec) — never stream conditions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DeviceError {
    /// One of the four raw handles is zero/null.
    NullHandle(&'static str),
    /// The decode queue family advertises no decode operation for the codec the
    /// decoder needs — either the device was created without that codec's
    /// extension, or `decode_qf` names the wrong family. Refusing here is what
    /// keeps `vkCreateVideoSessionKHR` from being called with a profile the
    /// device never enabled (the caps query alone would not catch it: it asks the
    /// PHYSICAL device, which answers for the hardware).
    NoCodecOperation {
        family: u32,
        /// The codec, spelled the way the caller would recognize it
        /// (`"H.264 decode"` / `"H.265 decode"`).
        wanted: &'static str,
    },
}

/// A device allocation that cannot proceed. Wraps the raw Vulkan failure OR the
/// memory-type miss that used to be silently papered over with index 0 — a wrong
/// type index is at best an immediate validation error and at worst a mapping of
/// the wrong heap, so a miss is an ERROR here, never a fallback.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum AllocError {
    Vk(vk::Result),
    /// No memory type satisfies (`type_bits`, `flags`) on this device.
    NoMemoryType {
        type_bits: u32,
        flags: vk::MemoryPropertyFlags,
    },
}

impl From<vk::Result> for AllocError {
    fn from(r: vk::Result) -> Self {
        AllocError::Vk(r)
    }
}

/// First memory type matching `bits` and `want` — an [`AllocError::NoMemoryType`]
/// when none does (the encoder's `find_mem` falls back to 0 there; here the miss
/// surfaces).
pub(crate) fn find_memory_type(
    props: &vk::PhysicalDeviceMemoryProperties,
    bits: u32,
    want: vk::MemoryPropertyFlags,
) -> Result<u32, AllocError> {
    for i in 0..props.memory_type_count {
        if (bits & (1 << i)) != 0 && props.memory_types[i as usize].property_flags.contains(want) {
            return Ok(i);
        }
    }
    Err(AllocError::NoMemoryType {
        type_bits: bits,
        flags: want,
    })
}

/// First memory type matching `bits` that also carries `prefer`; when none does,
/// the first type matching `bits` at all. A driver constrains `memoryTypeBits` to
/// where the allocation can legally live — NVIDIA (610.88) reports some video-
/// session bindings host-visible-ONLY, which is spec-legal, so a hard `prefer`
/// requirement there is unsatisfiable by construction. Still an
/// [`AllocError::NoMemoryType`] when `bits` selects nothing whatsoever (that
/// miss-is-error contract stays; only the property preference softens). Mapped
/// staging paths (the bitstream ring) must NOT use this: they require
/// `HOST_VISIBLE|HOST_COHERENT` as a hard property, not a preference.
pub(crate) fn find_memory_type_preferring(
    props: &vk::PhysicalDeviceMemoryProperties,
    bits: u32,
    prefer: vk::MemoryPropertyFlags,
) -> Result<u32, AllocError> {
    match find_memory_type(props, bits, prefer) {
        Ok(index) => Ok(index),
        Err(_) => find_memory_type(props, bits, vk::MemoryPropertyFlags::empty()),
    }
}

impl std::fmt::Display for DeviceError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DeviceError::NullHandle(which) => {
                write!(f, "DeviceHandles.{which} is null — a half-filled bundle")
            }
            DeviceError::NoCodecOperation { family, wanted } => {
                write!(
                    f,
                    "decode queue family {family} advertises no {wanted} operation \
                     (extension not enabled on the device, or the wrong family)"
                )
            }
        }
    }
}

impl std::error::Error for DeviceError {}

/// The borrowed device with ash function tables loaded: the object every other
/// module in this crate makes its Vulkan calls through.
///
/// Clone is cheap-ish (ash tables are plain structs of function pointers) and safe:
/// clones share the same borrowed handles under the same liveness contract.
#[derive(Clone)]
pub struct DecodeDevice {
    instance: ash::Instance,
    device: ash::Device,
    physical_device: vk::PhysicalDevice,
    video_queue_instance: ash::khr::video_queue::Instance,
    video_queue: ash::khr::video_queue::Device,
    video_decode_queue: ash::khr::video_decode_queue::Device,
    decode_queue: vk::Queue,
    decode_qf: u32,
    graphics_qf: u32,
    /// The decode family advertises `queryResultStatusSupport`: per-op
    /// RESULT_STATUS queries are legal in its video coding scopes. FALSE on RADV
    /// (2026-08, .25: recording one anyway hangs the VCN ring) — the decoder
    /// must skip queries entirely there and fall back to timeline-completion
    /// verdicts.
    result_status_queries: bool,
    /// `VkQueueFamilyVideoPropertiesKHR::videoCodecOperations` of the decode
    /// family — which codecs this queue can actually run decode ops for.
    ///
    /// Read from the SAME `vkGetPhysicalDeviceQueueFamilyProperties2` call as the
    /// status-query support (one extra chained struct, no extra round trip). It
    /// is the only honest answer to "may I create a DECODE_H265 session here?":
    /// `vkGetPhysicalDeviceVideoCapabilitiesKHR` is a PHYSICAL-device query and
    /// succeeds on capable hardware whether or not the VkDevice was created with
    /// `VK_KHR_video_decode_h265` enabled, so caps derivation alone would happily
    /// lead into `vkCreateVideoSessionKHR` on an unenabled codec — undefined
    /// behaviour instead of a clean demote to the next decoder rung.
    ///
    /// Empty (no bits) is treated as "this family decodes nothing" and refuses.
    /// That is safe to rely on because every driver hosting Vulkan Video fills
    /// this struct — it is how applications pick a decode queue in the first
    /// place (FFmpeg's own `vulkan_video.c` selects its family by exactly this
    /// field, on the very drivers the shipping FFmpeg-Vulkan rung runs on).
    decode_codec_ops: vk::VideoCodecOperationFlagsKHR,
}

impl DecodeDevice {
    /// Load ash function tables over the borrowed handles.
    ///
    /// # Safety
    ///
    /// The full [`DeviceHandles`] caller contract: live handles (outliving `self`
    /// and everything created through it), the video-decode extensions/features
    /// enabled at creation, and truthful queue-family fields. Null handles are
    /// rejected here; everything else cannot be checked and is trusted.
    pub unsafe fn wrap(handles: &DeviceHandles) -> Result<Self, DeviceError> {
        if handles.get_instance_proc_addr == 0 {
            return Err(DeviceError::NullHandle("get_instance_proc_addr"));
        }
        if handles.instance == 0 {
            return Err(DeviceError::NullHandle("instance"));
        }
        if handles.physical_device == 0 {
            return Err(DeviceError::NullHandle("physical_device"));
        }
        if handles.device == 0 {
            return Err(DeviceError::NullHandle("device"));
        }

        // SAFETY: the usize is non-zero (checked above) and the caller contract says
        // it is the loader's PFN_vkGetInstanceProcAddr; fn pointers and usize share
        // size/ABI on every supported target.
        let gipa: vk::PFN_vkGetInstanceProcAddr = unsafe {
            std::mem::transmute::<usize, vk::PFN_vkGetInstanceProcAddr>(
                handles.get_instance_proc_addr,
            )
        };
        // SAFETY: `gipa` is a valid Vulkan-1.0-conformant loader entry point per the
        // caller contract, valid for the returned Entry's lifetime (handle liveness).
        let entry = unsafe {
            ash::Entry::from_static_fn(ash::StaticFn {
                get_instance_proc_addr: gipa,
            })
        };
        // SAFETY: `handles.instance` is a live VkInstance created through this very
        // loader (caller contract), so loading instance-level functions against it
        // is exactly the ash::Instance::load contract.
        let instance = unsafe {
            ash::Instance::load(
                entry.static_fn(),
                vk::Instance::from_raw(handles.instance as u64),
            )
        };
        // SAFETY: `handles.device` is a live VkDevice of that instance (caller
        // contract) — the ash::Device::load contract.
        let device = unsafe {
            ash::Device::load(
                instance.fp_v1_0(),
                vk::Device::from_raw(handles.device as u64),
            )
        };
        let video_queue_instance = ash::khr::video_queue::Instance::new(&entry, &instance);
        let video_queue = ash::khr::video_queue::Device::new(&instance, &device);
        let video_decode_queue = ash::khr::video_decode_queue::Device::new(&instance, &device);
        // SAFETY: the caller contract guarantees `decode_qf`/`decode_queue_index`
        // name a queue the device was created with.
        let decode_queue =
            unsafe { device.get_device_queue(handles.decode_qf, handles.decode_queue_index) };

        // The two per-family video facts, from ONE query: whether RESULT_STATUS
        // queries are legal here, and which codec operations this family can run
        // (struct field docs for both). An out-of-range decode family — a bundle
        // naming a queue this physical device does not have — answers "no" to
        // both, and the codec check then refuses the decoder outright.
        let physical_device = vk::PhysicalDevice::from_raw(handles.physical_device as u64);
        // SAFETY: live physical device (caller contract); the two-call form fills
        // the chained per-family structs.
        let family_count =
            unsafe { instance.get_physical_device_queue_family_properties2_len(physical_device) };
        let (result_status_queries, decode_codec_ops) = if (handles.decode_qf as usize)
            < family_count
        {
            let mut status_props =
                vec![vk::QueueFamilyQueryResultStatusPropertiesKHR::default(); family_count];
            let mut video_props = vec![vk::QueueFamilyVideoPropertiesKHR::default(); family_count];
            let mut families: Vec<vk::QueueFamilyProperties2<'_>> = status_props
                .iter_mut()
                .zip(video_props.iter_mut())
                .map(|(status, video)| {
                    vk::QueueFamilyProperties2::default()
                        .push_next(status)
                        .push_next(video)
                })
                .collect();
            // SAFETY: as above, arrays sized to the reported count.
            unsafe {
                instance
                    .get_physical_device_queue_family_properties2(physical_device, &mut families)
            };
            drop(families);
            let family = handles.decode_qf as usize;
            (
                status_props[family].query_result_status_support != vk::FALSE,
                video_props[family].video_codec_operations,
            )
        } else {
            (false, vk::VideoCodecOperationFlagsKHR::NONE)
        };

        // `entry` is only the ladder the tables above were loaded through; nothing
        // needs it afterwards (ash tables own their function pointers).
        drop(entry);

        Ok(Self {
            instance,
            device,
            physical_device,
            video_queue_instance,
            video_queue,
            video_decode_queue,
            decode_queue,
            decode_qf: handles.decode_qf,
            graphics_qf: handles.graphics_qf,
            result_status_queries,
            decode_codec_ops,
        })
    }

    pub(crate) fn ash(&self) -> &ash::Device {
        &self.device
    }

    /// Whether the decode family supports per-op RESULT_STATUS queries (struct
    /// field docs — FALSE on RADV, where recording one hangs the VCN).
    pub(crate) fn result_status_queries(&self) -> bool {
        self.result_status_queries
    }

    /// The codec operations the decode family advertises (struct field docs).
    pub fn decode_codec_ops(&self) -> vk::VideoCodecOperationFlagsKHR {
        self.decode_codec_ops
    }

    /// Refuse unless the decode family advertises `op`.
    ///
    /// The decoders' first act, before any caps query: a physical-device caps
    /// query answers for the HARDWARE and would happily green-light a codec the
    /// VkDevice never enabled the extension for, at which point
    /// `vkCreateVideoSessionKHR` is undefined behaviour. This turns that into the
    /// ladder's clean, named demote.
    pub(crate) fn require_codec_op(
        &self,
        op: vk::VideoCodecOperationFlagsKHR,
        what: &'static str,
    ) -> Result<(), DeviceError> {
        if self.decode_codec_ops.contains(op) {
            Ok(())
        } else {
            Err(DeviceError::NoCodecOperation {
                family: self.decode_qf,
                wanted: what,
            })
        }
    }

    pub(crate) fn physical_device(&self) -> vk::PhysicalDevice {
        self.physical_device
    }

    pub(crate) fn video_queue_instance(&self) -> &ash::khr::video_queue::Instance {
        &self.video_queue_instance
    }

    pub(crate) fn video_queue(&self) -> &ash::khr::video_queue::Device {
        &self.video_queue
    }

    pub(crate) fn video_decode_queue(&self) -> &ash::khr::video_decode_queue::Device {
        &self.video_decode_queue
    }

    pub(crate) fn decode_queue(&self) -> vk::Queue {
        self.decode_queue
    }

    pub(crate) fn decode_qf(&self) -> u32 {
        self.decode_qf
    }

    /// The queue families image sharing spans: empty (EXCLUSIVE) when decode and
    /// graphics are one family, both otherwise (CONCURRENT — the presenter samples
    /// decode output on its own family and per-frame ownership transfers would buy
    /// latency for nothing at punktfunk's frame rates).
    pub(crate) fn sharing_families(&self) -> Vec<u32> {
        if self.decode_qf == self.graphics_qf {
            Vec::new()
        } else {
            vec![self.decode_qf, self.graphics_qf]
        }
    }

    /// The device's memory properties (queried fresh; cheap and stateless).
    pub(crate) fn memory_properties(&self) -> vk::PhysicalDeviceMemoryProperties {
        // SAFETY: `physical_device` is live per the DeviceHandles contract; the call
        // fills a plain struct and touches nothing else.
        unsafe {
            self.instance
                .get_physical_device_memory_properties(self.physical_device)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn a_half_filled_bundle_is_rejected_before_any_ffi() {
        let mut handles = DeviceHandles {
            get_instance_proc_addr: 0,
            instance: 1,
            physical_device: 1,
            device: 1,
            decode_qf: 0,
            decode_queue_index: 0,
            graphics_qf: 0,
        };
        // SAFETY: wrap rejects the null handle before making any Vulkan call, so no
        // part of the liveness contract is exercised. (`Err` matched by hand: the
        // Ok side holds ash tables, which carry no Debug for unwrap_err.)
        let result = unsafe { DecodeDevice::wrap(&handles) };
        let Err(err) = result else {
            panic!("a null gipa must be rejected")
        };
        assert_eq!(err, DeviceError::NullHandle("get_instance_proc_addr"));

        handles.get_instance_proc_addr = 1;
        handles.device = 0;
        // SAFETY: as above — the null device handle is rejected before any FFI.
        let result = unsafe { DecodeDevice::wrap(&handles) };
        let Err(err) = result else {
            panic!("a null device must be rejected")
        };
        assert_eq!(err, DeviceError::NullHandle("device"));
    }

    #[test]
    fn a_memory_type_miss_is_an_error_never_a_fallback_to_index_zero() {
        let mut props = vk::PhysicalDeviceMemoryProperties {
            memory_type_count: 2,
            ..Default::default()
        };
        props.memory_types[0].property_flags = vk::MemoryPropertyFlags::DEVICE_LOCAL;
        props.memory_types[1].property_flags =
            vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT;

        // A hit resolves to the matching index, not the first.
        assert_eq!(
            find_memory_type(
                &props,
                0b11,
                vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT
            ),
            Ok(1)
        );
        // A type excluded by the requirement bits does not count as a hit.
        assert_eq!(
            find_memory_type(&props, 0b01, vk::MemoryPropertyFlags::HOST_VISIBLE),
            Err(AllocError::NoMemoryType {
                type_bits: 0b01,
                flags: vk::MemoryPropertyFlags::HOST_VISIBLE
            })
        );
        // Flags nothing advertises: an error carrying the miss, never index 0.
        assert_eq!(
            find_memory_type(&props, 0b11, vk::MemoryPropertyFlags::PROTECTED),
            Err(AllocError::NoMemoryType {
                type_bits: 0b11,
                flags: vk::MemoryPropertyFlags::PROTECTED
            })
        );
    }

    #[test]
    fn preferring_picks_the_preferred_type_and_falls_back_inside_the_bits() {
        let mut props = vk::PhysicalDeviceMemoryProperties {
            memory_type_count: 4,
            ..Default::default()
        };
        props.memory_types[0].property_flags =
            vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT;
        props.memory_types[1].property_flags = vk::MemoryPropertyFlags::DEVICE_LOCAL;
        props.memory_types[2].property_flags = vk::MemoryPropertyFlags::DEVICE_LOCAL;
        props.memory_types[3].property_flags =
            vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT;

        // The preferred property wins over a lower-indexed non-preferred type.
        assert_eq!(
            find_memory_type_preferring(&props, 0b0011, vk::MemoryPropertyFlags::DEVICE_LOCAL),
            Ok(1)
        );
        // The NVIDIA session-binding shape: `memoryTypeBits` names only a
        // host-visible type — honor the bits instead of erroring.
        assert_eq!(
            find_memory_type_preferring(&props, 0b1000, vk::MemoryPropertyFlags::DEVICE_LOCAL),
            Ok(3)
        );
        // Bits selecting nothing remain a hard miss, never index 0.
        assert_eq!(
            find_memory_type_preferring(&props, 0b0000, vk::MemoryPropertyFlags::DEVICE_LOCAL),
            Err(AllocError::NoMemoryType {
                type_bits: 0b0000,
                flags: vk::MemoryPropertyFlags::empty()
            })
        );
    }

    #[test]
    fn the_queue_submit_guard_brackets_the_lock() {
        use std::sync::atomic::AtomicI32;
        use std::sync::atomic::Ordering;

        #[derive(Default)]
        struct CountingLock {
            depth: AtomicI32,
            peak: AtomicI32,
        }
        impl QueueLock for CountingLock {
            fn lock(&self) {
                let d = self.depth.fetch_add(1, Ordering::SeqCst) + 1;
                self.peak.fetch_max(d, Ordering::SeqCst);
            }
            fn unlock(&self) {
                self.depth.fetch_sub(1, Ordering::SeqCst);
            }
        }

        let lock = CountingLock::default();
        {
            let _guard = QueueSubmitGuard::acquire(&lock);
            assert_eq!(lock.depth.load(Ordering::SeqCst), 1);
        }
        assert_eq!(lock.depth.load(Ordering::SeqCst), 0, "released on drop");
        assert_eq!(lock.peak.load(Ordering::SeqCst), 1);
    }
}