- commit
- 28acfe16eaa673520ef05c10a1ebf6c59e3904a6
- parent
- f53de5e47c6a2cadc2e3f951d89281f1ef61adef
- author
- Dami <dami.etoile@gmail.com>
| M | crates/ | 189 | ++++++++++++++++++++++++++++-- |
@@ -5,37 +5,83 @@ use super::MAX_ENCODE_BUFFER;
55 66 pub(super) struct AnnexBAccessUnitParser {77 buffer: Vec<u8>,8+ next_aud_scan: usize,9+ has_access_unit_start: bool,10+ 11+ scanned_candidates: usize,812 }913 1014 impl AnnexBAccessUnitParser {1115 pub(super) fn push(&mut self, bytes: &[u8]) -> Result<Vec<Vec<u8>>> {12- self.buffer.extend_from_slice(bytes);16+ if bytes.is_empty() {17+ return Ok(Vec::new());18+ }19+ let encoded_len = self.buffer.len().checked_add(bytes.len());1320 ensure!(14- self.buffer.len() <= MAX_ENCODE_BUFFER,21+ encoded_len.is_some_and(|len| len <= MAX_ENCODE_BUFFER),1522 "encoded stream exceeded {MAX_ENCODE_BUFFER} bytes without a frame boundary"1623 );24+ self.buffer.extend_from_slice(bytes);1725 1826 let mut frames = Vec::new();19- while let Some(first) = find_h264_aud(&self.buffer, 0) {20- if first > 0 {21- self.buffer.drain(..first);27+ loop {28+ let scan = find_h264_aud(&self.buffer, self.next_aud_scan);29+ 30+ {31+ self.scanned_candidates += scan.examined;2232 }23- let Some(next) = find_h264_aud(&self.buffer, 4) else {33+ let Some(aud) = scan.found else {34+ self.next_aud_scan = scan.resume_at;2435 break;2536 };26- let frame: Vec<u8> = self.buffer.drain(..next).collect();37+38+ 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+ }50+51+ let frame: Vec<u8> = self.buffer.drain(..aud).collect();2752 if !frame.is_empty() {2853 frames.push(frame);2954 }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;3058 }3159 Ok(frames)3260 }3361 }3462 35-fn find_h264_aud(data: &[u8], start: usize) -> Option<usize> {63+struct AudScan {64+ found: Option<usize>,65+ resume_at: usize,66+ 67+ examined: usize,68+}69+70+fn find_h264_aud(data: &[u8], start: usize) -> AudScan {3671 let mut index = start;72+ 73+ let mut examined = 0;3774 while index + 4 <= data.len() {75+ 76+ {77+ examined += 1;78+ }3879 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+ }3985 Some(index + 4)4086 } else if data.get(index..index + 3) == Some(&[0, 0, 1]) {4187 Some(index + 3)@@ -45,11 +91,21 @@ fn find_h264_aud(data: &[u8], start: usize) -> Option<usize> {
4591 if let Some(nal_offset) = nal_offset4692 && data.get(nal_offset).is_some_and(|byte| byte & 0x1f == 9)4793 {48- return Some(index);94+ return AudScan {95+ found: Some(index),96+ resume_at: index,97+ 98+ examined,99+ };49100 }50101 index += 1;51102 }52- None103+ AudScan {104+ found: None,105+ resume_at: index,106+ 107+ examined,108+ }53109 }54110 55111 @@ -71,4 +127,117 @@ mod tests {
71127 72128 assert_eq!(frames, vec![first.to_vec(), second.to_vec()]);73129 }130+131+ 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();137+138+ 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());142+143+ assert_eq!(144+ frames,145+ vec![first.to_vec(), second.to_vec()],146+ "split at byte {split}"147+ );148+ }149+150+ 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+ }157+158+ 159+ fn scans_large_chunked_access_unit_in_linear_work() {160+ const PAYLOAD_LEN: usize = 1024 * 1024;161+ const CHUNK_LEN: usize = 257;162+163+ 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();170+171+ 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+ }181+182+ 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+ }190+191+ 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();211+212+ 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+ }219+220+ 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;229+230+ 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);240+241+ assert_eq!(parser.push(&second).unwrap(), vec![first.to_vec()]);242+ }74243 }