screen_cast/ios/Receiver/Rtp/H264Depacketizer.swift

import Foundation

/// Result of feeding one packet to the depacketizer.
struct DepacketizeResult {
    /// Completed access unit (Annex-B with 3-byte start codes) when the frame closed undamaged.
    var accessUnit: [UInt8]?
    /// True when this call discarded a frame as damaged (packet loss or unsupported packetization).
    var frameDropped = false
    /// True when the completed access unit carries SPS/PPS (a keyframe).
    var isKeyFrame = false
}

/// Reassembles RFC 6184 packet streams (single NAL unit packets and FU-A)
/// into Annex-B access units. Packets must arrive in order; frames damaged by
/// sequence gaps or missing fragments are reported via DepacketizeResult.
///
/// Mirrors the C++ `H264Depacketizer` (same state machine and start codes).
final class H264Depacketizer {
    static let fuA = 28

    private var lastSequenceNumber: Int?
    private var frameStarted = false
    private var frameDamaged = false
    private var frameTimestamp = 0
    // Growable byte accumulators: keyframes reach hundreds of KB.
    private var accessUnit: [UInt8] = []
    private var fuActive = false
    private var fuNal: [UInt8] = []

    /// Feed one packet (in sequence order, from the jitter buffer).
    func depacketize(_ packet: RtpPacket) -> DepacketizeResult {
        var result = DepacketizeResult()

        // Track sequence continuity: a gap means packets were lost.
        if let last = lastSequenceNumber {
            let expected = (last + 1) & 0xFFFF
            if packet.header.sequenceNumber != expected {
                fuActive = false
                fuNal.removeAll(keepingCapacity: true)
                if frameStarted { frameDamaged = true }
            }
        }
        lastSequenceNumber = packet.header.sequenceNumber

        // A timestamp change without a closing marker means the previous frame
        // lost its tail and can no longer be recovered.
        if frameStarted && packet.header.timestamp != frameTimestamp {
            dropFrame()
            result.frameDropped = true
        }
        if !frameStarted {
            frameStarted = true
            frameDamaged = false
            frameTimestamp = packet.header.timestamp
            accessUnit.removeAll(keepingCapacity: true)
        }

        let payload = packet.payload
        if !payload.isEmpty {
            let type = Int(payload[0]) & 0x1F
            if type >= 1 && type <= 23 {
                // Single NAL unit packet.
                if fuActive {
                    frameDamaged = true
                    fuActive = false
                    fuNal.removeAll(keepingCapacity: true)
                }
                appendStartCode()
                accessUnit.append(contentsOf: payload)
            } else if type == Self.fuA {
                if payload.count < 2 {
                    frameDamaged = true
                } else {
                    let fuHeader = Int(payload[1])
                    let start = fuHeader & 0x80 != 0
                    let end = fuHeader & 0x40 != 0
                    let fragment = Array(payload[2...])
                    if start {
                        if fuActive {
                            // The previous fragmented NAL lost its end packet.
                            frameDamaged = true
                        }
                        fuActive = true
                        fuNal.removeAll(keepingCapacity: true)
                        // The FU indicator keeps the original NAL's F bit (0) and NRI,
                        // and declares type 28; the FU header carries S/E plus the real type.
                        fuNal.append(UInt8((Int(payload[0]) & 0xE0) | (fuHeader & 0x1F)))
                        fuNal.append(contentsOf: fragment)
                    } else if !fuActive {
                        // Continuation without a start: the head of the NAL is lost.
                        frameDamaged = true
                    } else {
                        fuNal.append(contentsOf: fragment)
                        if end {
                            appendStartCode()
                            accessUnit.append(contentsOf: fuNal)
                            fuActive = false
                            fuNal.removeAll(keepingCapacity: true)
                        }
                    }
                }
            } else {
                // Unsupported packetization mode (STAP-A, MTAP, FU-B).
                frameDamaged = true
            }
        }

        if !packet.header.marker {
            return result
        }

        if fuActive {
            // The marker arrived while a NAL was still fragmented.
            frameDamaged = true
            fuActive = false
            fuNal.removeAll(keepingCapacity: true)
        }

        if !frameDamaged && !accessUnit.isEmpty {
            let unit = accessUnit
            result.accessUnit = unit
            result.isKeyFrame = Self.containsParameterSets(unit)
        } else {
            // The frame that just ended is unusable.
            result.frameDropped = true
        }
        dropFrame()
        return result
    }

    private func appendStartCode() {
        accessUnit.append(0)
        accessUnit.append(0)
        accessUnit.append(1)
    }

    /// The sender repeats SPS/PPS in-band at every keyframe; sniff for NAL types 7/8.
    private static func containsParameterSets(_ unit: [UInt8]) -> Bool {
        guard unit.count >= 4 else { return false }
        var i = 0
        while i <= unit.count - 4 {
            if unit[i] == 0 && unit[i + 1] == 0 && unit[i + 2] == 1 {
                let nalType = Int(unit[i + 3]) & 0x1F
                if nalType == 7 || nalType == 8 {
                    return true
                }
            }
            i += 1
        }
        return false
    }

    private func dropFrame() {
        frameStarted = false
        frameDamaged = false
        accessUnit.removeAll(keepingCapacity: true)
        fuActive = false
        fuNal.removeAll(keepingCapacity: true)
    }
}