arcly-stream 0.1.6

//! VVC (H.266) bitstream helpers: NAL classification and SPS resolution
//! extraction. Shares the Annex-B scanner and conformance-crop math with the
//! other NAL codecs.
//!
//! The SPS parser handles both carriages of the profile-tier-level block: when
//! `sps_ptl_dpb_hrd_params_present_flag == 1` it parses `profile_tier_level()`
//! (extracting `general_profile_idc` / `general_tier_flag` / `general_level_idc`
//! and stepping over the variable `general_constraint_info()` by bit count) to
//! reach the dimension fields; when absent, `profile`/`level`/`tier` report `0`.
//! Random-access detection — the signal the engine needs for GOP caching and
//! segmentation — works for all VVC streams.

use super::bitreader::BitReader;
use super::nal::{conformance_dims, iter_nals, unescape_rbsp};
use super::parser::{CodecParser, VideoParams};
use crate::CodecId;

/// NAL unit type for a video parameter set.
pub const NAL_VPS: u8 = 14;
/// NAL unit type for a sequence parameter set.
pub const NAL_SPS: u8 = 15;
/// NAL unit type for a picture parameter set.
pub const NAL_PPS: u8 = 16;

/// VVC NAL unit type, from `nuh_unit_type` in the 2-byte NAL header.
fn nal_type(nal: &[u8]) -> Option<u8> {
    nal.get(1).map(|b| (b >> 3) & 0x1F)
}

/// VVC IRAP VCL NAL types: IDR_W_RADL=7, IDR_N_LP=8, CRA_NUT=9.
fn is_irap(t: u8) -> bool {
    (7..=9).contains(&t)
}

/// Number of fixed `general_constraint_info()` flag bits when `gci_present_flag`
/// is set — the constraint-flag block from `gci_intra_only_constraint_flag`
/// through `gci_no_virtual_boundaries_constraint_flag` (ITU-T H.266 §7.3.3.2),
/// preceding `gci_num_additional_bits`.
const GCI_FIXED_FLAG_BITS: usize = 71;

/// Parse `profile_tier_level(profileTierPresentFlag = 1, max_sublayers_minus1)`
/// (ITU-T H.266 §7.3.3.1), returning `(profile, tier, level)` and leaving the
/// reader positioned at the field that follows. The variable
/// `general_constraint_info()` is stepped over by bit count, not interpreted.
fn parse_profile_tier_level(r: &mut BitReader, max_sublayers_minus1: u32) -> Option<(u8, u8, u8)> {
    let general_profile_idc = r.read_bits(7)? as u8;
    let general_tier_flag = r.read_bit()? as u8;
    let general_level_idc = r.read_bits(8)? as u8;
    let _ptl_frame_only_constraint_flag = r.read_bit()?;
    let _ptl_multilayer_enabled_flag = r.read_bit()?;

    // general_constraint_info() (profileTierPresentFlag == 1).
    if r.read_bit()? == 1 {
        // gci_present_flag: a fixed flag block, then num_additional_bits more.
        r.skip_bits(GCI_FIXED_FLAG_BITS)?;
        let gci_num_additional_bits = r.read_bits(8)? as usize;
        r.skip_bits(gci_num_additional_bits)?;
    }
    r.align_to_byte(); // gci_alignment_zero_bit(s)

    // ptl_sublayer_level_present_flag[ MaxNumSubLayersMinus1-1 .. 0 ].
    let mut sublayer_present = vec![false; max_sublayers_minus1 as usize];
    for i in (0..max_sublayers_minus1 as usize).rev() {
        sublayer_present[i] = r.read_bit()? == 1;
    }
    r.align_to_byte(); // ptl_reserved_zero_bit(s)
    for present in sublayer_present.iter().rev() {
        if *present {
            let _sublayer_level_idc = r.read_bits(8)?;
        }
    }

    // profileTierPresentFlag == 1 → sub-profiles.
    let num_sub_profiles = r.read_bits(8)?;
    for _ in 0..num_sub_profiles {
        let _general_sub_profile_idc = r.read_bits(32)?;
    }

    Some((general_profile_idc, general_tier_flag, general_level_idc))
}

