Replace bitlab dependency with homegrown bit twiddling

bitlab was a bit of an odd dependency, providing some convenience
wrappers for writing bits across byte slices. It additionally
pulled in the num crate. Replace it by writing our own version of
the bit reading/writing routines, which should amount to much less
code and is more optimized for our use case since a number of
checks can be elided and it has a simpler error path.

quinn-h3/Cargo.toml

@@ -25,7 +25,6 @@ codecov = { repository = "djc/quinn" }
 maintenance = { status = "experimental" }
 
 [dependencies]
-bitlab = "0.8.1"
 bytes = "0.5.2"
 err-derive = "0.2.3"
 futures = "0.3.1"

quinn-h3/src/qpack/prefix_string/decode.rs

@@ -1,6 +1,3 @@
-#![allow(clippy::ptr_arg)] // bitlab doesn't handle `&[u8]` for some reason
-use bitlab::*;
-
 use super::BitWindow;
 
 #[derive(Debug, PartialEq)]
@@ -22,7 +19,7 @@ struct HuffmanDecoder {
 }
 
 impl HuffmanDecoder {
-    fn check_eof(&self, bit_pos: &mut BitWindow, input: &Vec<u8>) -> Result<Option<u32>, Error> {
+    fn check_eof(&self, bit_pos: &mut BitWindow, input: &[u8]) -> Result<Option<u32>, Error> {
         use std::cmp::Ordering;
         match ((bit_pos.byte + 1) as usize).cmp(&input.len()) {
             // Position is out-of-range
@@ -33,9 +30,9 @@ impl HuffmanDecoder {
             Ordering::Equal => {
                 let side = bit_pos.opposite_bit_window();
 
-                let rest = match input.get_u8(side.byte, side.bit, side.count) {
+                let rest = match read_bits(input, side.byte, side.bit, side.count) {
                     Ok(x) => x,
-                    Err(_) => {
+                    Err(()) => {
                         return Err(Error::MissingBits(side));
                     }
                 };
@@ -50,14 +47,14 @@ impl HuffmanDecoder {
         Err(Error::MissingBits(bit_pos.clone()))
     }
 
-    fn fetch_value(&self, bit_pos: &mut BitWindow, input: &Vec<u8>) -> Result<Option<u32>, Error> {
-        match input.get_u32(bit_pos.byte, bit_pos.bit, bit_pos.count) {
-            Ok(value) => Ok(Some(value)),
-            Err(_) => self.check_eof(bit_pos, &input),
+    fn fetch_value(&self, bit_pos: &mut BitWindow, input: &[u8]) -> Result<Option<u32>, Error> {
+        match read_bits(input, bit_pos.byte, bit_pos.bit, bit_pos.count) {
+            Ok(value) => Ok(Some(value as u32)),
+            Err(()) => self.check_eof(bit_pos, &input),
         }
     }
 
-    fn decode_next(&self, bit_pos: &mut BitWindow, input: &Vec<u8>) -> Result<Option<u8>, Error> {
+    fn decode_next(&self, bit_pos: &mut BitWindow, input: &[u8]) -> Result<Option<u8>, Error> {
         bit_pos.forwards(self.lookup);
 
         let value = match self.fetch_value(bit_pos, input) {
@@ -78,6 +75,29 @@ impl HuffmanDecoder {
     }
 }
 
+/// Read `len` bits from the `src` slice at the specified position
+///
+/// Never read more than 8 bits at a time. `bit_offset` may be larger than 8.
+fn read_bits(src: &[u8], mut byte_offset: u32, mut bit_offset: u32, len: u32) -> Result<u8, ()> {
+    if len == 0 || len > 8 || src.len() as u32 * 8 < (byte_offset * 8) + bit_offset + len {
+        return Err(());
+    }
+
+    // Deal with `bit_offset` > 8
+    byte_offset += bit_offset / 8;
+    bit_offset -= (bit_offset / 8) * 8;
+
+    Ok(if bit_offset + len <= 8 {
+        // Read all the bits from a single byte
+        (src[byte_offset as usize] << bit_offset) >> (8 - len)
+    } else {
+        // The range of bits spans over 2 bytes
+        let mut result = (src[byte_offset as usize] as u16) << 8;
+        result |= src[byte_offset as usize + 1] as u16;
+        ((result << bit_offset) >> (16 - len)) as u8
+    })
+}
+
 macro_rules! bits_decode {
     // general way
     (
@@ -307,6 +327,43 @@ impl HpackStringDecode for Vec<u8> {
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_read_bits() {
+        // Basic case (within one byte, aligned with start)
+        assert_eq!(read_bits(&[0b1010_1010], 0, 0, 5), Ok(0b1_0101));
+        // Within one byte, aligned with end of byte
+        assert_eq!(read_bits(&[0b1010_1010], 0, 3, 5), Ok(0b1010));
+        // Within one byte, unaligned with either side
+        assert_eq!(read_bits(&[0b1010_1010], 0, 3, 3), Ok(0b10));
+        // `len` == 0
+        assert_eq!(read_bits(&[0b1010_1010], 0, 0, 0), Err(()));
+        // `len` > 8
+        assert_eq!(read_bits(&[0b1010_1010], 0, 0, 9), Err(()));
+
+        // `bit_offset` > 7
+        assert_eq!(
+            read_bits(&[0b1010_1010, 0b1010_1010], 0, 8, 8),
+            Ok(0b1010_1010)
+        );
+        // Read spanning two bytes
+        assert_eq!(
+            read_bits(&[0b1010_1010, 0b1010_1010], 0, 4, 8),
+            Ok(0b1010_1010)
+        );
+        // Read with non-zero `byte_offset`
+        assert_eq!(
+            read_bits(&[0b1010_1010, 0b1010_1010], 1, 0, 5),
+            Ok(0b1_0101)
+        );
+        // Read with `bit_offset` > 7, unaligned with either side
+        assert_eq!(
+            read_bits(&[0b1010_1010, 0b1010_1010], 0, 10, 5),
+            Ok(0b1_0101)
+        );
+        // Read with `bit_offset` > 7 past end of input slice
+        assert_eq!(read_bits(&[0b1010_1010, 0b1010_1010], 0, 16, 5), Err(()));
+    }
+
     macro_rules! decoding {
         [ $( $code:expr => $( $byte:expr ),* ; )* ] => { $( {
             let bytes = vec![$( $byte ),*];

quinn-h3/src/qpack/prefix_string/encode.rs

@@ -1,5 +1,3 @@
-use bitlab::*;
-
 use super::BitWindow;
 
 #[derive(Debug, PartialEq)]
@@ -30,15 +28,6 @@ impl HuffmanEncoder {
         }
     }
 
-    fn create_error(&self, text: String) -> Error {
-        Error {
-            buffer_pos: self.buffer_pos.clone(),
-            len: self.buffer.len(),
-            capacity: self.buffer.capacity(),
-            text,
-        }
-    }
-
     fn ensure_free_space(&mut self, bit_count: u32) {
         let mut end_range = self.buffer_pos.clone();
         end_range.forwards(bit_count);
@@ -75,14 +64,7 @@ impl HuffmanEncoder {
             self.buffer_pos.forwards(if rest < 8 { rest } else { 8 });
             rest -= self.buffer_pos.count;
 
-            self.buffer
-                .set(
-                    self.buffer_pos.byte,
-                    self.buffer_pos.bit,
-                    self.buffer_pos.count,
-                    part,
-                )
-                .map_err(|x| self.create_error(x))?;
+            write_bits(&mut self.buffer, &self.buffer_pos, part)
         }
 
         Ok(())
@@ -93,6 +75,43 @@ impl HuffmanEncoder {
     }
 }
 
+/// Write bits from `value` to the `out` slice
+///
+/// Write the least significant `pos.count` bits from `value` to the position specified by
+/// `(pos.byte, pos.bit)`. Writes may span multiple bytes. `out` is expected to be long enough
+/// to write these bits; this is ensured by `HuffmanEncoder::ensure_free_space()`, which is
+/// always called prior to calling this function.
+///
+/// The bits to be written to are expected to be set to 1 when calling this function. Similarly,
+/// this function maintains the invariant that unused bits in the output bytes are set to 1.
+fn write_bits(out: &mut [u8], pos: &BitWindow, value: u8) {
+    debug_assert!(pos.bit < 8);
+    debug_assert!(pos.count <= 8);
+    debug_assert!(pos.count > 0);
+
+    if (pos.bit + pos.count) <= 8 {
+        // Bits to be written to fit in a single byte
+        debug_assert_eq!(out[pos.byte as usize] | PAD_LEFT[pos.bit as usize], 255);
+        let pad_left = out[pos.byte as usize] | PAD_RIGHT[(8 - pos.bit) as usize];
+        let shifted = value << (8 - pos.bit - pos.count) | PAD_LEFT[pos.bit as usize];
+        let pad_right = PAD_RIGHT[(8 - pos.count - pos.bit) as usize];
+        out[pos.byte as usize] = (pad_left & shifted) | pad_right;
+    } else {
+        // Bits to be written to span two bytes
+        debug_assert_eq!(out[pos.byte as usize] | PAD_LEFT[pos.bit as usize], 255);
+        let split = 8 - pos.bit;
+        let pad_left = out[pos.byte as usize] | PAD_RIGHT[split as usize];
+        let shifted = (value >> (pos.count - split)) | PAD_LEFT[pos.bit as usize];
+        out[pos.byte as usize] = pad_left & shifted;
+
+        let rem = 8 - (pos.count - split);
+        out[(pos.byte + 1) as usize] = (value << rem) | PAD_RIGHT[rem as usize];
+    }
+}
+
+const PAD_RIGHT: [u8; 9] = [0, 1, 3, 7, 15, 31, 63, 127, 255];
+const PAD_LEFT: [u8; 9] = [0, 128, 192, 224, 240, 248, 252, 254, 255];
+
 macro_rules! bits_encode {
     [ $( ( $len:expr => [ $( $byte:expr ),* ] ), )* ] => {
         [ $(
@@ -381,6 +400,42 @@ impl HpackStringEncode for Vec<u8> {
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_set_bits() {
+        let mut buf = [0b1111_1111; 16];
+        let mut pos = BitWindow::default();
+
+        // Write a full 8 bits into a single byte
+        pos.count = 8;
+        write_bits(&mut buf, &pos, 0b1_0101);
+        assert_eq!(&buf[..1], &[0b1_0101]);
+        pos.byte += 1;
+
+        // 7-bit byte-spanning writes at each possible bit offset
+        pos.count = 7;
+        for _ in 0..8 {
+            write_bits(&mut buf, &pos, 0b101_0101);
+            pos.forwards(7);
+        }
+        assert_eq!(
+            &buf[1..8],
+            &[
+                0b1010_1011,
+                0b0101_0110,
+                0b1010_1101,
+                0b0101_1010,
+                0b1011_0101,
+                0b0110_1010,
+                0b1101_0101
+            ]
+        );
+
+        // Single-write partial bits, aligned with byte start
+        pos.count = 5;
+        write_bits(&mut buf, &pos, 0b1_0101);
+        assert_eq!(&buf[8..9], &[0b1010_1111]);
+    }
+
     macro_rules! encoding {
         [ $( $code:expr => $( $byte:expr ),* ; )* ] => { $( {
             let bytes = vec![$( $byte ),*];