ez-ffmpeg 0.16.0

//! Codec-neutral NAL unit framing: streaming walkers over Annex-B and 4-byte
//! length-prefixed (AVCC) access units.
//!
//! Both walkers validate boundaries, classify NAL types and hand each NAL
//! slice to a caller closure in ONE traversal, allocating nothing themselves —
//! the packet hot path stays allocation-free (the Annex-B rewrite writes into
//! a caller-reserved scratch buffer; the already-length-prefixed and AAC paths
//! do no byte work at all).
//!
//! Boundary conventions mirror libavformat (`nal.c`):
//! * a NAL ends at the next `00 00 01` triple, backing up one byte when the
//!   preceding byte is zero (4-byte start-code attribution) — identical
//!   across FFmpeg versions;
//! * `trailing_zero_8bits` are trimmed from every NAL before it is emitted,
//!   matching FFmpeg master's (n8.2+) `nal_parse_units`. FFmpeg 7.1/8.1
//!   length-prefix the NAL unchanged, carrying the padding into the sample;
//!   the divergence is fixture-only, since real encoders emit no
//!   trailing_zero_8bits.

/// NAL unit types the strict tier cares about (H.264 table 7-1).
pub(crate) const NAL_IDR: u8 = 5;
pub(crate) const NAL_SPS: u8 = 7;
pub(crate) const NAL_PPS: u8 = 8;

/// Byte length of the AVCC length prefix the strict tier emits and accepts
/// (`lengthSizeMinusOne == 3`). Both FFmpeg construction paths (libx264
/// `set_avcc_extradata`, `ff_isom_write_avcc` for Annex-B input) hardcode it.
pub(crate) const NAL_LENGTH_SIZE: usize = 4;

/// What a NAL walk over one access unit observed.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub(crate) struct AuScan {
    pub(crate) nal_count: usize,
    pub(crate) has_idr: bool,
    pub(crate) has_parameter_set: bool,
}

impl AuScan {
    pub(crate) fn note(&mut self, header: u8) {
        let nal_type = header & 0x1F;
        self.nal_count += 1;
        if nal_type == NAL_IDR {
            self.has_idr = true;
        }
        if nal_type == NAL_SPS || nal_type == NAL_PPS {
            self.has_parameter_set = true;
        }
    }
}

/// Walks an Annex-B byte stream NAL by NAL, invoking `on_nal` with each
/// trimmed NAL slice, and returns the accumulated scan.
///
/// Validation (identical to the previous materializing splitter): the payload
/// must begin with a start code (`(00)+ 01`), separators must be well-formed,
/// and a NAL that is empty (including empty after the trailing-zero trim) is
/// rejected.
pub(crate) fn walk_annexb<'a>(
    data: &'a [u8],
    mut on_nal: impl FnMut(&'a [u8]),
) -> Result<AuScan, String> {
    if data.len() < 4 {
        return Err(format!("Annex-B payload too short ({} bytes)", data.len()));
    }
    // Leading start code: (00)+ 01.
    let mut pos = 0;
    while pos < data.len() && data[pos] == 0 {
        pos += 1;
    }
    if pos < 2 || pos >= data.len() || data[pos] != 1 {
        return Err("payload does not begin with an Annex-B start code".to_string());
    }
    pos += 1;

    let mut scan = AuScan::default();
    loop {
        let boundary = find_startcode(data, pos).unwrap_or(data.len());
        // FFmpeg master's (n8.2+) nal_parse_units (libavformat/nal.c) trims
        // trailing_zero_8bits from every NAL before length-prefixing; this
        // trim matches it. FFmpeg 7.1/8.1 length-prefix the NAL unchanged,
        // carrying the padding into the sample. The divergence is
        // fixture-only: real encoders emit no trailing_zero_8bits.
        let mut end = boundary;
        while end > pos && data[end - 1] == 0 {
            end -= 1;
        }
        if end == pos {
            return Err("empty NAL unit".to_string());
        }
        scan.note(data[pos]);
        on_nal(&data[pos..end]);
        if !data[boundary..].iter().any(|&b| b != 0) {
            // Only trailing zeros remain: the stream ends after this NAL.
            break;
        }
        // Skip the separator start code ((00)+ 01; the trim above may have
        // folded this NAL's trailing zeros into the run before the 1).
        let mut next = boundary;
        while next < data.len() && data[next] == 0 {
            next += 1;
        }
        if next >= data.len() || data[next] != 1 {
            return Err("malformed start code between NAL units".to_string());
        }
        pos = next + 1;
        if pos >= data.len() {
            return Err("trailing start code without a NAL unit".to_string());
        }
    }
    Ok(scan)
}

