screen_cast/ios/Receiver/Decode/H264VideoToolboxDecoder.swift

import VideoToolbox
import CoreMedia
import CoreGraphics
import Foundation

/// Hardware H.264 decoder using VideoToolbox. The stream is self-describing:
/// the sender repeats SPS/PPS in-band at every keyframe, so no out-of-band
/// codec data is needed. On a keyframe the in-band SPS/PPS (re)establish the
/// stream size; when that size changes, the decode session is recreated — the
/// Android C2 "in-band SPS reconfigure" pattern, which also avoids the startup
/// squish (we never report a placeholder size before the first real keyframe).
///
/// The decode output callback may run on a worker thread, so session and
/// format-description access is guarded by a lock.
final class H264VideoToolboxDecoder {
    private let renderSink: RenderSink
    private let stateLock = NSLock()
    private var session: VTDecompressionSession?
    private var formatDescription: CMVideoFormatDescription?
    private var streamSize = CGSize.zero
    private var realSizeSeen = false

    init(renderSink: RenderSink) {
        self.renderSink = renderSink
    }

    /// The decoded resolution, once the first keyframe configured the session
    /// (nil before that — the caller must not drive layout off a placeholder).
    func outputSize() -> CGSize? {
        guard realSizeSeen, streamSize != .zero else { return nil }
        return streamSize
    }

    /// Decodes one access unit (Annex-B) and renders the output. Returns false
    /// when a decode error occurred and the caller should request a keyframe.
    func decode(accessUnit annexB: [UInt8], rtpTimestamp: Int, isKeyFrame: Bool) -> Bool {
        if isKeyFrame {
            guard let sets = NalExtractor.parameterSets(annexB) else { return false }
            configureIfNeeded(sps: sets.sps, pps: sets.pps)
        }
        stateLock.lock()
        let session = self.session
        let cd = self.formatDescription
        stateLock.unlock()
        guard let session, let cd else { return false }
        guard let avcc = AvccConverter.toAvcc(annexB) else { return false }

        let pts = CMTime(value: CMTimeValue(rtpTimestamp), timescale: 90000)
        guard let blockBuffer = makeBlockBuffer(avcc) else { return false }

        var infoFlags: VTDecodeInfoFlags = []
        let status = VTDecompressionSessionDecodeFrame(
            session,
            blockBuffer,
            isKeyFrame ? kVTDecodeFrameFlags_EnableFastPath : 0,
            &infoFlags,
            pts)
        CFRelease(blockBuffer)

        if status != noErr {
            // Recoverable: the next keyframe (SPS/PPS + IDR) re-primes it.
            teardownSession()
            return false
        }
        realSizeSeen = true
        return true
    }

    func release() {
        stateLock.lock()
        teardownSessionLocked()
        if let cd = formatDescription { CFRelease(cd) }
        formatDescription = nil
        stateLock.unlock()
        renderSink.detach()
        streamSize = .zero
        realSizeSeen = false
    }

    // MARK: - Internals

    private func configureIfNeeded(sps: [UInt8], pps: [UInt8]) {
        guard let cd = H264FormatDescription.create(sps: sps, pps: pps) else { return }
        var dims = CMVideoDimensions()
        guard CMVideoFormatDescriptionGetDimensions(cd, &dims) == noErr else {
            CFRelease(cd)
            return
        }
        let size = CGSize(width: CGFloat(dims.width), height: CGFloat(dims.height))
        stateLock.lock()
        // Same size and a live session: keep it (the sender only changes the
        // stream when the receiver's display size changes).
        if session != nil && size == streamSize {
            CFRelease(cd)
            stateLock.unlock()
            return
        }
        teardownSessionLocked()

        var outSession: VTDecompressionSession?
        let status = VTDecompressionSessionCreate(
            kCFAllocatorDefault,
            cd,
            vtOutputCallback,
            Unmanaged.passUnretained(self).toOpaque(),
            &outSession)
        guard status == noErr, let newSession = outSession else {
            CFRelease(cd)
            stateLock.unlock()
            return
        }
        // Low-latency decode: emit as soon as the frame is complete.
        VTSessionSetProperty(newSession, kVTDecompressionPropertyKey_RealTime, kCFBooleanTrue)

        renderSink.setFormatDescription(cd)
        formatDescription = cd // we hold the +1 from H264FormatDescription.create
        session = newSession
        streamSize = size
        realSizeSeen = false
        stateLock.unlock()
    }

    // Caller holds stateLock.
    private func teardownSessionLocked() {
        if let s = session {
            VTDecompressionSessionInvalidate(s)
            CFRelease(s)
        }
        session = nil
    }

    // Caller does not hold the lock.
    private func teardownSession() {
        stateLock.lock()
        teardownSessionLocked()
        stateLock.unlock()
    }

    private func makeBlockBuffer(_ bytes: [UInt8]) -> CMBlockBuffer? {
        let cfData = Data(bytes) as CFData
        var blockBuffer: CMBlockBuffer?
        let status = CMBlockBufferCreateWithData(kCFAllocatorDefault, cfData, &blockBuffer)
        guard status == noErr, let bb = blockBuffer else { return nil }
        return bb
    }

    // Runs on whatever thread VideoToolbox uses for the output callback.
    private func handleOutput(_ pixelBuffer: CVPixelBuffer?, _ presentationTime: CMTime?) {
        stateLock.lock()
        let cd = formatDescription
        stateLock.unlock()
        guard let cd, let pixelBuffer else { return }
        let pts = presentationTime ?? CMTime(value: 0, timescale: 600)
        var sampleBuffer: CMSampleBuffer?
        let status = CMSampleBufferCreate(kCFAllocatorDefault, nil, 0, pts, .invalid, 1, 0, nil, &sampleBuffer)
        guard status == noErr, let sb = sampleBuffer else { return }
        guard CMSampleBufferSetDataBufferFromPixelBuffer(sb, pixelBuffer) == noErr else {
            CFRelease(sb)
            return
        }
        renderSink.enqueue(sb)
    }
}

/// C-compatible decode output callback; recovers the decoder from the refCon.
private func vtOutputCallback(_ refCon: UnsafeMutableRawPointer?,
                              _ pixelBuffer: CVPixelBuffer?,
                              _ presentationTime: CMTime?,
                              _ duration: CMTime?,
                              _ infoFlags: VTDecodeInfoFlags) {
    guard let refCon = refCon else { return }
    let decoder = Unmanaged<H264VideoToolboxDecoder>.fromOpaque(refCon).takeUnretainedValue()
    decoder.handleOutput(pixelBuffer, presentationTime: presentationTime)
}