arcly-stream 0.1.7

//! H.265 (HEVC) bitstream helpers: NAL classification and SPS resolution
//! extraction. Shares the Annex-B scanner and conformance-crop math with H.264.

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 = 32;
/// NAL unit type for a sequence parameter set.
pub const NAL_SPS: u8 = 33;
/// NAL unit type for a picture parameter set.
pub const NAL_PPS: u8 = 34;

/// HEVC NAL unit type, read from the 2-byte NAL header.
fn nal_type(nal: &[u8]) -> Option<u8> {
    nal.first().map(|b| (b >> 1) & 0x3F)
}

/// HEVC IRAP (random-access) VCL NAL types are `16..=23`
/// (BLA/IDR/CRA). A frame containing any such VCL NAL is a random-access point.
fn is_irap(t: u8) -> bool {
    (16..=23).contains(&t)
}

/// Skip the HEVC `profile_tier_level`, capturing `(profile_idc, tier, level_idc)`.
fn read_profile_tier_level(r: &mut BitReader, max_sub_layers_minus1: u32) -> Option<(u8, u8, u8)> {
    // general profile portion is 88 bits, then 8-bit general_level_idc.
    let _general_profile_space = r.read_bits(2)?;
    let general_tier_flag = r.read_bit()? as u8;
    let general_profile_idc = r.read_bits(5)? as u8;
    let _compat = r.read_bits(32)?; // 32 compatibility flags
    let _constraint_hi = r.read_bits(32)?; // 48 constraint/reserved/inbld bits …
    let _constraint_lo = r.read_bits(16)?; // … = 48 total
    let general_level_idc = r.read_bits(8)? as u8;

    // Sub-layer present flags.
    let mut sub_profile = [false; 8];
    let mut sub_level = [false; 8];
    for i in 0..max_sub_layers_minus1 as usize {
        sub_profile[i] = r.read_bit()? == 1;
        sub_level[i] = r.read_bit()? == 1;
    }
    if max_sub_layers_minus1 > 0 {
        for _ in max_sub_layers_minus1..8 {
            let _reserved = r.read_bits(2)?;
        }
    }
    for i in 0..max_sub_layers_minus1 as usize {
        if sub_profile[i] {
            // 88-bit sub-layer profile block.
            r.read_bits(32)?;
            r.read_bits(32)?;
            r.read_bits(24)?;
        }
        if sub_level[i] {
            r.read_bits(8)?;
        }
    }
    Some((general_profile_idc, general_tier_flag, general_level_idc))
}

/// Parse an HEVC 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_video_parameter_set_id = r.read_bits(4)?;
    let max_sub_layers_minus1 = r.read_bits(3)?;
    let _sps_temporal_id_nesting = r.read_bit()?;

    let (profile, tier, level) = read_profile_tier_level(&mut r, max_sub_layers_minus1)?;

    let _sps_seq_parameter_set_id = r.read_ue()?;
    let chroma_format_idc = r.read_ue()?;
    if chroma_format_idc == 3 {
        let _separate_colour_plane = 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,
    })
}

/// [`CodecParser`] implementation for H.265 / HEVC.
pub struct H265;

impl CodecParser for H265 {
    const CODEC: CodecId = CodecId::H265;

    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 <= 31 && 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 {
        // hvc1.{profile_space + profile_idc}.{compat}.{tier}{level}.{constraints}
        let tier = if p.tier == 1 { 'H' } else { 'L' };
        format!("hvc1.{}.6.{}{}.B0", p.profile, tier, p.level)
    }
}

