#[allow(unused_imports)]
use super::encoder::bitstream_writer::{BitSink, BitWriter}; use super::tables::{encode_coeff_token, encode_run_before, encode_total_zeros};
use super::H264Error;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CavlcBlockType {
Luma4x4,
Intra16x16Ac,
ChromaAc,
ChromaDc,
}
impl CavlcBlockType {
pub fn max_coeffs(self) -> u8 {
match self {
Self::Luma4x4 => 16,
Self::Intra16x16Ac | Self::ChromaAc => 15,
Self::ChromaDc => 4,
}
}
}
pub fn encode_cavlc_block<S: BitSink>(
writer: &mut S,
coeffs: &[i32],
nc: i8,
block_type: CavlcBlockType,
) -> Result<(), H264Error> {
let max_coeffs = block_type.max_coeffs();
if coeffs.len() != max_coeffs as usize {
return Err(H264Error::CavlcError(format!(
"encode_cavlc_block: expected {max_coeffs} coeffs, got {}",
coeffs.len()
)));
}
let (total_coeffs, trailing_ones, levels, t1_signs, runs, total_zeros) =
derive_cavlc_fields(coeffs)?;
let (bits, len) = encode_coeff_token(total_coeffs, trailing_ones, nc).ok_or_else(|| {
H264Error::CavlcError(format!(
"encode_coeff_token missing entry: tc={total_coeffs} t1={trailing_ones} nc={nc}"
))
})?;
writer.write_bits(bits as u32, len);
if total_coeffs == 0 {
return Ok(());
}
for k in 0..trailing_ones {
let sign_bit = (t1_signs >> k) & 1;
writer.write_bit(sign_bit != 0);
}
let mut suffix_length: u8 = if total_coeffs > 10 && trailing_ones < 3 {
1
} else {
0
};
let thresholds: [u32; 6] = [3, 6, 12, 24, 48, u32::MAX];
for (k, &level) in levels.iter().enumerate() {
let mut mag = level.unsigned_abs() as i64;
let neg = level < 0;
if k == 0 && trailing_ones < 3 {
mag -= 1; }
debug_assert!(mag >= 1, "derive_cavlc_fields should have ensured |level|>=1");
let level_code = 2 * mag - if neg { 1 } else { 2 };
debug_assert!(level_code >= 0, "level_code underflow at k={k} level={level}");
emit_level_code(writer, level_code as u64, suffix_length)?;
let abs_level = level.unsigned_abs();
if suffix_length == 0 {
suffix_length = 1;
}
if suffix_length < 6 && abs_level > thresholds[suffix_length as usize - 1] {
suffix_length += 1;
}
}
if total_coeffs < max_coeffs {
let (tz_bits, tz_len) = encode_total_zeros(total_zeros, total_coeffs, max_coeffs)
.ok_or_else(|| {
H264Error::CavlcError(format!(
"encode_total_zeros missing: tz={total_zeros} tc={total_coeffs} max={max_coeffs}"
))
})?;
writer.write_bits(tz_bits as u32, tz_len);
}
let mut zeros_left = total_zeros;
for i in 0..(total_coeffs as usize - 1) {
if zeros_left == 0 {
break;
}
let run = runs[i];
let (rb_bits, rb_len) = encode_run_before(run, zeros_left).ok_or_else(|| {
H264Error::CavlcError(format!(
"encode_run_before missing: run={run} zeros_left={zeros_left}"
))
})?;
writer.write_bits(rb_bits as u32, rb_len);
zeros_left = zeros_left.saturating_sub(run);
}
Ok(())
}
fn derive_cavlc_fields(
coeffs: &[i32],
) -> Result<(u8, u8, Vec<i32>, u8, Vec<u8>, u8), H264Error> {
let positions: Vec<usize> = coeffs
.iter()
.enumerate()
.filter(|(_, c)| **c != 0)
.map(|(i, _)| i)
.collect();
let total_coeffs = positions.len() as u8;
if total_coeffs == 0 {
return Ok((0, 0, Vec::new(), 0, Vec::new(), 0));
}
let levels_rev: Vec<i32> = positions.iter().rev().map(|&p| coeffs[p]).collect();
let mut trailing_ones = 0u8;
let mut t1_signs = 0u8;
for (k, &v) in levels_rev.iter().enumerate() {
if trailing_ones >= 3 || v.abs() != 1 {
break;
}
if v < 0 {
t1_signs |= 1 << k;
}
trailing_ones += 1;
}
let levels: Vec<i32> = levels_rev
.iter()
.skip(trailing_ones as usize)
.copied()
.collect();
let mut runs = Vec::with_capacity(total_coeffs as usize);
let tc = total_coeffs as usize;
for i in 0..tc {
let run_val = if i == tc - 1 {
positions[0]
} else {
let p_hi = positions[tc - 1 - i];
let p_lo_next = positions[tc - 2 - i];
p_hi - p_lo_next - 1
};
if run_val > 255 {
return Err(H264Error::CavlcError(format!(
"run_before overflow: {run_val}"
)));
}
runs.push(run_val as u8);
}
let total_zeros: u16 = runs.iter().map(|&r| r as u16).sum();
if total_zeros > 255 {
return Err(H264Error::CavlcError(format!(
"total_zeros overflow: {total_zeros}"
)));
}
Ok((
total_coeffs,
trailing_ones,
levels,
t1_signs,
runs,
total_zeros as u8,
))
}
fn emit_level_code<S: BitSink>(
writer: &mut S,
level_code: u64,
suffix_length: u8,
) -> Result<(), H264Error> {
let lc = level_code;
if suffix_length == 0 {
if lc < 14 {
emit_unary_prefix(writer, lc as u8);
return Ok(());
}
if lc < 14 + 16 {
emit_unary_prefix(writer, 14);
writer.write_bits((lc - 14) as u32, 4);
return Ok(());
}
let escape = lc - 30;
if escape > 0xFFF {
return Err(H264Error::CavlcError(format!(
"level_code too large for 12-bit escape: lc={lc}"
)));
}
emit_unary_prefix(writer, 15);
writer.write_bits(escape as u32, 12);
return Ok(());
}
let sl = suffix_length as u32;
let boundary = 15u64 << sl; if lc < boundary {
let prefix = (lc >> sl) as u8;
let suffix = (lc & ((1u64 << sl) - 1)) as u32;
emit_unary_prefix(writer, prefix);
writer.write_bits(suffix, suffix_length);
return Ok(());
}
let escape = lc - boundary;
if escape > 0xFFF {
return Err(H264Error::CavlcError(format!(
"level_code too large for 12-bit escape: lc={lc} sl={sl}"
)));
}
emit_unary_prefix(writer, 15);
writer.write_bits(escape as u32, 12);
Ok(())
}
fn emit_unary_prefix<S: BitSink>(writer: &mut S, n: u8) {
for _ in 0..n {
writer.write_bit(false);
}
writer.write_bit(true);
}
#[cfg(test)]
mod tests {
use super::*;
use crate::codec::h264::bitstream::{EpByteMap, RbspReader};
use crate::codec::h264::cavlc::decode_cavlc_block;
fn identity_ep_map(len: usize) -> EpByteMap {
EpByteMap {
rbsp_to_raw: (0..len).collect(),
}
}
fn encode_to_bytes(coeffs: &[i32], nc: i8, bt: CavlcBlockType) -> Vec<u8> {
let mut w = BitWriter::new();
encode_cavlc_block(&mut w, coeffs, nc, bt).unwrap();
while !w.byte_aligned() {
w.write_bit(false);
}
w.finish()
}
fn check_round_trip(coeffs: &[i32], nc: i8, bt: CavlcBlockType) {
let bytes = encode_to_bytes(coeffs, nc, bt);
let mut padded = bytes;
padded.extend_from_slice(&[0u8; 4]);
let ep_map = identity_ep_map(padded.len());
let mut reader = RbspReader::new(&padded);
let (block, _positions) =
decode_cavlc_block(&mut reader, nc, &ep_map, &padded, bt.max_coeffs())
.expect("decode should succeed on valid CAVLC");
assert_eq!(
block.total_coeffs as usize,
coeffs.iter().filter(|c| **c != 0).count(),
"total_coeffs mismatch for input {coeffs:?}"
);
for i in 0..bt.max_coeffs() as usize {
assert_eq!(
block.coeffs[i], coeffs[i],
"scan pos {i} mismatch for input {coeffs:?} → decoded {:?}",
&block.coeffs[..bt.max_coeffs() as usize]
);
}
}
#[test]
fn dc_scan_img4138_mb_0_0_round_trip() {
let coeffs = [
79i32, 14, 34, 2, -1, 3, 8, 0, 13, -19, -4, 0, 5, -10, -5, -48,
];
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn dc_scan_deterministic_mb_0_0_round_trip() {
let coeffs = [
-167i32, -24, -49, 0, 0, 0, -15, 0, 0, -29, 0, 0, 0, 0, 0, 0,
];
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn nc_table_selection_round_trip_all_tables() {
let coeffs = [3i32, -2, 5, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
for nc in [0i8, 1, 2, 3, 4, 5, 6, 7, 8, 12, 16] {
check_round_trip(&coeffs, nc, CavlcBlockType::Luma4x4);
}
}
#[test]
fn empty_block_emits_single_bit() {
let coeffs = [0i32; 16];
let bytes = encode_to_bytes(&coeffs, 0, CavlcBlockType::Luma4x4);
assert_eq!(bytes[0] & 0b1000_0000, 0b1000_0000);
}
#[test]
fn empty_luma_round_trip() {
check_round_trip(&[0i32; 16], 0, CavlcBlockType::Luma4x4);
}
#[test]
fn single_plus_one_at_last_scan_pos() {
let mut coeffs = [0i32; 16];
coeffs[15] = 1;
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn single_minus_one_at_last_scan_pos() {
let mut coeffs = [0i32; 16];
coeffs[15] = -1;
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn two_trailing_ones_round_trip() {
let mut coeffs = [0i32; 16];
coeffs[14] = -1;
coeffs[15] = 1;
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn three_trailing_ones_round_trip() {
let mut coeffs = [0i32; 16];
coeffs[13] = 1;
coeffs[14] = -1;
coeffs[15] = 1;
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn small_level_round_trip() {
let mut coeffs = [0i32; 16];
coeffs[10] = 3;
coeffs[15] = 1;
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn negative_level_round_trip() {
let mut coeffs = [0i32; 16];
coeffs[8] = -5;
coeffs[15] = 1;
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn large_level_round_trip() {
let mut coeffs = [0i32; 16];
coeffs[0] = 100;
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn extreme_level_escape_path() {
let mut coeffs = [0i32; 16];
coeffs[0] = 1000;
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn block_with_gaps_round_trip() {
let mut coeffs = [0i32; 16];
coeffs[0] = 5;
coeffs[4] = -3;
coeffs[8] = 2;
coeffs[15] = 1;
check_round_trip(&coeffs, 0, CavlcBlockType::Luma4x4);
}
#[test]
fn nc_2_round_trip() {
let mut coeffs = [0i32; 16];
coeffs[10] = 2;
coeffs[15] = -1;
check_round_trip(&coeffs, 2, CavlcBlockType::Luma4x4);
}
#[test]
fn nc_4_round_trip() {
let mut coeffs = [0i32; 16];
coeffs[5] = 4;
coeffs[12] = -2;
coeffs[15] = 1;
check_round_trip(&coeffs, 4, CavlcBlockType::Luma4x4);
}
#[test]
fn nc_8_round_trip() {
let mut coeffs = [0i32; 16];
coeffs[0] = -7;
coeffs[14] = 1;
coeffs[15] = -1;
check_round_trip(&coeffs, 8, CavlcBlockType::Luma4x4);
}
#[test]
fn intra16x16_ac_round_trip() {
let mut coeffs = [0i32; 15];
coeffs[3] = 2;
coeffs[14] = 1;
check_round_trip(&coeffs, 0, CavlcBlockType::Intra16x16Ac);
}
#[test]
fn chroma_ac_round_trip() {
let mut coeffs = [0i32; 15];
coeffs[0] = 3;
coeffs[14] = -1;
check_round_trip(&coeffs, 1, CavlcBlockType::ChromaAc);
}
#[test]
fn chroma_dc_round_trip() {
let mut coeffs = [0i32; 4];
coeffs[0] = 2;
coeffs[3] = -1;
check_round_trip(&coeffs, -1, CavlcBlockType::ChromaDc);
}
#[test]
fn chroma_dc_all_zero_round_trip() {
check_round_trip(&[0i32; 4], -1, CavlcBlockType::ChromaDc);
}
#[test]
fn derive_empty_block() {
let (tc, t1, _lv, sg, _runs, tz) = derive_cavlc_fields(&[0i32; 16]).unwrap();
assert_eq!(tc, 0);
assert_eq!(t1, 0);
assert_eq!(sg, 0);
assert_eq!(tz, 0);
}
#[test]
fn derive_single_plus_one_at_top() {
let mut c = [0i32; 16];
c[15] = 1;
let (tc, t1, lv, sg, runs, tz) = derive_cavlc_fields(&c).unwrap();
assert_eq!(tc, 1);
assert_eq!(t1, 1);
assert!(lv.is_empty());
assert_eq!(sg, 0);
assert_eq!(runs, vec![15]);
assert_eq!(tz, 15);
}
#[test]
fn derive_single_plus_one_at_bottom() {
let mut c = [0i32; 16];
c[0] = 1;
let (tc, t1, lv, _sg, runs, tz) = derive_cavlc_fields(&c).unwrap();
assert_eq!(tc, 1);
assert_eq!(t1, 1);
assert!(lv.is_empty());
assert_eq!(runs, vec![0]);
assert_eq!(tz, 0);
}
#[test]
fn derive_three_nonzeros_spec_layout() {
let mut c = [0i32; 16];
c[2] = 3;
c[5] = 2;
c[10] = -4;
let (tc, _t1, _lv, _sg, runs, tz) = derive_cavlc_fields(&c).unwrap();
assert_eq!(tc, 3);
assert_eq!(runs, vec![4, 2, 2]);
assert_eq!(tz, 8);
}
}