arcly-stream 0.1.6

//! AV1 bitstream helpers: OBU sequence-header resolution extraction and
//! random-access detection.

use super::bitreader::BitReader;
use super::obu::iter_obus;
use super::parser::{CodecParser, VideoParams};
use crate::CodecId;

/// OBU type for a sequence header.
pub const OBU_SEQUENCE_HEADER: u8 = 1;

/// Parse an AV1 sequence-header OBU payload into [`VideoParams`].
///
/// Handles the common carriage cases (reduced still-picture headers and
/// timing-info-absent sequence headers). Returns `None` for sequence headers
/// carrying timing/decoder-model info, which require fuller parsing.
fn parse_sequence_header(payload: &[u8]) -> Option<VideoParams> {
    let mut r = BitReader::new(payload);
    let seq_profile = r.read_bits(3)? as u8;
    let _still_picture = r.read_bit()?;
    let reduced_still_picture_header = r.read_bit()?;

    let level;
    if reduced_still_picture_header == 1 {
        level = r.read_bits(5)? as u8; // seq_level_idx[0]
    } else {
        let timing_info_present_flag = r.read_bit()?;
        if timing_info_present_flag == 1 {
            return None; // timing_info()/decoder_model not parsed here
        }
        let initial_display_delay_present_flag = r.read_bit()?;
        let operating_points_cnt_minus1 = r.read_bits(5)?;
        let mut level0 = 0u8;
        for i in 0..=operating_points_cnt_minus1 {
            let _operating_point_idc = r.read_bits(12)?;
            let seq_level_idx = r.read_bits(5)?;
            if seq_level_idx > 7 {
                let _seq_tier = r.read_bit()?;
            }
            if initial_display_delay_present_flag == 1 {
                let present = r.read_bit()?;
                if present == 1 {
                    let _initial_display_delay_minus_1 = r.read_bits(4)?;
                }
            }
            if i == 0 {
                level0 = seq_level_idx as u8;
            }
        }
        level = level0;
    }

    let frame_width_bits_minus_1 = r.read_bits(4)?;
    let frame_height_bits_minus_1 = r.read_bits(4)?;
    let max_frame_width_minus_1 = r.read_bits(frame_width_bits_minus_1 + 1)?;
    let max_frame_height_minus_1 = r.read_bits(frame_height_bits_minus_1 + 1)?;

    Some(VideoParams {
        width: max_frame_width_minus_1 + 1,
        height: max_frame_height_minus_1 + 1,
        profile: seq_profile,
        level,
        tier: 0,
        bit_depth: 8,
    })
}

/// Encode an unsigned LEB128 integer onto `out`.
fn leb128_encode(mut v: u64, out: &mut Vec<u8>) {
    loop {
        let mut byte = (v & 0x7f) as u8;
        v >>= 7;
        if v != 0 {
            byte |= 0x80;
        }
        out.push(byte);
        if v == 0 {
            break;
        }
    }
}

/// Build the **AV1CodecConfigurationRecord** (`av1C` box body) from a config
/// access unit — the temporal unit carrying the sequence-header OBU — alongside
/// the parsed [`VideoParams`]. This is what the fMP4 muxer wraps in an `av1C`
/// box inside the `av01` sample entry.
///
/// Color configuration (bit depth, chroma subsampling) is not parsed from the
/// sequence header here; the record uses the **8-bit 4:2:0** defaults that match
/// `seq_profile` 0 streams. Returns `None` if no sequence header is present.
pub fn av1c_config_record(data: &[u8]) -> Option<(VideoParams, Vec<u8>)> {
    let seq_payload = iter_obus(data)
        .find(|o| o.obu_type == OBU_SEQUENCE_HEADER)?
        .payload;
    let params = parse_sequence_header(seq_payload)?;

    let mut rec = Vec::with_capacity(4 + 2 + seq_payload.len());
    rec.push(0x81); // marker(1) | version(7) = 1
    rec.push(((params.profile & 0x07) << 5) | (params.level & 0x1f)); // seq_profile | seq_level_idx_0
                                                                      // seq_tier_0 | high_bitdepth(0) | twelve_bit(0) | monochrome(0)
                                                                      //   | chroma_subsampling_x(1) | chroma_subsampling_y(1) | chroma_sample_position(0)
    rec.push(((params.tier & 1) << 7) | 0x0C);
    rec.push(0x00); // reserved(3) | initial_presentation_delay_present(0) | reserved(4)
                    // configOBUs: the sequence-header OBU, reframed in low-overhead (sized) form.
    rec.push(0x0A); // obu_type = 1 (seq header), obu_has_size_field = 1
    leb128_encode(seq_payload.len() as u64, &mut rec);
    rec.extend_from_slice(seq_payload);
    Some((params, rec))
}