/// Parse a VVC SPS NAL (including its 2-byte header) into [`VideoParams`].
fn parse_sps(nal: &[u8]) -> Option<VideoParams> {
    if nal.len() < 3 || nal_type(nal)? != NAL_SPS {
        return None;
    }
    let rbsp = unescape_rbsp(&nal[2..]);
    let mut r = BitReader::new(&rbsp);

    let _sps_seq_parameter_set_id = r.read_bits(4)?;
    let _sps_video_parameter_set_id = r.read_bits(4)?;
    let sps_max_sublayers_minus1 = r.read_bits(3)?;
    let chroma_format_idc = r.read_bits(2)?;
    let _sps_log2_ctu_size_minus5 = r.read_bits(2)?;
    let ptl_present = r.read_bit()?;
    let (profile, tier, level) = if ptl_present == 1 {
        parse_profile_tier_level(&mut r, sps_max_sublayers_minus1)?
    } else {
        (0, 0, 0)
    };

    let _sps_gdr_enabled_flag = r.read_bit()?;
    let sps_ref_pic_resampling = r.read_bit()?;
    if sps_ref_pic_resampling == 1 {
        let _sps_res_change_in_clvs_allowed = r.read_bit()?;
    }

    let width_luma = r.read_ue()?;
    let height_luma = r.read_ue()?;
    let (mut l, mut rr, mut t, mut b) = (0, 0, 0, 0);
    if r.read_bit()? == 1 {
        l = r.read_ue()?;
        rr = r.read_ue()?;
        t = r.read_ue()?;
        b = r.read_ue()?;
    }
    let (width, height) = conformance_dims(width_luma, height_luma, chroma_format_idc, l, rr, t, b);

    Some(VideoParams {
        width,
        height,
        profile,
        level,
        tier,
        bit_depth: 8,
    })
}

/// Build the **VvcDecoderConfigurationRecord** (`vvcC` box body) from an Annex-B
/// config access unit carrying the VVC parameter sets, alongside the parsed
/// [`VideoParams`]. This is what the fMP4 muxer wraps in a `vvcC` box inside the
/// `vvc1` sample entry.
///
/// Profile/tier/level come from [`Vvc::parse_config`]: real values when the SPS
/// carries a PTL block, or `0` when it does not. Chroma is assumed 8-bit 4:2:0
/// (the SPS color-config fields are read later in the SPS than the muxer needs).
/// The record embeds the actual VPS/SPS/PPS NAL units and fixes `num_sublayers
/// = 1` with a 1-byte constraint-info field so every field stays byte-aligned.
/// Returns `None` when no SPS is present.
pub fn vvcc_config_record(data: &[u8]) -> Option<(VideoParams, Vec<u8>)> {
    let params = Vvc::parse_config(data)?;

    // Collect parameter sets, preserving order within each type.
    let mut vps: Vec<&[u8]> = Vec::new();
    let mut sps: Vec<&[u8]> = Vec::new();
    let mut pps: Vec<&[u8]> = Vec::new();
    for nal in iter_nals(data) {
        match nal_type(nal) {
            Some(NAL_VPS) => vps.push(nal),
            Some(NAL_SPS) => sps.push(nal),
            Some(NAL_PPS) => pps.push(nal),
            _ => {}
        }
    }
    if sps.is_empty() {
        return None;
    }

    let chroma_format_idc = 1u8; // 4:2:0 (VideoParams carries no chroma)
    let bit_depth_minus8 = params.bit_depth.saturating_sub(8);

    // Fixed prefix (ptl_present): config flags + the byte-aligned VvcPTLRecord(1).
    let mut rec = vec![
        0xFF,                              // reserved(5) | LengthSizeMinusOne=3 | ptl_present_flag=1
        0x00,                              // ols_idx(9)=0 | num_sublayers(3)=1 | … (high bits)
        0x10 | (chroma_format_idc & 0x03), // … | constant_frame_rate(2)=0 | chroma(2)
        (bit_depth_minus8 << 5) | 0x1F,    // bit_depth_minus8(3) | reserved(5)
        0x01, // VvcPTLRecord: reserved(2) | num_bytes_constraint_info = 1
        ((params.profile & 0x7F) << 1) | (params.tier & 1), // profile(7) | tier(1)
        params.level, // general_level_idc
        0x00, // frame_only(1) | multilayer(1) | general_constraint_info(6)
        0x00, // num_ptl_sub_profiles = 0
    ];
    rec.extend_from_slice(&(params.width as u16).to_be_bytes()); // max_picture_width
    rec.extend_from_slice(&(params.height as u16).to_be_bytes()); // max_picture_height
    rec.extend_from_slice(&0u16.to_be_bytes()); // avg_frame_rate

    // NAL-unit arrays (VPS/SPS/PPS present), each a complete NAL incl. header.
    let groups = [(NAL_VPS, &vps), (NAL_SPS, &sps), (NAL_PPS, &pps)];
    let present: Vec<_> = groups.iter().filter(|(_, v)| !v.is_empty()).collect();
    rec.push(present.len() as u8); // num_of_arrays
    for (t, nals) in present {
        rec.push(0x80 | (t & 0x1F)); // array_completeness=1 | reserved(2) | NAL_unit_type(5)
        rec.extend_from_slice(&(nals.len() as u16).to_be_bytes());
        for nal in nals.iter() {
            rec.extend_from_slice(&(nal.len() as u16).to_be_bytes());
            rec.extend_from_slice(nal);
        }
    }
    Some((params, rec))
}

