punktfunk/clients/apple/Sources/PunktfunkKit/Video/VideoDecoder.swift

// Stage-2 presenter, decode half: explicit VideoToolbox decode of the host's AUs (H.264 /
// HEVC / AV1 — whatever the Welcome resolved).
//
// Stage-1 hands compressed samples to AVSampleBufferDisplayLayer, which decodes AND presents
// internally with no per-frame callback — so neither decode-completion nor present can be
// stamped, and frames can't be hand-paced. Here we drive VTDecompressionSession ourselves: the
// output callback delivers a decoded CVPixelBuffer, we stamp decode-completion, and push it into
// a ready ring the presenter's display link drains. See docs apple-stage2-presenter.md.

import CoreMedia
import CoreVideo
import Foundation
import VideoToolbox

/// A decoded frame's pixels — which present path they take. VideoToolbox codecs deliver a
/// biplanar `CVPixelBuffer` (NV12/P010/444v/x444); the PyroWave Metal decoder delivers three
/// separate R8 plane textures straight off its compute pass (there is no CVPixelBuffer — the
/// planes never leave the GPU).
public enum ReadyImage: @unchecked Sendable {
    /// 8-bit NV12 / 4:4:4 biplanar (SDR) or 10-bit P010 / x444 (HDR), Metal-compatible.
    /// `isHDR` = the stream is BT.2020 PQ and the presenter must configure EDR output.
    case video(CVPixelBuffer, isHDR: Bool)
    #if canImport(Metal)
    /// PyroWave planar output (Y full-res + Cb/Cr half-res, 8-bit SDR) with its precomputed
    /// CSC rows — presented by `MetalVideoPresenter.renderPlanar`.
    case planar(WaveletPlanes)
    #endif
}

/// One decoded frame waiting to be presented. Owns its image (a retained `CVPixelBuffer`, or
/// the PyroWave ring textures) until shown.
public struct ReadyFrame: @unchecked Sendable {
    /// Host capture clock (the AU's pts), in nanoseconds.
    public let ptsNs: UInt64
    /// Client `CLOCK_REALTIME` instant the AU was received (`AccessUnit.receivedNs`, threaded
    /// through the decode via the frame refcon), in nanoseconds. 0 when unknown (a caller that
    /// didn't stamp receipt) — the decode-stage meter then drops the sample via its sanity guard.
    public let receivedNs: Int64
    /// Client `CLOCK_REALTIME` instant decode completed, in nanoseconds.
    public let decodedNs: Int64
    /// The decoded image and which present path it takes.
    public let image: ReadyImage
    /// The AU's wire `user_flags` (`AccessUnit.flags`), threaded through the decode via the frame
    /// context so the re-anchor gate can classify this decoded frame (IDR / RFI anchor / recovery
    /// mark) at present time — the async decode callback has no other access to it. 0 when unknown.
    public let flags: UInt32

    /// The VideoToolbox path's buffer; nil for a PyroWave planar frame. (Kept as the accessor
    /// the decode round-trip tests assert against.)
    public var pixelBuffer: CVPixelBuffer? {
        if case .video(let buffer, _) = image { return buffer }
        return nil
    }

    /// Whether this frame presents on the HDR path. PyroWave planar frames are 8-bit SDR by
    /// contract.
    public var isHDR: Bool {
        if case .video(_, let hdr) = image { return hdr }
        return false
    }
}

/// Per-frame context threaded through the VideoToolbox frame refcon: the AU's receipt instant (for
/// the decode-stage meter) and its wire `user_flags` (for the re-anchor gate). Retained across the
/// async decode and reclaimed exactly once — by the output callback for every frame VideoToolbox
/// accepts, or by `decode`'s error branch for a frame `DecodeFrame` rejected outright (the callback
/// then never fires). A tiny per-frame allocation, the price of smuggling two values (a 64-bit
/// instant plus the flags) through the single `void*` a bit-pattern scalar can't hold.
private final class FrameContext {
    let receivedNs: Int64
    let flags: UInt32
    init(receivedNs: Int64, flags: UInt32) {
        self.receivedNs = receivedNs
        self.flags = flags
    }
}

