Merge pull request #32 from njames93/vbr-encoding

Add support for encoding LLVM style variable bit rate integers

Author: tuffy

src/read.rs

@@ -25,6 +25,34 @@ use super::{
 
 use core::convert::TryInto;
 
+/// An error returned if performing VBR read overflows
+#[derive(Copy, Clone, Debug)]
+pub(crate) struct VariableWidthOverflow;
+
+impl core::fmt::Display for VariableWidthOverflow {
+    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+        "variable bit rate overflowed".fmt(f)
+    }
+}
+
+impl core::error::Error for VariableWidthOverflow {}
+
+impl From<VariableWidthOverflow> for io::Error {
+    fn from(VariableWidthOverflow: VariableWidthOverflow) -> Self {
+        io::Error::new(
+            #[cfg(feature = "std")]
+            {
+                io::ErrorKind::StorageFull
+            },
+            #[cfg(not(feature = "std"))]
+            {
+                io::ErrorKind::Other
+            },
+            "variable bit rate overflow",
+        )
+    }
+}
+
 /// A trait for anything that can read a variable number of
 /// potentially un-aligned values from an input stream
 pub trait BitRead {
@@ -874,6 +902,97 @@ pub trait BitRead {
         T::from_bits(|| self.read_bit())
     }
 
+    /// Reads a number using a variable using a variable-width integer.
+    /// This optimises the case when the number is small.
+    ///
+    /// The integer is mapped to an unsigned value using zigzag encoding.
+    /// For an integer X:
+    ///   - if X >= 0 -> 2X
+    ///   - else -> -2X + 1
+    ///
+    /// # Errors
+    ///
+    /// Passes along any I/O error from the underlying stream.
+    /// Returns an error if the data read would overflow the size of the result
+    ///
+    /// # Example
+    /// ```
+    /// use bitstream_io::{BitReader, BitRead, BigEndian};
+    ///
+    /// let bytes: &[u8] = &[0b0111_1100, 0b1100_0001];
+    /// let mut r = BitReader::endian(bytes, BigEndian);
+    /// assert_eq!(r.read_unsigned_vbr::<4, u32>().unwrap(), 7);
+    /// assert_eq!(r.read_unsigned_vbr::<4, u32>().unwrap(), 100);
+    /// ```
+    /// ```
+    /// use bitstream_io::{BitReader, BitRead, BigEndian};
+    ///
+    /// let bytes: &[u8] = &[0b1111_1111, 0b0011_1000, 0b1000_0100, 0b1000_1000, 0b1000_0000];
+    /// let mut r = BitReader::endian(bytes, BigEndian);
+    /// assert_eq!(r.read_unsigned_vbr::<4, u8>().unwrap(), 255); // Tries to read <011><111><111>
+    /// assert!(r.read_unsigned_vbr::<4, u8>().is_err()); // Tries to read a value of <100><000><000>
+    /// assert!(r.read_unsigned_vbr::<4, u8>().is_err()); // Tries to read a value of <000><000><000><000>
+    /// ```
+    fn read_unsigned_vbr<const FIELD_SIZE: u32, U: UnsignedInteger>(&mut self) -> io::Result<U> {
+        const { assert!(FIELD_SIZE >= 2 && FIELD_SIZE < U::BITS_SIZE) };
+        let payload_bits = FIELD_SIZE - 1;
+        let mut value = U::ZERO;
+        let mut shift = 0u32;
+        loop {
+            let (data, continuation) = self.read_unsigned::<FIELD_SIZE, U>().map(|item| {
+                (
+                    item & ((U::ONE << payload_bits) - U::ONE),
+                    (item >> payload_bits) != U::ZERO,
+                )
+            })?;
+            let shifted = data << shift;
+            value |= shifted;
+            if !continuation {
+                if (data << shift) >> shift == data {
+                    break Ok(value);
+                } else {
+                    break Err(VariableWidthOverflow {}.into());
+                }
+            }
+            shift += payload_bits;
+            if shift >= U::BITS_SIZE {
+                break Err(VariableWidthOverflow {}.into());
+            }
+        }
+    }
+
+    /// Reads a number using a variable using a variable-width integer.
+    /// This optimises the case when the number is small.
+    ///
+    /// The integer is mapped to an unsigned value using zigzag encoding.
+    /// For an integer X:
+    ///   - if X >= 0 -> 2X
+    ///   - else -> -2X + 1
+    ///
+    /// # Errors
+    ///
+    /// Passes along any I/O error from the underlying stream.
+    /// Returns an error if the data read would overflow the size of the result
+    ///
+    /// # Example
+    /// ```
+    /// use bitstream_io::{BitReader, BitRead, BigEndian};
+    ///
+    /// let bytes: &[u8] = &[0b0110_1011, 0b1100_0001];
+    /// let mut r = BitReader::endian(bytes, BigEndian);
+    /// assert_eq!(r.read_signed_vbr::<4, i32>().unwrap(), 3);
+    /// assert_eq!(r.read_signed_vbr::<4, i32>().unwrap(), -50);
+    /// ```
+    fn read_signed_vbr<const FIELD_SIZE: u32, I: SignedInteger>(&mut self) -> io::Result<I> {
+        self.read_unsigned_vbr::<FIELD_SIZE, I::Unsigned>()
+            .map(|zig_zag| {
+                let shifted = zig_zag >> 1;
+                let complimented = zig_zag & <I::Unsigned as crate::Numeric>::ONE;
+                let neg = I::ZERO - complimented.as_non_negative();
+                shifted.as_non_negative() ^ neg
+            })
+    }
+
     /// Creates a "by reference" adaptor for this `BitRead`
     ///
     /// The returned adapter also implements `BitRead`

src/write.rs

@@ -962,6 +962,78 @@ pub trait BitWrite {
         T::to_bits(value, |b| self.write_bit(b))
     }
 
+    /// Writes a number using a variable using a variable-width integer.
+    /// This optimises the case when the number is small.
+    ///
+    /// Given a 4-bit VBR field, any 3-bit value (0 through 7) is encoded directly, with the high bit set to zero.
+    /// Values larger than N-1 bits emit their bits in a series of N-1 bit chunks, where all but the last set the high bit.
+    ///
+    /// # Errors
+    ///
+    /// Passes along any I/O error from the underlying stream.
+    ///
+    /// # Example
+    /// ```
+    /// use std::io::Write;
+    /// use bitstream_io::{BigEndian, BitWriter, BitWrite};
+    /// let mut writer = BitWriter::endian(Vec::new(), BigEndian);
+    /// writer.write_unsigned_vbr::<4,_>(7u32);
+    /// writer.write_unsigned_vbr::<4,_>(100u32);
+    /// assert_eq!(writer.into_writer(), [0b0111_1100, 0b1100_0001]);
+    /// ```
+    fn write_unsigned_vbr<const FIELD_SIZE: u32, U: UnsignedInteger>(
+        &mut self,
+        value: U,
+    ) -> io::Result<()> {
+        const { assert!(FIELD_SIZE >= 2 && FIELD_SIZE < U::BITS_SIZE) };
+        let payload_bits = FIELD_SIZE - 1;
+        let continuation_bit = U::ONE.shl(payload_bits);
+        let payload_mask = continuation_bit.sub(U::ONE);
+        let mut value = value;
+
+        loop {
+            let payload = value & payload_mask;
+            value >>= payload_bits;
+            if value != U::ZERO {
+                self.write_unsigned::<FIELD_SIZE, U>(payload | continuation_bit)?;
+            } else {
+                self.write_unsigned::<FIELD_SIZE, U>(payload)?;
+                break;
+            }
+        }
+        Ok(())
+    }
+
+    /// Writes a number using a variable using a variable-width integer.
+    /// This optimises the case when the number is small.
+    ///
+    /// The integer is mapped to an unsigned value using zigzag encoding.
+    /// For an integer X:
+    ///   - if X >= 0 -> 2X
+    ///   - else -> -2X + 1
+    ///
+    /// # Errors
+    ///
+    /// Passes along any I/O error from the underlying stream.
+    ///
+    /// # Example
+    /// ```
+    /// use std::io::Write;
+    /// use bitstream_io::{BigEndian, BitWriter, BitWrite};
+    /// let mut writer = BitWriter::endian(Vec::new(), BigEndian);
+    /// writer.write_signed_vbr::<4,_>(3);
+    /// writer.write_signed_vbr::<4,_>(-50);
+    /// assert_eq!(writer.into_writer(), [0b0110_1011, 0b1100_0001]);
+    /// ```
+    #[inline]
+    fn write_signed_vbr<const FIELD_SIZE: u32, I: SignedInteger>(
+        &mut self,
+        value: I,
+    ) -> io::Result<()> {
+        let zig_zag = value.shl(1).bitxor(value.shr(I::BITS_SIZE - 1));
+        self.write_unsigned_vbr::<FIELD_SIZE, _>(zig_zag.as_non_negative())
+    }
+
     /// Creates a "by reference" adaptor for this `BitWrite`
     ///
     /// The returned adapter also implements `BitWrite`

tests/read.rs

@@ -110,6 +110,20 @@ fn test_reader_be() {
     assert_eq!(r.read_unary::<1>().unwrap(), 3);
     assert_eq!(r.read_unary::<1>().unwrap(), 0);
 
+    // reading unsigned vbr
+    let mut r = BitReader::endian(actual_data.as_slice(), BigEndian);
+    assert_eq!(r.read_unsigned_vbr::<4, u8>().unwrap(), 11);
+    assert_eq!(r.read_unsigned_vbr::<4, u8>().unwrap(), 238);
+    assert_eq!(r.read_unsigned_vbr::<4, u8>().unwrap(), 99);
+    assert!(r.read_unsigned_vbr::<4, u8>().is_err());
+
+    // reading signed vbr
+    let mut r = BitReader::endian(actual_data.as_slice(), BigEndian);
+    assert_eq!(r.read_signed_vbr::<4, i8>().unwrap(), -6);
+    assert_eq!(r.read_signed_vbr::<4, i8>().unwrap(), 119);
+    assert_eq!(r.read_signed_vbr::<4, i8>().unwrap(), -50);
+    assert!(r.read_signed_vbr::<4, i8>().is_err());
+
     // byte aligning
     let mut r = BitReader::endian(actual_data.as_slice(), BigEndian);
     assert_eq!(r.read_var::<u32>(3).unwrap(), 5);

tests/write.rs

@@ -134,6 +134,20 @@ fn test_writer_be() {
     w.write_unary::<1>(5).unwrap();
     assert_eq!(w.into_writer().as_slice(), &final_data);
 
+    // writing unsigned vbr
+    let mut w = BitWriter::endian(Vec::with_capacity(4), BigEndian);
+    w.write_unsigned_vbr::<4, _>(11u8).unwrap(); // 001 011 -> <1>011 <0>001
+    w.write_unsigned_vbr::<4, _>(238u8).unwrap(); // 011 101 110 -> <1>110 <1>101 <0>011
+    w.write_unsigned_vbr::<4, _>(99u8).unwrap(); // 001 100 011 -> <1>011 <1>100 <0>001
+    assert_eq!(w.into_writer().as_slice(), &final_data);
+
+    // writing signed vbr
+    let mut w = BitWriter::endian(Vec::with_capacity(4), BigEndian);
+    w.write_signed_vbr::<4, _>(-6i16).unwrap(); // 001 011 -> <1>011 <0>001
+    w.write_signed_vbr::<4, _>(119i16).unwrap(); // 011 101 110 -> <1>110 <1>101 <0>011
+    w.write_signed_vbr::<4, _>(-50i16).unwrap(); // 001 100 011 -> <1>011 <1>100 <0>001
+    assert_eq!(w.into_writer().as_slice(), &final_data);
+
     // byte aligning
     let aligned_data = [0xA0, 0xE0, 0x3B, 0xC0];
     let mut w = BitWriter::endian(Vec::with_capacity(4), BigEndian);