use std::io::{self, Write};
pub(crate) const START_CODE: [u8; 4] = [0x00, 0x00, 0x00, 0x01];
#[allow(dead_code)]
const HEVC_NAL_TYPE_MASK: u8 = 0x3F;
pub struct HevcMux<W: Write> {
writer: W,
codec_private: Option<Vec<u8>>,
params_written: bool,
}
impl<W: Write> HevcMux<W> {
pub fn new(writer: W) -> Self {
Self {
writer,
codec_private: None,
params_written: false,
}
}
pub fn set_codec_private(&mut self, data: Vec<u8>) {
self.codec_private = Some(data);
}
pub fn write_frame(&mut self, _pts_ns: i64, data: &[u8]) -> io::Result<()> {
if !self.params_written {
self.params_written = true;
if let Some(cp) = &self.codec_private {
match hvcc_to_annex_b(cp) {
Some(params) => self.writer.write_all(¶ms)?,
None if !cp.is_empty() => {
return Err(crate::error::Error::HevcParamParse.into());
}
None => {}
}
}
}
let annex_b = length_prefixed_to_annex_b(data);
self.writer.write_all(&annex_b)
}
pub fn finish(&mut self) -> io::Result<()> {
self.writer.flush()
}
}
pub(crate) fn hvcc_to_annex_b(hvcc: &[u8]) -> Option<Vec<u8>> {
if hvcc.len() < 23 {
return None;
}
let num_arrays = hvcc[22] as usize;
let mut out = Vec::new();
let mut offset = 23;
let mut truncated = false;
for _ in 0..num_arrays {
if truncated || offset + 3 > hvcc.len() {
break;
}
offset += 1; let num_nalus = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize;
offset += 2;
for _ in 0..num_nalus {
if offset + 2 > hvcc.len() {
truncated = true;
break;
}
let nal_len = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize;
offset += 2;
if offset + nal_len > hvcc.len() {
truncated = true;
break;
}
if nal_len == 0 {
continue;
}
out.extend_from_slice(&START_CODE);
out.extend_from_slice(&hvcc[offset..offset + nal_len]);
offset += nal_len;
}
}
if out.is_empty() { None } else { Some(out) }
}
pub(crate) fn length_prefixed_to_annex_b(data: &[u8]) -> Vec<u8> {
if starts_with_start_code(data) {
return data.to_vec();
}
let mut out = Vec::with_capacity(data.len() + (data.len() / 32));
append_length_prefixed_as_annex_b(&mut out, data);
out
}
pub(crate) fn append_length_prefixed_as_annex_b(out: &mut Vec<u8>, data: &[u8]) {
let mut offset = 0;
let mut parsed_any = false;
while offset + 4 <= data.len() {
let len = u32::from_be_bytes([
data[offset],
data[offset + 1],
data[offset + 2],
data[offset + 3],
]) as usize;
offset += 4;
if offset + len > data.len() {
break;
}
parsed_any = true;
if len == 0 {
continue;
}
out.extend_from_slice(&START_CODE);
out.extend_from_slice(&data[offset..offset + len]);
offset += len;
}
if !parsed_any && !data.is_empty() {
out.extend_from_slice(data);
}
}
fn starts_with_start_code(data: &[u8]) -> bool {
data.starts_with(&START_CODE) || data.starts_with(&[0x00, 0x00, 0x01])
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn length_prefixed_converts_to_annex_b() {
let mut buf = Vec::new();
buf.extend_from_slice(&3u32.to_be_bytes());
buf.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
buf.extend_from_slice(&2u32.to_be_bytes());
buf.extend_from_slice(&[0xDD, 0xEE]);
let got = length_prefixed_to_annex_b(&buf);
let want = [
0x00, 0x00, 0x00, 0x01, 0xAA, 0xBB, 0xCC, 0x00, 0x00, 0x00, 0x01, 0xDD, 0xEE, ];
assert_eq!(&got[..], &want[..]);
}
#[test]
fn already_annex_b_passes_through_when_no_lengths_match() {
let raw = [0xAA, 0xBB, 0xCC];
let got = length_prefixed_to_annex_b(&raw);
assert_eq!(&got[..], &raw[..]);
}
#[test]
fn mid_nal_truncation_drops_trailing_nal_keeps_prefix() {
let mut raw = Vec::new();
raw.extend_from_slice(&2u32.to_be_bytes());
raw.extend_from_slice(&[0x11, 0x22]);
raw.extend_from_slice(&100u32.to_be_bytes());
raw.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
let got = length_prefixed_to_annex_b(&raw);
let want = [0x00, 0x00, 0x00, 0x01, 0x11, 0x22];
assert_eq!(&got[..], &want[..]);
}
#[test]
fn leading_annex_b_start_code_passes_through() {
let raw = [
0x00, 0x00, 0x00, 0x01, 0x26, 0x01, 0xDE, 0xAD, 0x00, 0x00, 0x00, 0x01, 0x02, 0x01, 0xBE, 0xEF, ];
let got = length_prefixed_to_annex_b(&raw);
assert_eq!(
&got[..],
&raw[..],
"Annex B input must pass through verbatim"
);
}
#[test]
fn leading_three_byte_start_code_passes_through() {
let raw = [0x00, 0x00, 0x01, 0x26, 0x01, 0xDE, 0xAD];
let got = length_prefixed_to_annex_b(&raw);
assert_eq!(&got[..], &raw[..]);
}
#[test]
fn hvcc_skips_zero_length_nal_entries() {
let mut hvcc = vec![0u8; 22];
hvcc.push(1); hvcc.push(33); hvcc.extend_from_slice(&2u16.to_be_bytes()); hvcc.extend_from_slice(&0u16.to_be_bytes()); hvcc.extend_from_slice(&3u16.to_be_bytes()); hvcc.extend_from_slice(&[0x42, 0x01, 0x01]);
let annex_b = hvcc_to_annex_b(&hvcc).expect("one valid NAL");
let want = [0x00, 0x00, 0x00, 0x01, 0x42, 0x01, 0x01];
assert_eq!(&annex_b[..], &want[..]);
}
#[test]
fn write_frame_errors_on_unparseable_non_empty_hvcc() {
let mut sink: Vec<u8> = Vec::new();
let mut mux = HevcMux::new(&mut sink);
mux.set_codec_private(vec![0xDE, 0xAD]); let err = mux.write_frame(0, &[]).unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn zero_length_nal_is_skipped_not_bare_start_code() {
let mut buf = Vec::new();
buf.extend_from_slice(&3u32.to_be_bytes());
buf.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
buf.extend_from_slice(&0u32.to_be_bytes()); buf.extend_from_slice(&2u32.to_be_bytes());
buf.extend_from_slice(&[0xDD, 0xEE]);
let got = length_prefixed_to_annex_b(&buf);
let want = [
0x00, 0x00, 0x00, 0x01, 0xAA, 0xBB, 0xCC, 0x00, 0x00, 0x00, 0x01, 0xDD, 0xEE, ];
assert_eq!(&got[..], &want[..]);
}
#[test]
fn all_zero_length_nals_emit_nothing() {
let mut buf = Vec::new();
buf.extend_from_slice(&0u32.to_be_bytes());
buf.extend_from_slice(&0u32.to_be_bytes());
let got = length_prefixed_to_annex_b(&buf);
assert!(got.is_empty(), "expected empty output, got {got:?}");
}
#[test]
fn hvcc_extracts_vps_sps_pps() {
let mut hvcc = vec![0u8; 22];
hvcc.push(3); for (nal_type, payload) in [
(32u8, [0x40, 0x01, 0x0C, 0x01]),
(33, [0x42, 0x01, 0x01, 0x01]),
(34, [0x44, 0x01, 0xC1, 0x72]),
] {
hvcc.push(nal_type & 0x3F);
hvcc.extend_from_slice(&1u16.to_be_bytes()); hvcc.extend_from_slice(&(payload.len() as u16).to_be_bytes());
hvcc.extend_from_slice(&payload);
}
let annex_b = hvcc_to_annex_b(&hvcc).expect("at least one NAL");
assert_eq!(annex_b.len(), 24);
assert_eq!(&annex_b[..4], &START_CODE);
assert_eq!(&annex_b[8..12], &START_CODE);
assert_eq!(&annex_b[16..20], &START_CODE);
assert_eq!(annex_b[4], 0x40); assert_eq!(annex_b[12], 0x42); assert_eq!(annex_b[20], 0x44); }
#[test]
fn hvcc_too_short_for_header_returns_none() {
assert!(hvcc_to_annex_b(&[0u8; 22]).is_none());
assert!(hvcc_to_annex_b(&[]).is_none());
}
#[test]
fn hvcc_zero_arrays_returns_none() {
let mut hvcc = vec![0u8; 22];
hvcc.push(0); assert!(hvcc_to_annex_b(&hvcc).is_none());
}
#[test]
fn hvcc_array_with_multiple_nalus() {
let mut hvcc = vec![0u8; 22];
hvcc.push(1); hvcc.push(33); hvcc.extend_from_slice(&2u16.to_be_bytes()); hvcc.extend_from_slice(&2u16.to_be_bytes()); hvcc.extend_from_slice(&[0x42, 0x01]);
hvcc.extend_from_slice(&3u16.to_be_bytes()); hvcc.extend_from_slice(&[0x44, 0x02, 0x03]);
let out = hvcc_to_annex_b(&hvcc).expect("two NALs");
let want = [
0x00, 0x00, 0x00, 0x01, 0x42, 0x01, 0x00, 0x00, 0x00, 0x01, 0x44, 0x02, 0x03, ];
assert_eq!(&out[..], &want[..]);
}
#[test]
fn hvcc_truncated_nal_length_stops_cleanly() {
let mut hvcc = vec![0u8; 22];
hvcc.push(2); hvcc.push(32);
hvcc.extend_from_slice(&1u16.to_be_bytes());
hvcc.extend_from_slice(&3u16.to_be_bytes());
hvcc.extend_from_slice(&[0x40, 0x01, 0x02]);
hvcc.push(33);
hvcc.extend_from_slice(&1u16.to_be_bytes());
hvcc.extend_from_slice(&100u16.to_be_bytes());
hvcc.extend_from_slice(&[0xAA, 0xBB]);
let out = hvcc_to_annex_b(&hvcc).expect("the one valid NAL");
assert_eq!(
&out[..],
&[0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x02],
"truncated trailing NAL dropped, valid prefix kept"
);
}
#[test]
fn hvcc_truncated_length_field_itself_stops() {
let mut hvcc = vec![0u8; 22];
hvcc.push(1);
hvcc.push(33);
hvcc.extend_from_slice(&2u16.to_be_bytes()); hvcc.extend_from_slice(&2u16.to_be_bytes()); hvcc.extend_from_slice(&[0x42, 0x01]);
hvcc.push(0x00); let out = hvcc_to_annex_b(&hvcc).expect("NAL0");
assert_eq!(&out[..], &[0x00, 0x00, 0x00, 0x01, 0x42, 0x01]);
}
#[test]
fn hvcc_array_header_truncated_stops_outer_loop() {
let mut hvcc = vec![0u8; 22];
hvcc.push(3); hvcc.push(32);
hvcc.extend_from_slice(&1u16.to_be_bytes());
hvcc.extend_from_slice(&2u16.to_be_bytes());
hvcc.extend_from_slice(&[0x40, 0x01]);
let out = hvcc_to_annex_b(&hvcc).expect("the one present NAL");
assert_eq!(&out[..], &[0x00, 0x00, 0x00, 0x01, 0x40, 0x01]);
}
#[test]
fn empty_input_yields_empty() {
assert!(length_prefixed_to_annex_b(&[]).is_empty());
}
#[test]
fn single_nal_length_prefix() {
let mut buf = 5u32.to_be_bytes().to_vec();
buf.extend_from_slice(&[0x26, 0x01, 0xAA, 0xBB, 0xCC]);
let got = length_prefixed_to_annex_b(&buf);
let mut want = START_CODE.to_vec();
want.extend_from_slice(&[0x26, 0x01, 0xAA, 0xBB, 0xCC]);
assert_eq!(got, want);
}
#[test]
fn non_length_prefixed_three_plus_bytes_passes_through() {
let raw = [0xFF, 0xFF, 0xFF, 0xFF, 0x11, 0x22];
let got = length_prefixed_to_annex_b(&raw);
assert_eq!(&got[..], &raw[..], "unparseable → passed through verbatim");
}
#[test]
fn append_into_caller_buffer_preserves_existing() {
let mut out = vec![0xDE, 0xAD];
let mut nal = 2u32.to_be_bytes().to_vec();
nal.extend_from_slice(&[0x11, 0x22]);
append_length_prefixed_as_annex_b(&mut out, &nal);
let mut want = vec![0xDE, 0xAD];
want.extend_from_slice(&START_CODE);
want.extend_from_slice(&[0x11, 0x22]);
assert_eq!(out, want);
}
#[test]
fn starts_with_start_code_detects_both_forms() {
assert!(starts_with_start_code(&[0x00, 0x00, 0x00, 0x01, 0x42]));
assert!(starts_with_start_code(&[0x00, 0x00, 0x01, 0x42]));
assert!(!starts_with_start_code(&[0x00, 0x00, 0x02, 0x42]));
assert!(!starts_with_start_code(&[0x42, 0x00, 0x00, 0x01]));
assert!(!starts_with_start_code(&[]));
}
#[test]
fn mux_empty_codec_private_emits_no_params() {
let mut sink: Vec<u8> = Vec::new();
let mut mux = HevcMux::new(&mut sink);
mux.set_codec_private(Vec::new());
let mut frame = 2u32.to_be_bytes().to_vec();
frame.extend_from_slice(&[0xAA, 0xBB]);
mux.write_frame(0, &frame).unwrap();
mux.finish().unwrap();
let mut want = START_CODE.to_vec();
want.extend_from_slice(&[0xAA, 0xBB]);
assert_eq!(sink, want);
}
#[test]
fn mux_no_codec_private_writes_frames_only() {
let mut sink: Vec<u8> = Vec::new();
let mut mux = HevcMux::new(&mut sink);
let mut frame = 2u32.to_be_bytes().to_vec();
frame.extend_from_slice(&[0xAA, 0xBB]);
mux.write_frame(0, &frame).unwrap();
let mut want = START_CODE.to_vec();
want.extend_from_slice(&[0xAA, 0xBB]);
assert_eq!(sink, want);
}
#[test]
fn mux_params_not_re_emitted_after_unparseable_error() {
let mut sink: Vec<u8> = Vec::new();
let mut mux = HevcMux::new(&mut sink);
mux.set_codec_private(vec![0xDE, 0xAD]); assert!(mux.write_frame(0, &[]).is_err(), "first frame errors");
let mut frame = 2u32.to_be_bytes().to_vec();
frame.extend_from_slice(&[0xAA, 0xBB]);
mux.write_frame(0, &frame).unwrap();
let mut want = START_CODE.to_vec();
want.extend_from_slice(&[0xAA, 0xBB]);
assert_eq!(sink, want, "params not re-emitted after the error");
}
#[test]
fn mux_annex_b_frame_passes_through() {
let mut sink: Vec<u8> = Vec::new();
let mut mux = HevcMux::new(&mut sink);
let frame = [0x00, 0x00, 0x00, 0x01, 0x26, 0x01, 0xDE];
mux.write_frame(0, &frame).unwrap();
assert_eq!(sink, frame);
}
#[test]
fn mux_writes_params_then_frames() {
let mut hvcc = vec![0u8; 22];
hvcc.push(1);
hvcc.push(33);
hvcc.extend_from_slice(&1u16.to_be_bytes());
hvcc.extend_from_slice(&3u16.to_be_bytes());
hvcc.extend_from_slice(&[0x42, 0x01, 0x01]);
let mut frame_data = Vec::new();
frame_data.extend_from_slice(&2u32.to_be_bytes());
frame_data.extend_from_slice(&[0xAA, 0xBB]);
let mut sink: Vec<u8> = Vec::new();
let mut mux = HevcMux::new(&mut sink);
mux.set_codec_private(hvcc);
mux.write_frame(0, &frame_data).unwrap();
mux.write_frame(40_000_000, &frame_data).unwrap();
mux.finish().unwrap();
assert_eq!(sink.len(), 7 + 6 + 6);
assert_eq!(&sink[0..4], &START_CODE);
assert_eq!(&sink[7..11], &START_CODE);
assert_eq!(&sink[13..17], &START_CODE);
assert_eq!(sink[4], 0x42);
assert_eq!(sink[11], 0xAA);
assert_eq!(sink[17], 0xAA);
}
}