/// Build an `hvcC` (HEVCDecoderConfigurationRecord, ISO/IEC 14496-15 §8.3.3.1)
/// from an Annex-B VPS/SPS/PPS config access unit, returning it alongside the
/// parsed [`VideoParams`]. Returns `None` if no SPS is present or its
/// profile_tier_level can't be located.
///
/// The 12-byte general `profile_tier_level` is copied verbatim from the SPS RBSP
/// (after un-escaping), so the profile/compatibility/constraint/level fields are
/// exact rather than re-derived. Used by the fMP4/CMAF muxer.
pub fn hvcc_config_record(data: &[u8]) -> Option<(VideoParams, Vec<u8>)> {
    let params = H265::parse_config(data)?;

    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),
            _ => {}
        }
    }
    let first_sps = *sps.first()?;

    // SPS RBSP layout: sps_vps_id(4) + max_sub_layers_minus1(3) + nesting(1) = 1
    // byte, then the 12-byte general profile_tier_level.
    let rbsp = unescape_rbsp(first_sps.get(2..)?);
    let ptl = rbsp.get(1..13)?; // profile(1) | compat(4) | constraint(6) | level(1)

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

    let mut rec = Vec::with_capacity(32 + first_sps.len());
    rec.push(1); // configurationVersion
    rec.extend_from_slice(ptl); // 12-byte general profile_tier_level
    rec.extend_from_slice(&[0xF0, 0x00]); // reserved | min_spatial_segmentation_idc=0
    rec.push(0xFC); // reserved(6) | parallelismType(2)=0
    rec.push(0xFC | (chroma_format_idc & 0x03)); // reserved(6) | chromaFormat(2)
    rec.push(0xF8 | (bit_depth_minus8 & 0x07)); // reserved(5) | bitDepthLumaMinus8(3)
    rec.push(0xF8 | (bit_depth_minus8 & 0x07)); // reserved(5) | bitDepthChromaMinus8(3)
    rec.extend_from_slice(&0u16.to_be_bytes()); // avgFrameRate

    // numTemporalLayers=1 | lengthSizeMinusOne=3 (constantFrameRate=0, temporalIdNested=0).
    rec.push((1 << 3) | 0x03);

    // NAL-unit arrays (VPS/SPS/PPS), each a complete NAL incl. its 2-byte 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); // numOfArrays
    for (t, nals) in present {
        rec.push(0x80 | (t & 0x3F)); // array_completeness=1 | reserved(1)=0 | type(6)
        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))
}

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

    /// Synthesize an HEVC SPS NAL for 1920x1080 (Main, level 4.0), 4:2:0 with a
    /// 4-row conformance crop (coded 1088 → displayed 1080).
    fn sps_1920x1080() -> Vec<u8> {
        let mut w = BitWriter::default();
        w.bits(0, 4); // sps_video_parameter_set_id
        w.bits(0, 3); // sps_max_sub_layers_minus1 = 0 (no sub-layer loops)
        w.bit(0); // sps_temporal_id_nesting_flag
                  // profile_tier_level (96 bits): profile_space(2) tier(1) idc(5) …
        w.bits(0, 2); // general_profile_space
        w.bit(0); // general_tier_flag = 0 (Main tier)
        w.bits(1, 5); // general_profile_idc = 1 (Main)
        w.bits(0, 32); // compatibility flags
        w.bits(0, 32); // constraint/reserved (hi)
        w.bits(0, 16); // constraint/reserved (lo) → 48 total
        w.bits(120, 8); // general_level_idc = 120 (level 4.0)
        w.ue(0); // sps_seq_parameter_set_id
        w.ue(1); // chroma_format_idc = 1 (4:2:0)
        w.ue(1920); // pic_width_in_luma_samples
        w.ue(1088); // pic_height_in_luma_samples
        w.bit(1); // conformance_window_flag
        w.ue(0); // conf_win_left_offset
        w.ue(0); // conf_win_right_offset
        w.ue(0); // conf_win_top_offset
        w.ue(2); // conf_win_bottom_offset → 1088 - 2*2 = 1084? see below
        let mut nal = vec![0x42u8, 0x01]; // NAL header: type 33 (SPS), tid+1=1
        nal.extend_from_slice(&w.bytes());
        nal
    }

    #[test]
    fn parse_sps_extracts_resolution() {
        let p = parse_sps(&sps_1920x1080()).expect("parse");
        // 4:2:0 crop unit Y = 2, bottom offset 2 → 1088 - 2*2 = 1084.
        assert_eq!((p.width, p.height), (1920, 1084));
        assert_eq!((p.profile, p.tier, p.level), (1, 0, 120));
    }

    #[test]
    fn hvcc_config_record_copies_ptl_and_lists_arrays() {
        let mut au = vec![0, 0, 0, 1, 0x40, 0x01, 0xAA]; // VPS (type 32)
        au.extend_from_slice(&[0, 0, 0, 1]);
        au.extend_from_slice(&sps_1920x1080()); // SPS (type 33)
        au.extend_from_slice(&[0, 0, 0, 1, 0x44, 0x01, 0xBB]); // PPS (type 34)

        let (params, rec) = hvcc_config_record(&au).expect("hvcC built");
        assert_eq!((params.width, params.profile, params.level), (1920, 1, 120));

        assert_eq!(rec[0], 1, "configurationVersion");
        assert_eq!(rec[1] & 0x1F, 1, "general_profile_idc copied from SPS");
        assert_eq!(rec[12], 120, "general_level_idc copied from SPS");
        // Fixed header is 22 bytes; numOfArrays follows. VPS+SPS+PPS = 3.
        assert_eq!(rec[22], 3, "three NAL arrays (VPS/SPS/PPS)");
    }

    #[test]
    fn hvcc_config_record_requires_an_sps() {
        // VPS only — no SPS — cannot build a config record.
        assert!(hvcc_config_record(&[0, 0, 0, 1, 0x40, 0x01, 0xAA]).is_none());
    }

    #[test]
    fn classifies_irap_and_config() {
        let mut au = vec![0, 0, 0, 1];
        au.extend_from_slice(&sps_1920x1080());
        // IDR_W_RADL VCL NAL (type 19): byte0 = 19<<1 = 0x26.
        au.extend_from_slice(&[0, 0, 0, 1, 0x26, 0x01, 0xAA]);

        assert!(H265::is_random_access_point(&au));
        assert!(H265::carries_config(&au));
        let p = H265::parse_config(&au).expect("params");
        assert_eq!(p.width, 1920);
        assert_eq!(H265::hls_codec_string(&p), "hvc1.1.6.L120.B0");

        // A TRAIL_R slice (type 1) alone is not a RAP.
        assert!(!H265::is_random_access_point(&[
            0, 0, 0, 1, 0x02, 0x01, 0xAA
        ]));
    }
}