/// First boundary at or after `from` where a start code begins. Mirrors
/// `ff_nal_find_startcode`: locate the next `00 00 01` triple, then back up by
/// exactly one byte when the preceding byte is zero (the leading zero of a
/// 4-byte start code belongs to the start code, not to the previous NAL).
///
/// The scan strides 3 bytes at a time on the classic observation that a
/// triple starting at `i`, `i + 1` or `i + 2` needs `data[i + 2]` to be
/// 1, 0 and 0 respectively: whenever that probe byte exceeds 1 — the
/// overwhelmingly common case in compressed payloads — all three starts
/// are ruled out at once. The probe tests are mask compares on one
/// little-endian `u32` covering `data[i..i + 4]`, keeping the 0-vs-1
/// disambiguation and the hit test branch-free: a branchy probe
/// comparison chain would pay a ~25% mispredicted branch per probe on
/// 0x01-dense payloads and drop below the byte scan on the rewrite walk.
/// The single word load per position also replaces the byte scan's three
/// dependent loads, so even the non-striding arms (a zero probe advancing
/// one byte across long zero runs) stay ahead of it.
///
/// Marked `#[inline(always)]`, not split behind a call: the previous
/// byte-by-byte finder was a small private fn that LLVM folded into
/// [`walk_annexb`], so the walker paid no call per NAL. This wider finder
/// is past the inliner's automatic size threshold (a plain `#[inline]`
/// hint leaves an out-of-line copy that the walker still `call`s once per
/// NAL — verified in release disassembly), which is exactly the
/// separator-dense penalty a real deployment would hit: streams of tiny
/// NALs invoke the finder once per NAL. Forcing the inline restores the
/// parent's lowering — the finder folds into the walker, no boundary —
/// and the benchmark's paired production rows confirm parity by timing
/// this walker against the parent's own inlined composition. Payload bytes
/// come from remote senders, which makes the variant choice worst-corpus
/// driven (see `bench_nal_scan`): word-at-a-time zero skipping (SWAR) and
/// memchr-style candidate search beat this scan on typical entropy but
/// collapse below the byte-by-byte loop on zero-dense and 0x01-dense
/// payloads respectively; this shape stays at or above the byte scan on
/// every measured corpus and both walks.
#[inline(always)]
pub(crate) fn find_startcode(data: &[u8], from: usize) -> Option<usize> {
    let n = data.len();
    if n < 3 {
        return None;
    }
    // Last index where a `00 00 01` triple can still start.
    let end = n - 2;
    let mut i = from;
    while i < end {
        if i + 4 <= n {
            let w = u32::from_le_bytes(data[i..i + 4].try_into().expect("4-byte chunk"));
            if (w & 0x00FF_0000) == 0 {
                // Probe byte data[i + 2] is 0: starts at `i + 1` and
                // `i + 2` stay possible, only the start at `i` (which
                // needs a 1 there) is ruled out; the constant step keeps
                // this arm a tight loop across long zero runs.
                i += 1;
                continue;
            }
            if (w & 0x00FF_FFFF) == 0x0001_0000 {
                // Bytes i..i + 3 are exactly 00 00 01.
                return Some(attribute(data, from, i));
            }
            // Non-zero probe with the triple test missed: a 1 rules out
            // starts at `i + 1` / `i + 2` (both need a 0 there) and the
            // start at `i` was just checked; anything above 1 rules out
            // all three starts directly.
            i += 3;
        } else {
            // i == n - 3 here: the only in-range start whose 4-byte window
            // would overrun the slice. Check the final triple byte-wise.
            if data[i] == 0 && data[i + 1] == 0 && data[i + 2] == 1 {
                return Some(attribute(data, from, i));
            }
            i += 1;
        }
    }
    None
}