/// [`CodecParser`] implementation for VVC / H.266.
pub struct Vvc;

impl CodecParser for Vvc {
    const CODEC: CodecId = CodecId::VVC;

    fn parse_config(data: &[u8]) -> Option<VideoParams> {
        iter_nals(data).find_map(parse_sps)
    }

    fn is_random_access_point(data: &[u8]) -> bool {
        iter_nals(data).any(|n| nal_type(n).is_some_and(|t| t <= 11 && is_irap(t)))
    }

    fn carries_config(data: &[u8]) -> bool {
        iter_nals(data)
            .any(|n| nal_type(n).is_some_and(|t| matches!(t, NAL_VPS | NAL_SPS | NAL_PPS)))
    }

    fn hls_codec_string(p: &VideoParams) -> String {
        format!("vvc1.{}.L{}", p.profile, p.level)
    }
}

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

    /// VVC SPS NAL for 1920x1080, 4:2:0, with the PTL block absent.
    fn sps_1920x1080() -> Vec<u8> {
        let mut w = BitWriter::default();
        w.bits(0, 4); // sps_seq_parameter_set_id
        w.bits(0, 4); // sps_video_parameter_set_id
        w.bits(0, 3); // sps_max_sublayers_minus1
        w.bits(1, 2); // sps_chroma_format_idc = 1 (4:2:0)
        w.bits(0, 2); // sps_log2_ctu_size_minus5
        w.bit(0); // sps_ptl_dpb_hrd_params_present_flag = 0
        w.bit(0); // sps_gdr_enabled_flag
        w.bit(0); // sps_ref_pic_resampling_enabled_flag
        w.ue(1920); // sps_pic_width_max_in_luma_samples
        w.ue(1080); // sps_pic_height_max_in_luma_samples
        w.bit(0); // sps_conformance_window_flag
        let mut nal = vec![0x00u8, 0x79]; // NAL header: nuh_unit_type 15 (SPS)
        nal.extend_from_slice(&w.bytes());
        nal
    }

    /// VVC SPS for 1280x720 with the profile-tier-level block **present**.
    /// `gci_present` selects whether `general_constraint_info` carries its fixed
    /// flag block (true) or is empty (false).
    fn sps_with_ptl(profile: u32, tier: u8, level: u32, gci_present: bool) -> Vec<u8> {
        let mut w = BitWriter::default();
        w.bits(0, 4); // sps_seq_parameter_set_id
        w.bits(0, 4); // sps_video_parameter_set_id
        w.bits(0, 3); // sps_max_sublayers_minus1 = 0
        w.bits(1, 2); // chroma_format_idc = 4:2:0
        w.bits(0, 2); // sps_log2_ctu_size_minus5
        w.bit(1); // sps_ptl_dpb_hrd_params_present_flag = 1
                  // profile_tier_level(1, 0)
        w.bits(profile, 7); // general_profile_idc
        w.bit(tier); // general_tier_flag
        w.bits(level, 8); // general_level_idc
        w.bit(0); // ptl_frame_only_constraint_flag
        w.bit(0); // ptl_multilayer_enabled_flag
        w.bit(gci_present as u8); // gci_present_flag
        if gci_present {
            for _ in 0..71 {
                w.bit(0); // fixed general_constraint_info flag block
            }
            w.bits(0, 8); // gci_num_additional_bits = 0
        }
        w.align(); // gci_alignment_zero_bit(s)
                   // no sublayer flags (max_sublayers_minus1 = 0); align is a no-op
        w.align();
        w.bits(0, 8); // ptl_num_sub_profiles = 0
                      // back in the SPS
        w.bit(0); // sps_gdr_enabled_flag
        w.bit(0); // sps_ref_pic_resampling_enabled_flag
        w.ue(1280); // sps_pic_width_max_in_luma_samples
        w.ue(720); // sps_pic_height_max_in_luma_samples
        w.bit(0); // sps_conformance_window_flag
        let mut nal = vec![0x00u8, 0x79]; // NAL header: nuh_unit_type 15 (SPS)
        nal.extend_from_slice(&w.bytes());
        nal
    }

    #[test]
    fn parse_sps_extracts_resolution() {
        let p = parse_sps(&sps_1920x1080()).expect("parse");
        assert_eq!((p.width, p.height), (1920, 1080));
        // PTL absent → profile/level/tier report 0.
        assert_eq!((p.profile, p.level, p.tier), (0, 0, 0));
    }

    #[test]
    fn parse_sps_with_ptl_extracts_profile_level_and_dims() {
        // gci_present_flag = 0 (the common case).
        let p = parse_sps(&sps_with_ptl(1, 0, 51, false)).expect("parse no-gci");
        assert_eq!((p.width, p.height), (1280, 720));
        assert_eq!((p.profile, p.level, p.tier), (1, 51, 0));

        // gci_present_flag = 1 with a full fixed constraint block + high tier.
        let p = parse_sps(&sps_with_ptl(1, 1, 51, true)).expect("parse gci");
        assert_eq!((p.width, p.height), (1280, 720));
        assert_eq!((p.profile, p.level, p.tier), (1, 51, 1));
        assert_eq!(Vvc::hls_codec_string(&p), "vvc1.1.L51");
    }

    #[test]
    fn classifies_irap_and_config() {
        let mut au = vec![0, 0, 0, 1];
        au.extend_from_slice(&sps_1920x1080());
        // IDR_W_RADL (type 7): byte1 = (7<<3)|1 = 0x39.
        au.extend_from_slice(&[0, 0, 0, 1, 0x00, 0x39, 0xAA]);

        assert!(Vvc::is_random_access_point(&au));
        assert!(Vvc::carries_config(&au));
        assert_eq!(Vvc::parse_config(&au).expect("params").width, 1920);

        // TRAIL (type 0): byte1 = 0x01 → not a RAP.
        assert!(!Vvc::is_random_access_point(&[
            0, 0, 0, 1, 0x00, 0x01, 0xAA
        ]));
    }

    #[test]
    fn vvcc_record_embeds_param_sets_and_dims() {
        // A VPS (type 14, byte1 = (14<<3)|1 = 0x71) plus the SPS.
        let vps = [0x00u8, 0x71, 0xAB, 0xCD];
        let sps = sps_1920x1080();
        let mut au = vec![0, 0, 0, 1];
        au.extend_from_slice(&vps);
        au.extend_from_slice(&[0, 0, 0, 1]);
        au.extend_from_slice(&sps);

        let (params, rec) = vvcc_config_record(&au).expect("record");
        assert_eq!((params.width, params.height), (1920, 1080));
        assert_eq!(rec[0], 0xFF); // LengthSizeMinusOne=3 | ptl_present=1
        assert_eq!(rec[2] & 0x03, 1); // chroma_format_idc = 4:2:0

        // num_of_arrays = 2 (VPS + SPS). It sits after the fixed 18-byte prefix
        // (1 + 2 + 1 + 5 PTL + 6 dims/fps = 15)… assert structurally via search:
        assert!(rec.windows(vps.len()).any(|w| w == vps), "VPS embedded");
        assert!(rec.windows(sps.len()).any(|w| w == sps), "SPS embedded");
        // The SPS array marker: array_completeness=1 | NAL_SPS(15) = 0x8F.
        assert!(rec.contains(&0x8F));
    }
}