Faster H264Reader (webrtc-rs#506)

* perf: speedup H264Parser

* refactor: split code for buffer management

Author: anatawa12

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,62 +132,105 @@ 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<R: Read> {
     reader: R,
-    nal_buffer: BytesMut,
-    count_of_consecutive_zero_bytes: usize,
+    // reading buffers
+    buffer: ReadBuffer,
+    // for reading
     nal_prefix_parsed: bool,
-    read_buffer: Vec<u8>,
-    temp_buf: Vec<u8>,
+    count_of_consecutive_zero_bytes: usize,
+    nal_buffer: BytesMut,
 }
 
 impl<R: Read> H264Reader<R> {
     /// new creates new `H264Reader` with `capacity` sized read buffer.
     pub fn new(reader: R, capacity: usize) -> H264Reader<R> {
         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<Bytes> {
-        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<Option<u8>> {
+        if self.buffer.in_buffer() == 0 {
+            self.buffer.fill_buffer(&mut self.reader)?;
 
-        Ok(Bytes::from(
-            self.read_buffer
-                .drain(..num_should_read)
-                .collect::<Vec<u8>>(),
-        ))
+            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<usize> {
-        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<R: Read> H264Reader<R> {
         }
 
         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);