/// 4-byte start-code attribution: the zero immediately before a found
/// triple belongs to the start code, not to the previous NAL — unless the
/// triple sits at the scan origin itself.
fn attribute(data: &[u8], from: usize, i: usize) -> usize {
    if i > from && data[i - 1] == 0 {
        i - 1
    } else {
        i
    }
}

/// Walks an already 4-byte length-prefixed AVCC access unit NAL by NAL,
/// invoking `on_nal` with each NAL slice, and returns the accumulated scan.
/// Every NAL must be fully contained; zero-length NAL units and trailing
/// garbage are rejected.
pub(crate) fn walk_length_prefixed<'a>(
    data: &'a [u8],
    mut on_nal: impl FnMut(&'a [u8]),
) -> Result<AuScan, String> {
    let mut scan = AuScan::default();
    let mut pos = 0usize;
    while pos < data.len() {
        if data.len() - pos < NAL_LENGTH_SIZE {
            return Err("truncated NAL length prefix".to_string());
        }
        let len = u32::from_be_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]])
            as usize;
        pos += NAL_LENGTH_SIZE;
        if len == 0 {
            return Err("zero-length NAL unit".to_string());
        }
        if data.len() - pos < len {
            return Err(format!(
                "NAL length {len} overruns the packet ({} bytes remain)",
                data.len() - pos
            ));
        }
        scan.note(data[pos]);
        on_nal(&data[pos..pos + len]);
        pos += len;
    }
    if scan.nal_count == 0 {
        return Err("packet contains no NAL units".to_string());
    }
    Ok(scan)
}

/// Appends one NAL as `u32 big-endian length + payload` to `out` — the write
/// half of the Annex-B rewrite, fed from [`walk_annexb`].
pub(crate) fn push_length_prefixed(nal: &[u8], out: &mut Vec<u8>) {
    debug_assert!(nal.len() <= u32::MAX as usize);
    out.extend_from_slice(&(nal.len() as u32).to_be_bytes());
    out.extend_from_slice(nal);
}

#[cfg(test)]
pub(crate) fn collect_annexb(data: &[u8]) -> Result<Vec<&[u8]>, String> {
    let mut nals = Vec::new();
    walk_annexb(data, |n| nals.push(n))?;
    Ok(nals)
}

