use crate::codec::h264::bitstream::insert_emulation_prevention;
use crate::codec::h264::NalType;
pub trait BitSink {
fn write_bits(&mut self, value: u32, n: u8);
fn write_bit(&mut self, b: bool);
fn write_ue(&mut self, value: u32);
fn write_se(&mut self, value: i32);
fn bits_written(&self) -> usize;
}
#[derive(Debug, Default, Clone, Copy)]
pub struct BitSizer {
bits: usize,
}
impl BitSizer {
pub fn new() -> Self {
Self::default()
}
pub fn reset(&mut self) {
self.bits = 0;
}
}
impl BitSink for BitSizer {
#[inline]
fn write_bits(&mut self, _value: u32, n: u8) {
self.bits += n as usize;
}
#[inline]
fn write_bit(&mut self, _b: bool) {
self.bits += 1;
}
#[inline]
fn write_ue(&mut self, value: u32) {
let plus_one = value as u64 + 1;
let leading_zeros = (63u8 - plus_one.leading_zeros() as u8).min(31);
self.bits += 2 * leading_zeros as usize + 1;
}
#[inline]
fn write_se(&mut self, value: i32) {
let code_num = if value > 0 {
(2 * value - 1) as u32
} else {
(-2 * value) as u32
};
self.write_ue(code_num);
}
#[inline]
fn bits_written(&self) -> usize {
self.bits
}
}
#[derive(Debug, Default, Clone)]
pub struct BitWriter {
buffer: Vec<u8>,
pending: u32,
pending_bits: u8,
}
impl BitWriter {
pub fn new() -> Self {
Self::default()
}
pub fn with_capacity(cap: usize) -> Self {
Self {
buffer: Vec::with_capacity(cap),
..Self::default()
}
}
pub fn write_bits(&mut self, value: u32, n: u8) {
debug_assert!(n <= 32, "write_bits called with n={n} > 32");
let mut remaining = n;
while remaining > 0 {
let take = remaining.min(8 - self.pending_bits);
let shift = remaining - take;
let chunk = (value >> shift) & ((1u32 << take) - 1);
self.pending = (self.pending << take) | chunk;
self.pending_bits += take;
remaining -= take;
if self.pending_bits == 8 {
self.buffer.push(self.pending as u8);
self.pending = 0;
self.pending_bits = 0;
}
}
}
#[inline]
pub fn write_bit(&mut self, b: bool) {
self.write_bits(b as u32, 1);
}
pub fn write_ue(&mut self, value: u32) {
let plus_one = value as u64 + 1;
let leading_zeros = (63u8 - plus_one.leading_zeros() as u8).min(31);
for _ in 0..leading_zeros {
self.write_bit(false);
}
self.write_bit(true);
if leading_zeros > 0 {
let suffix = plus_one as u32 & ((1u32 << leading_zeros) - 1);
self.write_bits(suffix, leading_zeros);
}
}
pub fn write_se(&mut self, value: i32) {
let code_num = if value > 0 {
(2 * value - 1) as u32
} else {
(-2 * value) as u32
};
self.write_ue(code_num);
}
pub fn write_rbsp_trailing(&mut self) {
self.write_bit(true);
while self.pending_bits != 0 {
self.write_bit(false);
}
}
#[inline]
pub fn byte_aligned(&self) -> bool {
self.pending_bits == 0
}
#[inline]
pub fn bits_written(&self) -> usize {
self.buffer.len() * 8 + self.pending_bits as usize
}
}
impl BitSink for BitWriter {
#[inline]
fn write_bits(&mut self, value: u32, n: u8) {
BitWriter::write_bits(self, value, n);
}
#[inline]
fn write_bit(&mut self, b: bool) {
BitWriter::write_bit(self, b);
}
#[inline]
fn write_ue(&mut self, value: u32) {
BitWriter::write_ue(self, value);
}
#[inline]
fn write_se(&mut self, value: i32) {
BitWriter::write_se(self, value);
}
#[inline]
fn bits_written(&self) -> usize {
BitWriter::bits_written(self)
}
}
impl BitWriter {
pub fn finish(mut self) -> Vec<u8> {
if self.pending_bits > 0 {
let pad = 8 - self.pending_bits;
self.pending <<= pad;
self.buffer.push(self.pending as u8);
}
self.buffer
}
}
#[derive(Debug, Clone, Copy)]
pub struct SpsParams {
pub width_pixels: u32,
pub height_pixels: u32,
pub sps_id: u8,
pub max_num_ref_frames: u8,
pub pic_order_cnt_type: u8,
pub log2_max_pic_order_cnt_lsb_minus4: u8,
pub vui: Option<VuiParams>,
}
impl Default for SpsParams {
fn default() -> Self {
Self {
width_pixels: 1280,
height_pixels: 720,
sps_id: 0,
max_num_ref_frames: 1,
pic_order_cnt_type: 2,
log2_max_pic_order_cnt_lsb_minus4: 4,
vui: None,
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct VuiParams {
pub aspect_ratio_idc: u8,
pub colour_primaries: u8,
pub transfer_characteristics: u8,
pub matrix_coefficients: u8,
pub video_full_range: bool,
pub num_units_in_tick: u32,
pub time_scale: u32,
pub max_num_reorder_frames: u8,
pub max_dec_frame_buffering: u8,
}
impl VuiParams {
pub fn handbrake_x264(fps_num: u32, fps_den: u32) -> Self {
Self {
aspect_ratio_idc: 1, colour_primaries: 1, transfer_characteristics: 1, matrix_coefficients: 1, video_full_range: false, num_units_in_tick: fps_den,
time_scale: 2u32.saturating_mul(fps_num),
max_num_reorder_frames: 2, max_dec_frame_buffering: 3, }
}
}
fn derive_level_idc(width: u32, height: u32) -> u8 {
let mb = (width / 16) * (height / 16);
match mb {
0..=99 => 10, 100..=396 => 20, 397..=1620 => 30, 1621..=3600 => 31, 3601..=5120 => 32, 5121..=8192 => 40, 8193..=8704 => 42, 8705..=22080 => 50, 22081..=36864 => 51, _ => 52, }
}
pub fn build_sps_baseline(p: &SpsParams) -> Vec<u8> {
build_sps_inner(p, 66, 0b1000_0000)
}
pub fn build_sps_main(p: &SpsParams) -> Vec<u8> {
build_sps_inner(p, 77, 0b0000_0000)
}
pub fn build_sps_high(p: &SpsParams) -> Vec<u8> {
build_sps_inner_high(p)
}
fn build_sps_inner(p: &SpsParams, profile_idc: u32, constraint_byte: u32) -> Vec<u8> {
let mut w = BitWriter::with_capacity(16);
w.write_bits(profile_idc, 8);
w.write_bits(constraint_byte, 8);
w.write_bits(derive_level_idc(p.width_pixels, p.height_pixels) as u32, 8);
w.write_ue(p.sps_id as u32);
w.write_ue(0);
w.write_ue(p.pic_order_cnt_type as u32);
if p.pic_order_cnt_type == 0 {
w.write_ue(p.log2_max_pic_order_cnt_lsb_minus4 as u32);
}
w.write_ue(p.max_num_ref_frames as u32);
w.write_bit(false);
w.write_ue((p.width_pixels / 16) - 1);
w.write_ue((p.height_pixels / 16) - 1);
w.write_bit(true);
w.write_bit(true);
w.write_bit(false);
write_vui_section(&mut w, p);
w.write_rbsp_trailing();
w.finish()
}
fn build_sps_inner_high(p: &SpsParams) -> Vec<u8> {
let mut w = BitWriter::with_capacity(16);
w.write_bits(100, 8); w.write_bits(0b0000_0000, 8); w.write_bits(derive_level_idc(p.width_pixels, p.height_pixels) as u32, 8);
w.write_ue(p.sps_id as u32);
w.write_ue(1); w.write_ue(0); w.write_ue(0); w.write_bit(false); w.write_bit(false);
w.write_ue(0); w.write_ue(p.pic_order_cnt_type as u32);
if p.pic_order_cnt_type == 0 {
w.write_ue(p.log2_max_pic_order_cnt_lsb_minus4 as u32);
}
w.write_ue(p.max_num_ref_frames as u32);
w.write_bit(false); w.write_ue((p.width_pixels / 16) - 1);
w.write_ue((p.height_pixels / 16) - 1);
w.write_bit(true); w.write_bit(true); w.write_bit(false); write_vui_section(&mut w, p);
w.write_rbsp_trailing();
w.finish()
}
fn write_vui_section(w: &mut BitWriter, p: &SpsParams) {
if let Some(vui) = p.vui {
w.write_bit(true); write_vui_body(w, &vui);
} else {
w.write_bit(false);
}
}
fn write_vui_body(w: &mut BitWriter, v: &VuiParams) {
w.write_bit(true);
w.write_bits(v.aspect_ratio_idc as u32, 8);
if v.aspect_ratio_idc == 255 {
w.write_bits(0, 16); w.write_bits(0, 16); }
w.write_bit(false);
w.write_bit(true);
w.write_bits(5, 3); w.write_bit(v.video_full_range);
w.write_bit(true); w.write_bits(v.colour_primaries as u32, 8);
w.write_bits(v.transfer_characteristics as u32, 8);
w.write_bits(v.matrix_coefficients as u32, 8);
w.write_bit(false);
w.write_bit(true);
w.write_bits(v.num_units_in_tick, 32);
w.write_bits(v.time_scale, 32);
w.write_bit(true);
w.write_bit(false);
w.write_bit(false);
w.write_bit(false);
w.write_bit(true);
w.write_bit(true);
w.write_ue(0); w.write_ue(0); w.write_ue(16); w.write_ue(16); w.write_ue(v.max_num_reorder_frames as u32);
w.write_ue(v.max_dec_frame_buffering as u32);
}
#[derive(Debug, Clone, Copy)]
pub struct PpsParams {
pub pps_id: u8,
pub sps_id: u8,
pub pic_init_qp: u8,
pub deblocking_filter_control_present: bool,
}
impl Default for PpsParams {
fn default() -> Self {
Self {
pps_id: 0,
sps_id: 0,
pic_init_qp: 26,
deblocking_filter_control_present: false,
}
}
}
pub fn build_pps_cavlc(p: &PpsParams) -> Vec<u8> {
build_pps_inner(p, false, false)
}
pub fn build_pps_cabac(p: &PpsParams) -> Vec<u8> {
build_pps_inner(p, true, false)
}
pub fn build_pps_cabac_high(p: &PpsParams) -> Vec<u8> {
build_pps_inner(p, true, true)
}
fn build_pps_inner(p: &PpsParams, entropy_coding_mode: bool, high_profile_suffix: bool) -> Vec<u8> {
let mut w = BitWriter::with_capacity(8);
w.write_ue(p.pps_id as u32);
w.write_ue(p.sps_id as u32);
w.write_bit(entropy_coding_mode);
w.write_bit(false);
w.write_ue(0);
w.write_ue(0);
w.write_ue(0);
w.write_bit(false);
w.write_bits(0, 2);
w.write_se(p.pic_init_qp as i32 - 26);
w.write_se(0);
w.write_se(0);
w.write_bit(p.deblocking_filter_control_present);
w.write_bit(false);
w.write_bit(false);
if high_profile_suffix {
w.write_bit(true);
w.write_bit(false);
w.write_se(0);
}
w.write_rbsp_trailing();
w.finish()
}
#[derive(Debug, Clone, Copy)]
pub struct ISliceHeaderParams {
pub is_idr: bool,
pub pps_id: u8,
pub frame_num: u8,
pub idr_pic_id: u16,
pub slice_qp_delta: i32,
pub disable_deblocking: bool,
pub pic_order_cnt_lsb: Option<u32>,
pub log2_max_pic_order_cnt_lsb_minus4: u8,
}
impl Default for ISliceHeaderParams {
fn default() -> Self {
Self {
is_idr: true,
pps_id: 0,
frame_num: 0,
idr_pic_id: 0,
slice_qp_delta: 0,
disable_deblocking: false,
pic_order_cnt_lsb: None,
log2_max_pic_order_cnt_lsb_minus4: 4,
}
}
}
pub fn build_slice_header_i(p: &ISliceHeaderParams) -> Vec<u8> {
let mut w = BitWriter::with_capacity(8);
w.write_ue(0);
w.write_ue(7);
w.write_ue(p.pps_id as u32);
w.write_bits(p.frame_num as u32 & 0xF, 4);
if p.is_idr {
w.write_ue(p.idr_pic_id as u32);
}
if let Some(lsb) = p.pic_order_cnt_lsb {
let bits: u8 = p.log2_max_pic_order_cnt_lsb_minus4 + 4;
let mask = (1u32 << bits as u32) - 1;
w.write_bits(lsb & mask, bits);
}
if p.is_idr {
w.write_bit(false); w.write_bit(false); } else {
w.write_bit(false);
}
w.write_se(p.slice_qp_delta);
if p.disable_deblocking {
w.write_ue(1);
} else {
w.write_ue(0);
w.write_se(0); w.write_se(0); }
w.finish()
}
pub fn continue_slice_header_i(w: &mut BitWriter, p: &ISliceHeaderParams) {
w.write_ue(0);
w.write_ue(7);
w.write_ue(p.pps_id as u32);
w.write_bits(p.frame_num as u32 & 0xF, 4);
if p.is_idr {
w.write_ue(p.idr_pic_id as u32);
}
if let Some(lsb) = p.pic_order_cnt_lsb {
let bits: u8 = p.log2_max_pic_order_cnt_lsb_minus4 + 4;
let mask = (1u32 << bits as u32) - 1;
w.write_bits(lsb & mask, bits);
}
if p.is_idr {
w.write_bit(false);
w.write_bit(false);
} else {
w.write_bit(false);
}
w.write_se(p.slice_qp_delta);
if p.disable_deblocking {
w.write_ue(1);
} else {
w.write_ue(0);
w.write_se(0);
w.write_se(0);
}
}
#[derive(Debug, Clone, Copy)]
#[derive(Default)]
pub struct PSliceHeaderParams {
pub pps_id: u8,
pub frame_num: u8,
pub pic_order_cnt_lsb: Option<u32>,
pub log2_max_pic_order_cnt_lsb_minus4: u8,
pub slice_qp_delta: i32,
pub disable_deblocking: bool,
pub num_ref_idx_active_override: bool,
pub num_ref_idx_l0_active_minus1: u8,
pub cabac_init_idc: Option<u8>,
}
pub fn continue_slice_header_p(w: &mut BitWriter, p: &PSliceHeaderParams) {
w.write_ue(0);
w.write_ue(5);
w.write_ue(p.pps_id as u32);
w.write_bits(p.frame_num as u32 & 0xF, 4);
if let Some(lsb) = p.pic_order_cnt_lsb {
let bits: u8 = p.log2_max_pic_order_cnt_lsb_minus4 + 4;
let mask = (1u32 << bits as u32) - 1;
w.write_bits(lsb & mask, bits);
}
w.write_bit(p.num_ref_idx_active_override);
if p.num_ref_idx_active_override {
w.write_ue(p.num_ref_idx_l0_active_minus1 as u32);
}
w.write_bit(false);
w.write_bit(false);
if let Some(idc) = p.cabac_init_idc {
debug_assert!(idc <= 2);
w.write_ue(idc as u32);
}
w.write_se(p.slice_qp_delta);
if p.disable_deblocking {
w.write_ue(1);
} else {
w.write_ue(0);
w.write_se(0);
w.write_se(0);
}
}
#[derive(Debug, Clone, Copy)]
#[derive(Default)]
pub struct BSliceHeaderParams {
pub pps_id: u8,
pub frame_num: u8,
pub pic_order_cnt_lsb: u32,
pub direct_spatial_mv_pred_flag: bool,
pub slice_qp_delta: i32,
pub disable_deblocking: bool,
pub num_ref_idx_active_override: bool,
pub num_ref_idx_l0_active_minus1: u8,
pub num_ref_idx_l1_active_minus1: u8,
pub cabac_init_idc: Option<u8>,
pub log2_max_pic_order_cnt_lsb_minus4: u8,
pub log2_max_frame_num_minus4: u8,
}
pub fn continue_slice_header_b(w: &mut BitWriter, p: &BSliceHeaderParams) {
let frame_num_bits: u8 = p.log2_max_frame_num_minus4 + 4;
let poc_lsb_bits: u8 = p.log2_max_pic_order_cnt_lsb_minus4 + 4;
w.write_ue(0);
w.write_ue(6);
w.write_ue(p.pps_id as u32);
let frame_num_mask = (1u32 << frame_num_bits as u32) - 1;
w.write_bits(p.frame_num as u32 & frame_num_mask, frame_num_bits);
let poc_mask = (1u32 << poc_lsb_bits as u32) - 1;
w.write_bits(p.pic_order_cnt_lsb & poc_mask, poc_lsb_bits);
w.write_bit(p.direct_spatial_mv_pred_flag);
w.write_bit(p.num_ref_idx_active_override);
if p.num_ref_idx_active_override {
w.write_ue(p.num_ref_idx_l0_active_minus1 as u32);
w.write_ue(p.num_ref_idx_l1_active_minus1 as u32);
}
w.write_bit(false); w.write_bit(false); if let Some(idc) = p.cabac_init_idc {
debug_assert!(idc <= 2);
w.write_ue(idc as u32);
}
w.write_se(p.slice_qp_delta);
if p.disable_deblocking {
w.write_ue(1);
} else {
w.write_ue(0);
w.write_se(0);
w.write_se(0);
}
}
#[derive(Debug, Clone, Copy)]
pub enum PrimaryPicType {
IOnly = 0,
IP = 1,
IPB = 2,
}
pub fn build_aud_rbsp(pic_type: PrimaryPicType) -> Vec<u8> {
let mut w = BitWriter::with_capacity(1);
w.write_bits(pic_type as u32, 3);
w.write_rbsp_trailing();
w.finish()
}
fn nal_header_byte(nal_type: NalType, nal_ref_idc: u8) -> u8 {
debug_assert!(nal_ref_idc <= 3);
((nal_ref_idc & 0x3) << 5) | (nal_type.0 & 0x1F)
}
pub fn wrap_rbsp_as_nal(rbsp: &[u8], nal_type: NalType, nal_ref_idc: u8) -> Vec<u8> {
let protected = insert_emulation_prevention(rbsp);
let mut nal = Vec::with_capacity(protected.len() + 1);
nal.push(nal_header_byte(nal_type, nal_ref_idc));
nal.extend_from_slice(&protected);
nal
}
#[cfg(test)]
mod tests {
use super::*;
use crate::codec::h264::bitstream::{parse_nal_unit, RbspReader};
use crate::codec::h264::sps::{parse_pps, parse_sps};
#[test]
fn bitwriter_bytes_and_bits() {
let mut w = BitWriter::new();
w.write_bits(0b1010, 4);
w.write_bits(0b0110, 4);
assert_eq!(w.finish(), vec![0xA6]);
}
#[test]
fn bitwriter_cross_byte_write() {
let mut w = BitWriter::new();
w.write_bits(0b101, 3);
w.write_bits(0b11_0000_1100, 10);
assert_eq!(w.finish(), vec![0xB8, 0x60]);
}
#[test]
fn bitwriter_ue_known_codewords() {
for (value, expected_bits) in &[
(0u32, "1"),
(1, "010"),
(2, "011"),
(3, "00100"),
(6, "00111"),
] {
let mut w = BitWriter::new();
w.write_ue(*value);
let mut padded = w;
padded.write_rbsp_trailing();
let bytes = padded.finish();
let mut r = RbspReader::new(&bytes);
assert_eq!(
r.read_ue().unwrap(),
*value,
"ue({value}) round-trip — expected encoding {expected_bits}"
);
}
}
#[test]
fn bitwriter_se_round_trip() {
for v in &[-10i32, -5, -1, 0, 1, 5, 10] {
let mut w = BitWriter::new();
w.write_se(*v);
w.write_rbsp_trailing();
let bytes = w.finish();
let mut r = RbspReader::new(&bytes);
assert_eq!(r.read_se().unwrap(), *v, "se({v}) round-trip");
}
}
#[test]
fn sps_baseline_round_trip() {
let p = SpsParams {
width_pixels: 320,
height_pixels: 240,
sps_id: 0,
max_num_ref_frames: 1,
..SpsParams::default()
};
let rbsp = build_sps_baseline(&p);
let sps = parse_sps(&rbsp).expect("parse_sps should accept our SPS");
assert_eq!(sps.profile_idc, 66);
assert_eq!(sps.level_idc, 20); assert_eq!(sps.sps_id, 0);
assert_eq!(sps.pic_width_in_mbs, 20);
assert_eq!(sps.pic_height_in_map_units, 15);
assert_eq!(sps.width_in_pixels, 320);
assert_eq!(sps.height_in_pixels, 240);
assert!(sps.frame_mbs_only_flag);
}
#[test]
fn sps_1080p_level_40() {
let p = SpsParams {
width_pixels: 1920,
height_pixels: 1088, ..SpsParams::default()
};
let rbsp = build_sps_baseline(&p);
let sps = parse_sps(&rbsp).unwrap();
assert_eq!(sps.level_idc, 40);
assert_eq!(sps.width_in_pixels, 1920);
assert_eq!(sps.height_in_pixels, 1088);
}
#[test]
fn pps_cavlc_round_trip() {
let p = PpsParams::default();
let rbsp = build_pps_cavlc(&p);
let pps = parse_pps(&rbsp).expect("parse_pps should accept our PPS");
assert_eq!(pps.pps_id, 0);
assert_eq!(pps.sps_id, 0);
assert!(!pps.entropy_coding_mode_flag);
assert_eq!(pps.pic_init_qp_minus26, 0);
assert!(!pps.deblocking_filter_control_present_flag);
}
#[test]
fn pps_with_deblocking_control_round_trip() {
let p = PpsParams {
deblocking_filter_control_present: true,
..PpsParams::default()
};
let rbsp = build_pps_cavlc(&p);
let pps = parse_pps(&rbsp).unwrap();
assert!(pps.deblocking_filter_control_present_flag);
}
#[test]
fn sps_main_profile_has_profile_idc_77() {
let p = SpsParams {
width_pixels: 320,
height_pixels: 240,
sps_id: 0,
max_num_ref_frames: 1,
..SpsParams::default()
};
let rbsp = build_sps_main(&p);
let sps = parse_sps(&rbsp).expect("parse_sps should accept our Main SPS");
assert_eq!(sps.profile_idc, 77);
assert_eq!(sps.level_idc, 20);
assert_eq!(sps.width_in_pixels, 320);
assert_eq!(sps.height_in_pixels, 240);
}
#[test]
fn sps_main_and_baseline_payloads_differ_only_in_profile_byte() {
let p = SpsParams::default();
let rb_base = build_sps_baseline(&p);
let rb_main = build_sps_main(&p);
assert_eq!(rb_base.len(), rb_main.len());
assert_eq!(rb_base[0], 66);
assert_eq!(rb_main[0], 77);
assert_eq!(&rb_base[2..], &rb_main[2..]);
}
#[test]
fn pps_cabac_sets_entropy_coding_mode_flag() {
let p = PpsParams::default();
let rbsp = build_pps_cabac(&p);
let pps = parse_pps(&rbsp).expect("parse_pps should accept CABAC PPS");
assert!(pps.entropy_coding_mode_flag);
assert_eq!(pps.pic_init_qp_minus26, 0);
}
#[test]
fn pps_cabac_and_cavlc_differ_only_at_entropy_coding_bit() {
let p = PpsParams::default();
let rbsp_cavlc = build_pps_cavlc(&p);
let rbsp_cabac = build_pps_cabac(&p);
assert_eq!(rbsp_cavlc.len(), rbsp_cabac.len());
let pps_cavlc = parse_pps(&rbsp_cavlc).unwrap();
let pps_cabac = parse_pps(&rbsp_cabac).unwrap();
assert!(!pps_cavlc.entropy_coding_mode_flag);
assert!(pps_cabac.entropy_coding_mode_flag);
assert_eq!(pps_cavlc.pps_id, pps_cabac.pps_id);
assert_eq!(
pps_cavlc.deblocking_filter_control_present_flag,
pps_cabac.deblocking_filter_control_present_flag,
);
}
#[test]
fn p_slice_header_emits_cabac_init_idc_when_set() {
use crate::codec::h264::encoder::bitstream_writer::continue_slice_header_p;
let mut w_cavlc = BitWriter::new();
let p_cavlc = PSliceHeaderParams::default();
continue_slice_header_p(&mut w_cavlc, &p_cavlc);
let len_cavlc_bits = w_cavlc.bits_written();
let mut w_cabac = BitWriter::new();
let p_cabac = PSliceHeaderParams {
cabac_init_idc: Some(0),
..PSliceHeaderParams::default()
};
continue_slice_header_p(&mut w_cabac, &p_cabac);
let len_cabac_bits = w_cabac.bits_written();
assert_eq!(
len_cabac_bits,
len_cavlc_bits + 1,
"ue(0) = '1' adds exactly 1 bit"
);
}
#[test]
fn p_slice_header_cabac_init_idc_2_adds_three_bits() {
use crate::codec::h264::encoder::bitstream_writer::continue_slice_header_p;
let mut w_cavlc = BitWriter::new();
continue_slice_header_p(&mut w_cavlc, &PSliceHeaderParams::default());
let len_cavlc_bits = w_cavlc.bits_written();
let mut w_cabac = BitWriter::new();
continue_slice_header_p(
&mut w_cabac,
&PSliceHeaderParams {
cabac_init_idc: Some(2),
..PSliceHeaderParams::default()
},
);
let len_cabac_bits = w_cabac.bits_written();
assert_eq!(len_cabac_bits, len_cavlc_bits + 3);
}
#[test]
fn nal_wrap_sps_parses_back() {
let p = SpsParams::default();
let rbsp = build_sps_baseline(&p);
let nal = wrap_rbsp_as_nal(&rbsp, NalType::SPS, 3);
let parsed = parse_nal_unit(&nal).unwrap();
assert_eq!(parsed.nal_type, NalType::SPS);
assert_eq!(parsed.nal_ref_idc, 3);
assert_eq!(parsed.rbsp, rbsp);
}
#[test]
fn nal_wrap_emulation_prevention_inserts_03() {
let rbsp = vec![0x00, 0x00, 0x01, 0xFF];
let nal = wrap_rbsp_as_nal(&rbsp, NalType::SLICE_IDR, 3);
assert_eq!(nal, vec![0x65, 0x00, 0x00, 0x03, 0x01, 0xFF]);
}
#[test]
fn nal_header_byte_layout() {
assert_eq!(nal_header_byte(NalType::SPS, 3), 0x67);
assert_eq!(nal_header_byte(NalType::PPS, 3), 0x68);
assert_eq!(nal_header_byte(NalType::SLICE_IDR, 3), 0x65);
assert_eq!(nal_header_byte(NalType::SLICE, 2), 0x41);
}
#[test]
fn level_idc_bands() {
assert_eq!(derive_level_idc(176, 144), 10);
assert_eq!(derive_level_idc(352, 288), 20);
assert_eq!(derive_level_idc(720, 480), 30);
assert_eq!(derive_level_idc(1280, 720), 31);
assert_eq!(derive_level_idc(1920, 1088), 40);
assert_eq!(derive_level_idc(3840, 2176), 51);
}
#[test]
fn sps_high_profile_idc_is_100() {
let p = SpsParams {
width_pixels: 320,
height_pixels: 240,
sps_id: 0,
max_num_ref_frames: 1,
..SpsParams::default()
};
let rbsp = build_sps_high(&p);
let sps = parse_sps(&rbsp).expect("parse_sps should accept our High SPS");
assert_eq!(sps.profile_idc, 100);
assert_eq!(sps.level_idc, 20);
assert_eq!(sps.width_in_pixels, 320);
assert_eq!(sps.height_in_pixels, 240);
}
#[test]
fn sps_high_with_vui_parses_back() {
let p = SpsParams {
width_pixels: 1920,
height_pixels: 1088,
sps_id: 0,
max_num_ref_frames: 3,
pic_order_cnt_type: 0,
log2_max_pic_order_cnt_lsb_minus4: 4,
vui: Some(VuiParams::handbrake_x264(30, 1)),
};
let rbsp = build_sps_high(&p);
let sps = parse_sps(&rbsp).expect("parse_sps accepts VUI-bearing SPS");
assert_eq!(sps.profile_idc, 100);
assert_eq!(sps.width_in_pixels, 1920);
assert_eq!(sps.height_in_pixels, 1088);
assert_eq!(sps.max_num_ref_frames, 3);
}
#[test]
fn sps_with_vui_is_longer_than_sps_without() {
let p_novui = SpsParams::default();
let p_vui = SpsParams {
vui: Some(VuiParams::handbrake_x264(30, 1)),
..p_novui
};
let r_novui = build_sps_high(&p_novui);
let r_vui = build_sps_high(&p_vui);
assert!(
r_vui.len() > r_novui.len(),
"VUI-bearing SPS ({} bytes) should be longer than no-VUI SPS ({} bytes)",
r_vui.len(),
r_novui.len(),
);
}
#[test]
fn vui_handbrake_x264_30fps_matches_libx264_convention() {
let v = VuiParams::handbrake_x264(30, 1);
assert_eq!(v.num_units_in_tick, 1);
assert_eq!(v.time_scale, 60);
assert_eq!(v.aspect_ratio_idc, 1, "1:1 SAR");
assert_eq!(v.colour_primaries, 1, "BT.709 primaries");
assert_eq!(v.transfer_characteristics, 1, "BT.709 transfer");
assert_eq!(v.matrix_coefficients, 1, "BT.709 matrix");
assert!(!v.video_full_range, "limited range default");
}
#[test]
fn vui_handbrake_x264_29_97fps_matches_libx264_convention() {
let v = VuiParams::handbrake_x264(30000, 1001);
assert_eq!(v.num_units_in_tick, 1001);
assert_eq!(v.time_scale, 60000);
}
#[test]
fn sps_high_emits_chroma_format_420() {
let p = SpsParams::default();
let rbsp = build_sps_high(&p);
let sps = parse_sps(&rbsp).unwrap();
assert_eq!(sps.chroma_format_idc, 1);
assert_eq!(sps.bit_depth_luma, 8);
assert_eq!(sps.bit_depth_chroma, 8);
}
#[test]
fn sps_high_is_longer_than_main_due_to_suffix() {
let p = SpsParams::default();
let rb_main = build_sps_main(&p);
let rb_high = build_sps_high(&p);
assert!(
rb_high.len() >= rb_main.len(),
"High SPS ({}) shouldn't be shorter than Main ({})",
rb_high.len(),
rb_main.len(),
);
assert_eq!(rb_high[0], 100);
assert_eq!(rb_main[0], 77);
}
#[test]
fn pps_cabac_high_sets_transform_8x8_mode_flag() {
let p = PpsParams::default();
let rbsp = build_pps_cabac_high(&p);
let pps = parse_pps(&rbsp).expect("parse_pps should accept High CABAC PPS");
assert!(pps.entropy_coding_mode_flag);
assert!(pps.transform_8x8_mode_flag);
}
#[test]
fn pps_cabac_high_matches_cabac_except_suffix() {
let p = PpsParams::default();
let rb_main = build_pps_cabac(&p);
let rb_high = build_pps_cabac_high(&p);
assert!(
rb_high.len() >= rb_main.len(),
"High PPS ({}) shouldn't be shorter than CABAC Main ({})",
rb_high.len(),
rb_main.len(),
);
}
#[test]
fn b_slice_header_writer_emits_bytes() {
let mut w = BitWriter::new();
let p = BSliceHeaderParams {
pps_id: 0,
frame_num: 3,
pic_order_cnt_lsb: 6,
direct_spatial_mv_pred_flag: true,
slice_qp_delta: 2,
disable_deblocking: false,
num_ref_idx_active_override: false,
num_ref_idx_l0_active_minus1: 0,
num_ref_idx_l1_active_minus1: 0,
cabac_init_idc: Some(0),
log2_max_pic_order_cnt_lsb_minus4: 4,
log2_max_frame_num_minus4: 0,
};
continue_slice_header_b(&mut w, &p);
let bytes = w.finish();
assert!(!bytes.is_empty(), "B-slice header produced no bytes");
assert_eq!(
bytes[0] >> 4, 0x9,
"first nibble should encode first_mb_in_slice=0 + slice_type=6 prefix; got {:08b}",
bytes[0],
);
}
#[test]
fn b_slice_header_writer_override_extends_output() {
let base_params = BSliceHeaderParams {
pps_id: 0,
frame_num: 0,
pic_order_cnt_lsb: 0,
direct_spatial_mv_pred_flag: true,
slice_qp_delta: 0,
disable_deblocking: false,
num_ref_idx_active_override: false,
num_ref_idx_l0_active_minus1: 0,
num_ref_idx_l1_active_minus1: 0,
cabac_init_idc: Some(0),
log2_max_pic_order_cnt_lsb_minus4: 4,
log2_max_frame_num_minus4: 0,
};
let mut w_no_override = BitWriter::new();
continue_slice_header_b(&mut w_no_override, &base_params);
let no_override = w_no_override.finish();
let mut w_override = BitWriter::new();
let mut p = base_params;
p.num_ref_idx_active_override = true;
p.num_ref_idx_l0_active_minus1 = 1;
p.num_ref_idx_l1_active_minus1 = 1;
continue_slice_header_b(&mut w_override, &p);
let with_override = w_override.finish();
assert!(
with_override.len() >= no_override.len(),
"override path should not produce shorter output"
);
}
}