diff --git a/media/src/io/h264_reader/mod.rs b/media/src/io/h264_reader/mod.rs index 362ac1189..a489e1e7a 100644 --- a/media/src/io/h264_reader/mod.rs +++ b/media/src/io/h264_reader/mod.rs @@ -4,7 +4,7 @@ mod h264_reader_test; use std::fmt; use std::io::Read; -use bytes::{BufMut, Bytes, BytesMut}; +use bytes::{BufMut, BytesMut}; use crate::error::{Error, Result}; @@ -132,14 +132,53 @@ impl NAL { const NAL_PREFIX_3BYTES: [u8; 3] = [0, 0, 1]; const NAL_PREFIX_4BYTES: [u8; 4] = [0, 0, 0, 1]; +/// Wrapper class around reading buffer +struct ReadBuffer { + buffer: Box<[u8]>, + read_end: usize, + filled_end: usize, +} + +impl ReadBuffer { + fn new(capacity: usize) -> ReadBuffer { + Self { + buffer: vec![0u8; capacity].into_boxed_slice(), + read_end: 0, + filled_end: 0, + } + } + + #[inline] + fn in_buffer(&self) -> usize { + self.filled_end - self.read_end + } + + fn consume(&mut self, consume: usize) -> &[u8] { + debug_assert!(self.read_end + consume <= self.filled_end); + let result = &self.buffer[self.read_end..][..consume]; + self.read_end += consume; + result + } + + pub(crate) fn fill_buffer(&mut self, reader: &mut impl Read) -> Result<()> { + debug_assert_eq!(self.read_end, self.filled_end); + + self.read_end = 0; + self.filled_end = reader.read(&mut self.buffer)?; + + Ok(()) + } +} + /// H264Reader reads data from stream and constructs h264 nal units pub struct H264Reader { reader: R, - nal_buffer: BytesMut, - count_of_consecutive_zero_bytes: usize, + // reading buffers + buffer: ReadBuffer, + // for reading nal_prefix_parsed: bool, - read_buffer: Vec, - temp_buf: Vec, + count_of_consecutive_zero_bytes: usize, + nal_buffer: BytesMut, } impl H264Reader { @@ -147,47 +186,51 @@ impl H264Reader { pub fn new(reader: R, capacity: usize) -> H264Reader { H264Reader { reader, - nal_buffer: BytesMut::new(), - count_of_consecutive_zero_bytes: 0, nal_prefix_parsed: false, - read_buffer: vec![], - temp_buf: vec![0u8; capacity], + buffer: ReadBuffer::new(capacity), + count_of_consecutive_zero_bytes: 0, + nal_buffer: BytesMut::new(), } } - fn read(&mut self, num_to_read: usize) -> Result { - let buf = &mut self.temp_buf; - while self.read_buffer.len() < num_to_read { - let n = match self.reader.read(buf) { - Ok(n) => { - if n == 0 { - break; - } - n - } - Err(e) => return Err(Error::Io(e.into())), - }; + fn read4(&mut self) -> Result<([u8; 4], usize)> { + let mut result = [0u8; 4]; + let mut result_filled = 0; + loop { + let in_buffer = self.buffer.in_buffer(); + + if in_buffer + result_filled >= 4 { + let consume = 4 - result_filled; + result[result_filled..].copy_from_slice(self.buffer.consume(consume)); + return Ok((result, 4)); + } + + result[result_filled..][..in_buffer].copy_from_slice(self.buffer.consume(in_buffer)); + result_filled += in_buffer; + + self.buffer.fill_buffer(&mut self.reader)?; - self.read_buffer.extend_from_slice(&buf[0..n]); + if self.buffer.in_buffer() == 0 { + return Ok((result, result_filled)); + } } + } - let num_should_read = if num_to_read <= self.read_buffer.len() { - num_to_read - } else { - self.read_buffer.len() - }; + fn read1(&mut self) -> Result> { + if self.buffer.in_buffer() == 0 { + self.buffer.fill_buffer(&mut self.reader)?; - Ok(Bytes::from( - self.read_buffer - .drain(..num_should_read) - .collect::>(), - )) + if self.buffer.in_buffer() == 0 { + return Ok(None); + } + } + + Ok(Some(self.buffer.consume(1)[0])) } fn bit_stream_starts_with_h264prefix(&mut self) -> Result { - let prefix_buffer = self.read(4)?; + let (prefix_buffer, n) = self.read4()?; - let n = prefix_buffer.len(); if n == 0 { return Err(Error::ErrIoEOF); } @@ -229,13 +272,10 @@ impl H264Reader { } loop { - let buffer = self.read(1)?; - let n = buffer.len(); + let Some(read_byte) = self.read1()? else { + break + }; - if n != 1 { - break; - } - let read_byte = buffer[0]; let nal_found = self.process_byte(read_byte); if nal_found { let nal_unit_type = NalUnitType::from(self.nal_buffer[0] & 0x1F);