Update non-aligned byte reading for bulk operation

That is, instead of updating the queue on a byte-by-byte basis,
process the whole buffer in sensibly-sized chunks and update
the queue at the very end.

Author: tuffy

src/lib.rs

@@ -1077,6 +1077,24 @@ pub trait Endianness: Sized {
         W: BitWrite,
         S: SignedInteger;
 
+    /// Reads whole set of bytes to output buffer
+    #[inline]
+    fn read_bytes<R>(
+        reader: &mut R,
+        queue_value: &mut u8,
+        queue_bits: &mut u32,
+        buf: &mut [u8],
+    ) -> io::Result<()>
+    where
+        R: io::Read,
+    {
+        // a naive implementation which works anywhere
+        buf.iter_mut().try_for_each(|b| {
+            *b = Self::read_bits_fixed::<8, _, _>(reader.by_ref(), queue_value, queue_bits)?;
+            Ok(())
+        })
+    }
+
     /// Reads convertable numeric value from reader in this endianness
     fn read_primitive<R, V>(r: &mut R) -> io::Result<V>
     where
@@ -2100,6 +2118,59 @@ impl Endianness for BigEndian {
         }
     }
 
+    fn read_bytes<R>(
+        reader: &mut R,
+        queue_value: &mut u8,
+        queue_bits: &mut u32,
+        buf: &mut [u8],
+    ) -> io::Result<()>
+    where
+        R: io::Read,
+    {
+        const CHUNK_SIZE: usize = 1024;
+
+        // we don't modify the final queue_bits count
+        // but the naive implementation might
+        let queue_bits = *queue_bits;
+
+        let mut input_chunk: [u8; CHUNK_SIZE] = [0; CHUNK_SIZE];
+
+        for output_chunk in buf.chunks_mut(CHUNK_SIZE) {
+            let input_chunk = &mut input_chunk[0..output_chunk.len()];
+            reader.read_exact(input_chunk)?;
+
+            // shift down each byte in our input to eventually
+            // accomodate the contents of the bit queue
+            // and make that our output
+            output_chunk
+                .iter_mut()
+                .zip(input_chunk.iter())
+                .for_each(|(o, i)| {
+                    *o = i >> queue_bits;
+                });
+
+            // include leftover bits from the next byte
+            // shifted to the top
+            output_chunk[1..]
+                .iter_mut()
+                .zip(input_chunk.iter())
+                .for_each(|(o, i)| {
+                    *o |= *i << (u8::BITS_SIZE - queue_bits);
+                });
+
+            // finally, prepend the queue's contents
+            // to the first byte in the chunk
+            // while replacing those contents
+            // with the final byte of the input
+            output_chunk[0] |= core::mem::replace(
+                queue_value,
+                input_chunk.last().unwrap() << (u8::BITS_SIZE - queue_bits),
+            );
+        }
+
+        Ok(())
+    }
+
     #[inline]
     fn read_primitive<R, V>(r: &mut R) -> io::Result<V>
     where
@@ -2553,6 +2624,50 @@ impl Endianness for LittleEndian {
         }
     }
 
+    fn read_bytes<R>(
+        reader: &mut R,
+        queue_value: &mut u8,
+        queue_bits: &mut u32,
+        buf: &mut [u8],
+    ) -> io::Result<()>
+    where
+        R: io::Read,
+    {
+        const CHUNK_SIZE: usize = 1024;
+
+        // we don't modify the final queue_bits count
+        // but the naive implementation might
+        let queue_bits = *queue_bits;
+
+        let mut input_chunk: [u8; CHUNK_SIZE] = [0; CHUNK_SIZE];
+
+        for output_chunk in buf.chunks_mut(CHUNK_SIZE) {
+            let input_chunk = &mut input_chunk[0..output_chunk.len()];
+            reader.read_exact(input_chunk)?;
+
+            output_chunk
+                .iter_mut()
+                .zip(input_chunk.iter())
+                .for_each(|(o, i)| {
+                    *o = i << queue_bits;
+                });
+
+            output_chunk[1..]
+                .iter_mut()
+                .zip(input_chunk.iter())
+                .for_each(|(o, i)| {
+                    *o |= i >> (u8::BITS_SIZE - queue_bits);
+                });
+
+            output_chunk[0] |= core::mem::replace(
+                queue_value,
+                input_chunk.last().unwrap() >> (u8::BITS_SIZE - queue_bits),
+            );
+        }
+
+        Ok(())
+    }
+
     #[inline]
     fn read_primitive<R, V>(r: &mut R) -> io::Result<V>
     where

src/read.rs