#[cfg(test)]
pub(crate) fn collect_length_prefixed(data: &[u8]) -> Result<Vec<&[u8]>, String> {
    let mut nals = Vec::new();
    walk_length_prefixed(data, |n| nals.push(n))?;
    Ok(nals)
}

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

    // Encoder-produced SPS/PPS (x264 via the ffmpeg CLI, Constrained
    // Baseline, 320x240): realistic payloads whose emulation-prevention
    // (00 00 03) sequences the splitter must not mistake for start codes.
    const SPS: &[u8] = &[
        0x67, 0x42, 0xC0, 0x1E, 0xD9, 0x01, 0x41, 0xFB, 0x01, 0x10, 0x00, 0x00, 0x03, 0x00, 0x10,
        0x00, 0x00, 0x03, 0x03, 0x20, 0xF1, 0x62, 0xE4, 0x80,
    ];
    const PPS: &[u8] = &[0x68, 0xCB, 0x83, 0xCB, 0x20];

    fn annexb_config() -> Vec<u8> {
        let mut v = vec![0, 0, 0, 1];
        v.extend_from_slice(SPS);
        v.extend_from_slice(&[0, 0, 1]);
        v.extend_from_slice(PPS);
        v
    }

    #[test]
    fn splits_three_and_four_byte_start_codes() {
        let config = annexb_config();
        let nals = collect_annexb(&config).unwrap();
        assert_eq!(nals, vec![SPS, PPS]);
    }

    #[test]
    fn rejects_garbage_prefix_and_empty_nals() {
        assert!(collect_annexb(&[0x12, 0, 0, 1, 0x67]).is_err());
        assert!(collect_annexb(&[0, 0, 1]).is_err());
        // Two adjacent start codes -> empty NAL.
        assert!(collect_annexb(&[0, 0, 1, 0, 0, 1, 0x41, 0x9A]).is_err());
    }

    #[test]
    fn converts_annexb_au_to_length_prefixed_in_one_walk() {
        let mut au = vec![0, 0, 0, 1, 0x65, 0x88, 0x80];
        au.extend_from_slice(&[0, 0, 1, 0x06, 0x05, 0xFF]);
        let mut out = Vec::new();
        let scan = walk_annexb(&au, |nal| push_length_prefixed(nal, &mut out)).unwrap();
        assert!(scan.has_idr);
        assert!(!scan.has_parameter_set);
        assert_eq!(scan.nal_count, 2);
        assert_eq!(
            out,
            vec![0, 0, 0, 3, 0x65, 0x88, 0x80, 0, 0, 0, 3, 0x06, 0x05, 0xFF]
        );
        // Round-trip: the produced AU walks back into the same NALs.
        let back = collect_length_prefixed(&out).unwrap();
        assert_eq!(back, collect_annexb(&au).unwrap());
    }

    /// FFmpeg-equivalence fixture: trailing_zero_8bits after a NAL (before a
    /// following start code and at stream end) are trimmed before
    /// length-prefixing, like FFmpeg master's (n8.2+) `nal_parse_units`
    /// (libavformat/nal.c); FFmpeg 7.1/8.1 length-prefix the NAL unchanged
    /// and carry the padding into the sample, a divergence that stays
    /// fixture-only because real encoders emit no trailing_zero_8bits.
    #[test]
    fn trailing_zero_bytes_are_trimmed_like_ffmpeg_master() {
        // NAL [65 AA] + two trailing zeros + 3-byte startcode + NAL [06 05].
        let au = vec![0, 0, 0, 1, 0x65, 0xAA, 0, 0, 0, 0, 1, 0x06, 0x05];
        let nals = collect_annexb(&au).unwrap();
        assert_eq!(nals, vec![&[0x65u8, 0xAA][..], &[0x06u8, 0x05][..]]);
        let mut out = Vec::new();
        walk_annexb(&au, |nal| push_length_prefixed(nal, &mut out)).unwrap();
        assert_eq!(out, vec![0, 0, 0, 2, 0x65, 0xAA, 0, 0, 0, 2, 0x06, 0x05]);

        // Trailing zeros at stream end are trimmed from the final NAL too.
        let au = vec![0, 0, 0, 1, 0x65, 0xBB, 0, 0];
        let nals = collect_annexb(&au).unwrap();
        assert_eq!(nals, vec![&[0x65u8, 0xBB][..]]);

        // A NAL reduced to nothing by the trim is still an empty NAL.
        assert!(collect_annexb(&[0, 0, 0, 1, 0, 0]).is_err());
    }

    #[test]
    fn length_prefixed_walk_rejects_overruns() {
        assert!(collect_length_prefixed(&[0, 0, 0, 9, 0x65]).is_err());
        assert!(collect_length_prefixed(&[0, 0, 0, 0]).is_err());
        assert!(collect_length_prefixed(&[]).is_err());
        // Trailing partial prefix after a valid NAL.
        assert!(collect_length_prefixed(&[0, 0, 0, 1, 0x65, 0xFF]).is_err());
    }

    /// The hot-path walkers must not allocate: exercised here by walking with
    /// a counting closure only (no Vec materialization anywhere in the call).
    #[test]
    fn walkers_scan_without_materializing() {
        let mut au = vec![0, 0, 0, 1, 0x65, 0x88, 0x80];
        au.extend_from_slice(&[0, 0, 1, 0x06, 0x05, 0xFF]);
        let mut count = 0usize;
        let scan = walk_annexb(&au, |_| count += 1).unwrap();
        assert_eq!(count, scan.nal_count);
        let lp = vec![0, 0, 0, 1, 0x41];
        let mut count = 0usize;
        let scan = walk_length_prefixed(&lp, |_| count += 1).unwrap();
        assert_eq!((count, scan.nal_count), (1, 1));
    }
}