/// The C output callback can't capture context, so VideoToolbox hands it the refcon we set at
/// session creation — a pointer back to the owning `VideoDecoder`. The per-frame refcon is the
/// retained `FrameContext` set at submit; reclaim it here (balancing `passRetained`) and unpack the
/// AU's receipt instant (for the decode stage) and wire flags (for the re-anchor gate).
private let decoderOutputCallback: VTDecompressionOutputCallback = {
    refcon, frameRefcon, status, _, imageBuffer, pts, _ in
    guard let refcon else { return }
    let ctx = frameRefcon.map { Unmanaged<FrameContext>.fromOpaque($0).takeRetainedValue() }
    Unmanaged<VideoDecoder>.fromOpaque(refcon)
        .takeUnretainedValue()
        .handleDecoded(
            status: status, imageBuffer: imageBuffer, pts: pts,
            receivedNs: ctx?.receivedNs ?? 0, flags: ctx?.flags ?? 0)
}

/// Owns a `VTDecompressionSession` rebuilt whenever the format description changes (every IDR /
/// mode change, the same trigger stage-1 uses). Thread-safe: `decode` runs on the pump thread,
/// the output callback on a VT-managed thread; the only shared mutable state is the session +
/// format, guarded by `lock`. `@unchecked Sendable` — the lock enforces the contract.
public final class VideoDecoder: @unchecked Sendable {
    private let lock = NSLock()
    private var session: VTDecompressionSession?
    private var format: CMVideoFormatDescription?

    /// Called on the VT thread for each successfully decoded frame — stamp + enqueue, don't block.
    private let onDecoded: @Sendable (ReadyFrame) -> Void
    /// Called on the VT thread when a frame fails to decode (bad data / decoder reset) so the
    /// pump can re-gate on the next IDR.
    private let onDecodeError: @Sendable (OSStatus) -> Void

    /// Whether the negotiated stream is full-chroma 4:4:4 (`connection.isChroma444`), set once at
    /// session start before any decode. Selects the 4:4:4 decode pixel format (orthogonal to bit
    /// depth / HDR). Read inside `createSessionLocked` under `lock`.
    private var chroma444 = false

    /// The negotiated codec (`connection.videoCodec`), set once at session start. Drives the
    /// bitstream framing (H.264/HEVC NAL parsing vs AV1 OBU repack). Read under `lock`.
    private var codec: VideoCodec = .hevc

    public init(
        onDecoded: @escaping @Sendable (ReadyFrame) -> Void,
        onDecodeError: @escaping @Sendable (OSStatus) -> Void = { _ in }
    ) {
        self.onDecoded = onDecoded
        self.onDecodeError = onDecodeError
    }

    deinit { teardown() }

    /// Select the chroma subsampling of the decode output (4:2:0 vs full-chroma 4:4:4). Call once at
    /// session start, before decoding, from `connection.isChroma444`. Takes effect on the next
    /// session (re)build. Thread-safe.
    public func setChroma444(_ on: Bool) {
        lock.lock()
        chroma444 = on
        lock.unlock()
    }

    /// Select the negotiated codec (H.264 / HEVC / AV1). Call once at session start, before
    /// decoding, from `connection.videoCodec`. Thread-safe.
    public func setCodec(_ c: VideoCodec) {
        lock.lock()
        codec = c
        lock.unlock()
    }