@@ -1715,14 +1715,12 @@ impl<R: io::Read, E: Endianness> BitRead for BitReader<R, E> {
     /// assert!(reader.read_bytes(&mut buf).is_ok());
     /// assert_eq!(&buf, b"bar");
     /// ```
+    #[inline]
     fn read_bytes(&mut self, buf: &mut [u8]) -> io::Result<()> {
         if BitRead::byte_aligned(self) {
             self.reader.read_exact(buf)
         } else {
-            for b in buf.iter_mut() {
-                *b = self.read_unsigned::<8, _>()?;
-            }
-            Ok(())
+            E::read_bytes(&mut self.reader, &mut self.value, &mut self.bits, buf)
         }
     }
 

tests/random-3be.bin

@@ -0,0 +1 @@
+x^㍀���Mt���oi���~3�J+�xm�v�M�v�^��xk.Oˈ���`��kq����r��ϙ|�GG��~l��{P4
Hg'~\��#��ʼ��>T�"r[ushHikT!T��u���x6

tests/random-3le.bin

@@ -242,6 +242,44 @@ fn test_reader_huffman_be() {
     assert_eq!(r.read_huffman::<SomeTree>().unwrap(), 2);
 }
 
+#[test]
+fn test_read_chunks_be() {
+    use bitstream_io::{BigEndian, BitRead, BitReader};
+
+    let data: &[u8] = &[0b1011_0001, 0b1110_1101, 0b0011_1011, 0b1100_0001];
+    let mut chunk: [u8; 2] = [0; 2];
+
+    // test non-aligned chunk reading
+    let mut r = BitReader::endian(data, BigEndian);
+    assert_eq!(r.read::<2, u8>().unwrap(), 0b10);
+    r.read_bytes(&mut chunk).unwrap();
+    assert_eq!(&chunk, &[0b11_0001_11, 0b10_1101_00]);
+    assert_eq!(r.read::<14, u16>().unwrap(), 0b11_1011_1100_0001);
+
+    // test the smallest chunk
+    let mut chunk = 0;
+    let mut r = BitReader::endian(data, BigEndian);
+    assert_eq!(r.read::<2, u8>().unwrap(), 0b10);
+    r.read_bytes(core::slice::from_mut(&mut chunk)).unwrap();
+    assert_eq!(chunk, 0b11_0001_11);
+    r.read_bytes(core::slice::from_mut(&mut chunk)).unwrap();
+    assert_eq!(chunk, 0b10_1101_00);
+    assert_eq!(r.read::<14, u16>().unwrap(), 0b11_1011_1100_0001);
+
+    // test a larger chunk
+    let data = include_bytes!("random.bin");
+    let mut chunk: [u8; 127] = [0; 127];
+
+    let mut r = BitReader::endian(data.as_slice(), BigEndian);
+    assert_eq!(r.read::<3, u8>().unwrap(), 0b000);
+    r.read_bytes(&mut chunk).unwrap();
+    assert_eq!(
+        chunk.as_slice(),
+        include_bytes!("random-3be.bin").as_slice()
+    );
+    assert_eq!(r.read::<5, u8>().unwrap(), 0b10110);
+}
+
 #[test]
 fn test_reader_le() {
     use bitstream_io::{BitRead, BitReader, LittleEndian};
@@ -468,6 +506,44 @@ fn test_reader_huffman_le() {
     assert_eq!(r.read_huffman::<SomeTree>().unwrap(), 3);
 }
 
+#[test]
+fn test_read_chunks_le() {
+    use bitstream_io::{BitRead, BitReader, LittleEndian};
+
+    let data: &[u8] = &[0b1011_0001, 0b1110_1101, 0b0011_1011, 0b1100_0001];
+    let mut chunk: [u8; 2] = [0; 2];
+
+    // test non-aligned chunk reading
+    let mut r = BitReader::endian(data, LittleEndian);
+    assert_eq!(r.read::<2, u8>().unwrap(), 0b01);
+    r.read_bytes(&mut chunk).unwrap();
+    assert_eq!(&chunk, &[0b01_1011_00, 0b11_1110_11]);
+    assert_eq!(r.read::<14, u16>().unwrap(), 0b1100_0001_0011_10);
+
+    // test the smallest chunk
+    let mut chunk = 0;
+    let mut r = BitReader::endian(data, LittleEndian);
+    assert_eq!(r.read::<2, u8>().unwrap(), 0b01);
+    r.read_bytes(core::slice::from_mut(&mut chunk)).unwrap();
+    assert_eq!(chunk, 0b01_1011_00);
+    r.read_bytes(core::slice::from_mut(&mut chunk)).unwrap();
+    assert_eq!(chunk, 0b11_1110_11);
+    assert_eq!(r.read::<14, u16>().unwrap(), 0b1100_0001_0011_10);
+
+    // test a larger chunk
+    let data = include_bytes!("random.bin");
+    let mut chunk: [u8; 127] = [0; 127];
+
+    let mut r = BitReader::endian(data.as_slice(), LittleEndian);
+    assert_eq!(r.read::<3, u8>().unwrap(), 0b010);
+    r.read_bytes(&mut chunk).unwrap();
+    assert_eq!(
+        chunk.as_slice(),
+        include_bytes!("random-3le.bin").as_slice()
+    );
+    assert_eq!(r.read::<5, u8>().unwrap(), 0b00110);
+}
+
 #[test]
 fn test_reader_io_errors_be() {
     use bitstream_io::{BigEndian, BitRead, BitReader};