/// [`CodecParser`] implementation for AV1.
pub struct Av1;

impl CodecParser for Av1 {
    const CODEC: CodecId = CodecId::AV1;

    fn parse_config(data: &[u8]) -> Option<VideoParams> {
        iter_obus(data)
            .find(|o| o.obu_type == OBU_SEQUENCE_HEADER)
            .and_then(|o| parse_sequence_header(o.payload))
    }

    fn is_random_access_point(data: &[u8]) -> bool {
        // A sequence-header OBU is only sent at a random-access point; its
        // presence marks a clean decode start (and the keyframe that follows).
        iter_obus(data).any(|o| o.obu_type == OBU_SEQUENCE_HEADER)
    }

    fn carries_config(data: &[u8]) -> bool {
        Self::is_random_access_point(data)
    }

    fn hls_codec_string(p: &VideoParams) -> String {
        // av01.<profile>.<level><tier=M>.<bit_depth>
        format!("av01.{}.{:02}M.{:02}", p.profile, p.level, p.bit_depth)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::codec::testutil::BitWriter;

    /// Build a low-overhead sequence-header OBU for 1920x1080, profile 0.
    fn seq_header_obu_1920x1080() -> Vec<u8> {
        let mut w = BitWriter::default();
        w.bits(0, 3); // seq_profile
        w.bit(0); // still_picture
        w.bit(0); // reduced_still_picture_header
        w.bit(0); // timing_info_present_flag
        w.bit(0); // initial_display_delay_present_flag
        w.bits(0, 5); // operating_points_cnt_minus1 = 0
        w.bits(0, 12); // operating_point_idc[0]
        w.bits(1, 5); // seq_level_idx[0] = 1 (<= 7 → no seq_tier)
        w.bits(11, 4); // frame_width_bits_minus_1 = 11 → 12-bit width field
        w.bits(11, 4); // frame_height_bits_minus_1 = 11 → 12-bit height field
        w.bits(1919, 12); // max_frame_width_minus_1
        w.bits(1079, 12); // max_frame_height_minus_1
        w.align();
        let payload = w.bytes();

        let mut obu = vec![0x0A]; // header: type=1 (seq header), has_size_field=1
        obu.push(payload.len() as u8); // LEB128 size (< 128)
        obu.extend_from_slice(&payload);
        obu
    }

    #[test]
    fn parse_sequence_header_extracts_resolution() {
        let p = parse_sequence_header(&seq_header_obu_1920x1080()[2..]).expect("parse");
        assert_eq!((p.width, p.height), (1920, 1080));
        assert_eq!((p.profile, p.level), (0, 1));
    }

    #[test]
    fn classifies_and_extracts_via_obus() {
        // Temporal delimiter + sequence header + a frame OBU.
        let mut tu = vec![0x12, 0x00]; // OBU_TEMPORAL_DELIMITER (type 2), size 0
        tu.extend_from_slice(&seq_header_obu_1920x1080());
        tu.extend_from_slice(&[0x32, 0x02, 0xAA, 0xBB]); // OBU_FRAME (type 6), size 2

        assert!(Av1::is_random_access_point(&tu));
        assert!(Av1::carries_config(&tu));
        let p = Av1::parse_config(&tu).expect("params");
        assert_eq!((p.width, p.height), (1920, 1080));
        assert_eq!(Av1::hls_codec_string(&p), "av01.0.01M.08");

        // A temporal unit with no sequence header is not a RAP.
        let no_seq = [0x12, 0x00, 0x32, 0x02, 0xAA, 0xBB];
        assert!(!Av1::is_random_access_point(&no_seq));
    }

    #[test]
    fn av1c_record_carries_profile_level_and_seq_header() {
        let mut tu = vec![0x12, 0x00]; // temporal delimiter
        let seq = seq_header_obu_1920x1080();
        tu.extend_from_slice(&seq);
        let (params, rec) = av1c_config_record(&tu).expect("record");
        assert_eq!((params.width, params.height), (1920, 1080));
        assert_eq!(rec[0], 0x81); // marker | version
        assert_eq!(rec[1], 0x01); // seq_profile 0 (<<5) | seq_level_idx 1
        assert_eq!(rec[2], 0x0C); // 8-bit 4:2:0 defaults
                                  // configOBUs: the reframed sequence-header OBU follows the 4-byte header.
        assert_eq!(rec[4], 0x0A); // seq-header OBU header, has_size_field
        assert!(rec.ends_with(&seq[2..])); // the seq header payload
    }
}