    /// Submit one AU for asynchronous decode, (re)creating the session if `format` changed. The
    /// caller resolves `format` from the keyframe exactly as stage-1 does
    /// (`VideoCodec.formatDescription(fromKeyframe:)`). Returns false if the session couldn't be
    /// created or the frame couldn't be submitted.
    @discardableResult
    public func decode(au: AccessUnit, format newFormat: CMVideoFormatDescription) -> Bool {
        lock.lock()
        let needsNew: Bool = {
            guard let session, let format else { return true }
            if CMFormatDescriptionEqual(format, otherFormatDescription: newFormat) { return false }
            // A new desc that the live session can still accept (rare for HEVC) avoids a rebuild.
            return !VTDecompressionSessionCanAcceptFormatDescription(session, formatDescription: newFormat)
        }()
        if needsNew, !createSessionLocked(format: newFormat) {
            lock.unlock()
            return false
        }
        // Submit WHILE holding the lock so a concurrent reset()/teardown (main thread) can't
        // invalidate the session between here and DecodeFrame. The VT output callback takes the
        // ring lock, not this one, so there's no re-entrancy. DecodeFrame is async — non-blocking.
        guard let session,
              let sample = codec.sampleBuffer(au: au, format: newFormat)
        else { lock.unlock(); return false }
        var infoOut = VTDecodeInfoFlags()
        // The AU's receipt instant + wire flags ride through as a retained context; the output
        // callback reclaims it. Retain immediately before submit so no early return can leak it.
        let ctx = FrameContext(receivedNs: au.receivedNs, flags: au.flags)
        let refcon = Unmanaged.passRetained(ctx).toOpaque()
        let status = VTDecompressionSessionDecodeFrame(
            session,
            sampleBuffer: sample,
            flags: [._EnableAsynchronousDecompression],
            frameRefcon: refcon,
            infoFlagsOut: &infoOut)
        lock.unlock()
        if status != noErr {
            // DecodeFrame rejected the frame outright — the output callback will NOT fire, so
            // reclaim the context here (balancing passRetained) to avoid leaking it.
            Unmanaged<FrameContext>.fromOpaque(refcon).release()
            onDecodeError(status)
            return false
        }
        return true
    }

    /// Drop the session — the next `decode` rebuilds it. Used on stop and to recover from a
    /// wedged decoder (re-gates on the next in-band parameter sets, like stage-1's flush).
    public func reset() {
        lock.lock()
        teardownLocked()
        lock.unlock()
    }

    private func teardown() {
        lock.lock()
        teardownLocked()
        lock.unlock()
    }

    private func teardownLocked() {
        if let session {
            VTDecompressionSessionWaitForAsynchronousFrames(session)
            VTDecompressionSessionInvalidate(session)
        }
        session = nil
        format = nil
    }

    /// True when `newFormat` carries a PQ (SMPTE ST 2084) or HLG transfer function — i.e. the host
    /// is sending HDR (BT.2020). VideoToolbox populates the transfer-function extension from the
    /// HEVC VUI, so this picks the decode bit depth (10-bit P010/x444 vs 8-bit NV12/444v) from the
    /// stream — and can flip mid-session (a game entering HDR re-inits the host encoder). The
    /// presenter follows the decoded frame's resulting `isHDR`, not the Welcome's latched flag
    /// (`render` reconciles the layer per frame via the idempotent `configure(hdr:)`).
    static func isHDRFormat(_ format: CMVideoFormatDescription) -> Bool {
        guard
            let tf = CMFormatDescriptionGetExtension(
                format, extensionKey: kCMFormatDescriptionExtension_TransferFunction)
        else { return false }
        let s = tf as? String
        return s == (kCMFormatDescriptionTransferFunction_SMPTE_ST_2084_PQ as String)
            || s == (kCMFormatDescriptionTransferFunction_ITU_R_2100_HLG as String)
    }

