1use anyhow::{Result, ensure};23use super::MAX_ENCODE_BUFFER;456pub(super) struct AnnexBAccessUnitParser {7 buffer: Vec<u8>,8 next_aud_scan: usize,9 has_access_unit_start: bool,10 11 scanned_candidates: usize,12}1314impl AnnexBAccessUnitParser {15 pub(super) fn push(&mut self, bytes: &[u8]) -> Result<Vec<Vec<u8>>> {16 if bytes.is_empty() {17 return Ok(Vec::new());18 }19 let encoded_len = self.buffer.len().checked_add(bytes.len());20 ensure!(21 encoded_len.is_some_and(|len| len <= MAX_ENCODE_BUFFER),22 "encoded stream exceeded {MAX_ENCODE_BUFFER} bytes without a frame boundary"23 );24 self.buffer.extend_from_slice(bytes);2526 let mut frames = Vec::new();27 loop {28 let scan = find_h264_aud(&self.buffer, self.next_aud_scan);29 30 {31 self.scanned_candidates += scan.examined;32 }33 let Some(aud) = scan.found else {34 self.next_aud_scan = scan.resume_at;35 break;36 };3738 if !self.has_access_unit_start {39 if aud > 0 {40 self.buffer.drain(..aud);41 }42 self.has_access_unit_start = true;43 // Preserve the original parser's search semantics: once the44 // first AUD is at offset zero, look for its successor at45 // offset four. That is after a three-byte AUD's NAL header and46 // at a four-byte AUD's NAL header.47 self.next_aud_scan = 4;48 continue;49 }5051 let frame: Vec<u8> = self.buffer.drain(..aud).collect();52 if !frame.is_empty() {53 frames.push(frame);54 }55 // The AUD that ended the emitted access unit is now at offset56 // zero and begins the next one.57 self.next_aud_scan = 4;58 }59 Ok(frames)60 }61}6263struct AudScan {64 found: Option<usize>,65 resume_at: usize,66 67 examined: usize,68}6970fn find_h264_aud(data: &[u8], start: usize) -> AudScan {71 let mut index = start;72 73 let mut examined = 0;74 while index + 4 <= data.len() {75 76 {77 examined += 1;78 }79 let nal_offset = if data.get(index..index + 4) == Some(&[0, 0, 0, 1]) {80 // The start code is complete, but its NAL header may arrive in the81 // next chunk. Revisit this one candidate rather than skipping it.82 if index + 4 == data.len() {83 break;84 }85 Some(index + 4)86 } else if data.get(index..index + 3) == Some(&[0, 0, 1]) {87 Some(index + 3)88 } else {89 None90 };91 if let Some(nal_offset) = nal_offset92 && data.get(nal_offset).is_some_and(|byte| byte & 0x1f == 9)93 {94 return AudScan {95 found: Some(index),96 resume_at: index,97 98 examined,99 };100 }101 index += 1;102 }103 AudScan {104 found: None,105 resume_at: index,106 107 examined,108 }109}110111112mod tests {113 use super::*;114115 116 fn splits_access_units_on_aud() {117 let first = [0, 0, 0, 1, 9, 0x10, 0, 0, 1, 5, 1, 2, 3];118 let second = [0, 0, 0, 1, 9, 0x10, 0, 0, 1, 1, 4, 5];119 let third = [0, 0, 1, 9, 0x10, 0, 0, 1, 1, 6, 7];120 let mut parser = AnnexBAccessUnitParser::default();121122 let mut bytes = Vec::new();123 bytes.extend_from_slice(&first);124 bytes.extend_from_slice(&second);125 bytes.extend_from_slice(&third);126 let frames = parser.push(&bytes).unwrap();127128 assert_eq!(frames, vec![first.to_vec(), second.to_vec()]);129 }130131 132 fn recognizes_three_and_four_byte_start_codes_split_at_every_boundary() {133 let first = [0, 0, 0, 1, 9, 0x10, 0, 0, 1, 5, 1, 2, 3];134 let second = [0, 0, 1, 9, 0x10, 0, 0, 1, 1, 4, 5];135 let third = [0, 0, 0, 1, 9, 0x10, 0, 0, 1, 1, 6, 7];136 let stream = [first.as_slice(), second.as_slice(), third.as_slice()].concat();137138 for split in 1..stream.len() {139 let mut parser = AnnexBAccessUnitParser::default();140 let mut frames = parser.push(&stream[..split]).unwrap();141 frames.extend(parser.push(&stream[split..]).unwrap());142143 assert_eq!(144 frames,145 vec![first.to_vec(), second.to_vec()],146 "split at byte {split}"147 );148 }149150 let mut bytewise = AnnexBAccessUnitParser::default();151 let frames: Vec<Vec<u8>> = stream152 .chunks(1)153 .flat_map(|chunk| bytewise.push(chunk).unwrap())154 .collect();155 assert_eq!(frames, vec![first.to_vec(), second.to_vec()]);156 }157158 159 fn scans_large_chunked_access_unit_in_linear_work() {160 const PAYLOAD_LEN: usize = 1024 * 1024;161 const CHUNK_LEN: usize = 257;162163 let mut first = vec![0, 0, 0, 1, 9, 0x10, 0, 0, 1, 5];164 first.resize(first.len() + PAYLOAD_LEN, 0x55);165 let second = [0, 0, 1, 9, 0x10, 0, 0, 1, 1, 4, 5];166 let third = [0, 0, 0, 1, 9, 0x10];167 let stream = [first.as_slice(), second.as_slice(), third.as_slice()].concat();168 let mut parser = AnnexBAccessUnitParser::default();169 let mut frames = Vec::new();170171 for chunk in stream.chunks(CHUNK_LEN) {172 let before = parser.scanned_candidates;173 frames.extend(parser.push(chunk).unwrap());174 let work = parser.scanned_candidates - before;175 assert!(176 work <= chunk.len() + 8,177 "one {chunk_len}-byte append examined {work} candidates",178 chunk_len = chunk.len()179 );180 }181182 assert_eq!(frames, vec![first, second.to_vec()]);183 assert!(184 parser.scanned_candidates <= stream.len() + 4 * stream.len().div_ceil(CHUNK_LEN),185 "{} bytes caused {} candidate scans",186 stream.len(),187 parser.scanned_candidates188 );189 }190191 192 fn resets_scan_offset_after_emitting_each_access_unit() {193 let leading_junk = [0x44; 19];194 let first = [0, 0, 1, 9, 0x10, 0, 0, 1, 5, 1];195 let second = [0, 0, 0, 1, 9, 0x10, 0, 0, 1, 1, 2];196 let third = [0, 0, 1, 9, 0x10, 0, 0, 1, 1, 3];197 let fourth = [0, 0, 0, 1, 9, 0x10];198 let stream = [199 leading_junk.as_slice(),200 first.as_slice(),201 second.as_slice(),202 third.as_slice(),203 fourth.as_slice(),204 ]205 .concat();206 let mut parser = AnnexBAccessUnitParser::default();207 let frames: Vec<Vec<u8>> = stream208 .chunks(7)209 .flat_map(|chunk| parser.push(chunk).unwrap())210 .collect();211212 assert_eq!(213 frames,214 vec![first.to_vec(), second.to_vec(), third.to_vec()]215 );216 assert!(parser.has_access_unit_start);217 assert!(parser.next_aud_scan >= 4);218 }219220 221 fn oversized_append_is_rejected_without_allocating_or_poisoning_state() {222 let first = [0, 0, 0, 1, 9, 0x10, 0, 0, 1, 5, 1, 2, 3];223 let second = [0, 0, 1, 9, 0x10, 0, 0, 1, 1, 4, 5];224 let mut parser = AnnexBAccessUnitParser::default();225 assert!(parser.push(&first).unwrap().is_empty());226 let buffer_len = parser.buffer.len();227 let scan_offset = parser.next_aud_scan;228 let scan_work = parser.scanned_candidates;229230 let oversized = vec![0x55; MAX_ENCODE_BUFFER];231 let error = parser.push(&oversized).unwrap_err();232 assert!(233 error234 .to_string()235 .contains("encoded stream exceeded 16777216 bytes")236 );237 assert_eq!(parser.buffer.len(), buffer_len);238 assert_eq!(parser.next_aud_scan, scan_offset);239 assert_eq!(parser.scanned_candidates, scan_work);240241 assert_eq!(parser.push(&second).unwrap(), vec![first.to_vec()]);242 }243}