    /// `lock` held. Replace the session with one for `newFormat`. SDR streams decode to 8-bit NV12;
    /// HDR streams (BT.2020 PQ) decode to 10-bit P010 so the presenter can drive EDR.
    private func createSessionLocked(format newFormat: CMVideoFormatDescription) -> Bool {
        if let session {
            VTDecompressionSessionWaitForAsynchronousFrames(session)
            VTDecompressionSessionInvalidate(session)
        }
        session = nil
        format = nil

        // Decode pixel format is a 2×2 of (chroma, depth/HDR), both biplanar so the presenter binds
        // plane 0 = luma, plane 1 = interleaved chroma uniformly — 4:4:4 just delivers a full-size
        // chroma plane. 10-bit (P010 / `x444`) for HDR (PQ/HLG), 8-bit (NV12 / `444v`) otherwise.
        let hdr = Self.isHDRFormat(newFormat)
        let pixelFormat: OSType = {
            switch (chroma444, hdr) {
            case (false, false): return kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange // NV12
            case (false, true): return kCVPixelFormatType_420YpCbCr10BiPlanarVideoRange // P010
            case (true, false): return kCVPixelFormatType_444YpCbCr8BiPlanarVideoRange // 444v
            case (true, true): return kCVPixelFormatType_444YpCbCr10BiPlanarVideoRange // x444
            }
        }()
        let imageAttrs: [CFString: Any] = [
            kCVPixelBufferMetalCompatibilityKey: true,
            kCVPixelBufferPixelFormatTypeKey: pixelFormat,
        ]
        var callback = VTDecompressionOutputCallbackRecord(
            decompressionOutputCallback: decoderOutputCallback,
            decompressionOutputRefCon: Unmanaged.passUnretained(self).toOpaque())
        // 4:4:4 and AV1 sessions REQUIRE a hardware decoder: both are only advertised when the
        // hardware gate passed (the 4:4:4 probe / `AV1.hardwareDecodeSupported`), so a
        // hardware-incapable mode (e.g. a resolution past a HW ceiling) must fail HERE,
        // synchronously, letting the pump's backstop end the session — rather than silently
        // falling back to a software decoder far too slow for a real-time stream. 4:2:0
        // H.264/HEVC keeps the software fallback (nil spec) as a robustness net.
        let spec: CFDictionary? =
            chroma444 || codec == .av1
            ? [kVTVideoDecoderSpecification_RequireHardwareAcceleratedVideoDecoder: true] as CFDictionary
            : nil
        var newSession: VTDecompressionSession?
        let status = VTDecompressionSessionCreate(
            allocator: kCFAllocatorDefault,
            formatDescription: newFormat,
            decoderSpecification: spec,
            imageBufferAttributes: imageAttrs as CFDictionary,
            outputCallback: &callback,
            decompressionSessionOut: &newSession)
        guard status == noErr, let newSession else { return false }
        // Real-time hint: schedule this session for live-streaming latency rather than
        // throughput/efficiency. Best-effort — decoders that don't support the property
        // return an error, which is fine to ignore.
        VTSessionSetProperty(
            newSession, key: kVTDecompressionPropertyKey_RealTime, value: kCFBooleanTrue)
        session = newSession
        format = newFormat
        return true
    }

    /// VT thread. Stamp decode-completion and enqueue, or report the error. `receivedNs` is the
    /// AU's receipt instant and `flags` its wire `user_flags`, both threaded through the frame refcon
    /// (0 = unknown).
    fileprivate func handleDecoded(
        status: OSStatus, imageBuffer: CVImageBuffer?, pts: CMTime, receivedNs: Int64, flags: UInt32
    ) {
        guard status == noErr, let imageBuffer else {
            onDecodeError(status)
            return
        }
        var ts = timespec()
        clock_gettime(CLOCK_REALTIME, &ts)
        let decodedNs = Int64(ts.tv_sec) * 1_000_000_000 + Int64(ts.tv_nsec)
        // pts was stamped at timescale 1e9 (AnnexB.sampleBuffer); normalize defensively.
        let p = CMTimeConvertScale(pts, timescale: 1_000_000_000, method: .default)
        let ptsNs = p.value > 0 ? UInt64(p.value) : 0
        // HDR iff the decoder produced a 10-bit buffer (we only request a 10-bit format for PQ/HLG
        // streams). Covers 4:2:0 (P010) and 4:4:4 (`x444`), video- and full-range, so a 10-bit 4:4:4
        // HDR frame isn't misclassified as SDR. (The mastering metadata is applied to the presenter's
        // CAMetalLayer via CAEDRMetadata, not to this source buffer — a separate-drawable presenter
        // never composites the source buffer's attachments, so attaching them here would be dead.)
        let fmt = CVPixelBufferGetPixelFormatType(imageBuffer)
        let isHDR =
            fmt == kCVPixelFormatType_420YpCbCr10BiPlanarVideoRange
            || fmt == kCVPixelFormatType_420YpCbCr10BiPlanarFullRange
            || fmt == kCVPixelFormatType_444YpCbCr10BiPlanarVideoRange
            || fmt == kCVPixelFormatType_444YpCbCr10BiPlanarFullRange
        onDecoded(
            ReadyFrame(
                ptsNs: ptsNs, receivedNs: receivedNs, decodedNs: decodedNs,
                image: .video(imageBuffer, isHDR: isHDR), flags: flags))
    }
}