Desloppify code

Author: ArniDagur

src/hpack/huffman/decode.rs

@@ -0,0 +1,293 @@
+//! Huffman decoding for HPACK (RFC 7541 §5.2).
+//!
+//! The decoder has a fast path which decodes 12 bits a time. The decoder
+//! looks the next 12 bits up in `FAST_TABLE`, which maps directly to either
+//! one or two symbols to decode. Since ASCII letters and digits have 5-7 bit
+//! codes, one lookup usually decodes two symbols at once.
+//!
+//! Two situations fall off this fast path:
+//!
+//! - Codes longer than 12 bits: control characters and bytes >= 0x80
+//! - The end of the input, when fewer than 12 bits remain.
+//!
+//! These are decoded by walking the generated `DECODE_TABLE` one input byte at a time.
+
+use crate::hpack::huffman::table::{DECODE_TABLE, ENCODE_TABLE};
+use crate::hpack::DecoderError;
+
+use bytes::BytesMut;
+
+// DECODE_TABLE (in the generated `table.rs`) is a series of 256-entry
+// tables, walked one input byte at a time. A leaf entry (BRANCH bit clear)
+// holds a decoded symbol and the number of bits its code used. A branch
+// entry (BRANCH bit set) holds the index of the table for the next byte.
+// No valid code leads back to table 0, so a branch to 0 means the input is
+// not a valid code.
+const BRANCH: u16 = 0x8000;
+const TABLE_INDEX_MASK: u16 = 0x7f00;
+const TABLE_WIDTH: usize = 256;
+
+// Maps each possible 12-bit chunk of input to the symbol(s) it starts
+// with. Entry layout (u32):
+//
+//   bits  0..8   first decoded byte
+//   bits  8..16  second decoded byte (if any)
+//   bits 16..24  number of symbols decoded: 1, 2, or 0, where 0 means the
+//                code is longer than 12 bits and the slow path must run
+//   bits 24..32  how many of the 12 bits the decoded symbols used
+const FAST_BITS: usize = 12;
+const FAST_DECODE_TABLE: [u32; 1 << FAST_BITS] = build_fast_table();
+
+const fn build_fast_table() -> [u32; 1 << FAST_BITS] {
+    let mut table = [0; 1 << FAST_BITS];
+
+    // First fill in every index that starts with one whole code
+    let mut a = 0;
+    while a < 256 {
+        let (len1, code1) = ENCODE_TABLE[a];
+        if len1 <= FAST_BITS {
+            let rem = FAST_BITS - len1;
+            let base = (code1 as usize) << rem;
+            let single = (1 << 16) | ((len1 as u32) << 24) | a as u32;
+            let mut i = 0;
+            while i < (1 << rem) {
+                table[base + i] = single;
+                i += 1;
+            }
+        }
+        a += 1;
+    }
+
+    // Overwrite the indices where a second whole code fit right after the first
+    let mut a = 0;
+    while a < 256 {
+        let (len1, code1) = ENCODE_TABLE[a];
+        // A second code needs at least 5 more bits (the shortest code)
+        if len1 + 5 <= FAST_BITS {
+            let rem = FAST_BITS - len1;
+            let mut b = 0;
+            while b < 256 {
+                let (len2, code2) = ENCODE_TABLE[b];
+                if len2 <= rem {
+                    let rem2 = rem - len2;
+                    let base = ((code1 as usize) << rem) | ((code2 as usize) << rem2);
+                    let pair =
+                        (2 << 16) | (((len1 + len2) as u32) << 24) | ((b as u32) << 8) | a as u32;
+                    let mut i = 0;
+                    while i < (1 << rem2) {
+                        table[base + i] = pair;
+                        i += 1;
+                    }
+                }
+                b += 1;
+            }
+        }
+        a += 1;
+    }
+
+    table
+}
+
+// Decodes one code of any length by walking DECODE_TABLE one input byte at
+// a time, and writes the symbol to `dst[*o]`. Used for codes longer than
+// FAST_BITS bits and for the tail of the input.
+//
+// Inlining avoids a 35% performance regression on ASCII benchmarks.
+#[inline(always)]
+fn decode_code_slow(
+    acc: &mut u64,
+    bits: &mut usize,
+    dst: &mut [u8],
+    o: &mut usize,
+) -> Result<(), DecoderError> {
+    let mut table = 0;
+    loop {
+        let e = DECODE_TABLE[table * TABLE_WIDTH + (*acc >> 56) as usize];
+        if e & BRANCH == 0 {
+            let used = (e >> 8) as usize;
+            if used > *bits {
+                return Err(DecoderError::InvalidHuffmanCode);
+            }
+            dst[*o] = e as u8;
+            *o += 1;
+            *acc <<= used;
+            *bits -= used;
+            return Ok(());
+        }
+        if *bits < 8 {
+            return Err(DecoderError::InvalidHuffmanCode);
+        }
+        table = ((e & TABLE_INDEX_MASK) >> 8) as usize;
+        if table == 0 {
+            return Err(DecoderError::InvalidHuffmanCode);
+        }
+        *acc <<= 8;
+        *bits -= 8;
+    }
+}
+
+// Decodes a Huffman encoded string into the provided buffer.
+pub fn decode(src: &[u8], buf: &mut BytesMut) -> Result<BytesMut, DecoderError> {
+    let len = src.len();
+    let base_len = buf.len();
+
+    // Reserve the worst case output size. Every code is at least 5 bits,
+    // so `len` input bytes can't decode to more than `len * 8 / 5` output
+    // bytes. One more byte covers the fast path always writing two bytes
+    // even when it decoded only one symbol.
+    buf.resize(base_len + len * 8 / 5 + 1, 0);
+    let dst = &mut buf[base_len..];
+    let mut o = 0;
+
+    let mut acc: u64 = 0;
+    let mut bits: usize = 0;
+    let mut pos: usize = 0;
+
+    'outer: loop {
+        // Fill the bit buffer. While 8 or more input bytes remain, load
+        // the next 8 in one go and count as many whole bytes as fit under
+        // the bits already in the buffer (bringing it to 56-63 bits). The
+        // bytes that didn't fit still land in the low end of `acc`; the
+        // next refill just ORs the same values over them, so no harm done.
+        if pos + 8 <= len {
+            let word = u64::from_be_bytes(src[pos..pos + 8].try_into().unwrap());
+            acc |= word >> bits;
+            pos += (63 - bits) >> 3; // bytes now accounted for in `bits`
+            bits |= 56; // equals bits + 8 * (bytes added)
+        } else {
+            while bits <= 56 && pos < len {
+                acc |= (src[pos] as u64) << (56 - bits);
+                pos += 1;
+                bits += 8;
+            }
+        }
+
+        while bits >= FAST_BITS {
+            let entry = FAST_DECODE_TABLE[(acc >> (64 - FAST_BITS)) as usize];
+            let count = (entry >> 16) & 0xff;
+            if count == 0 {
+                // The next code is longer than 12 bits. Refill first if it
+                // might not be in the buffer yet (the longest code is 30
+                // bits), then decode it by walking DECODE_TABLE.
+                if bits < 30 && pos < len {
+                    continue 'outer;
+                }
+                decode_code_slow(&mut acc, &mut bits, dst, &mut o)?;
+                continue;
+            }
+            let consumed = (entry >> 24) as usize;
+            dst[o] = entry as u8;
+            dst[o + 1] = (entry >> 8) as u8;
+            o += count as usize;
+            acc <<= consumed;
+            bits -= consumed;
+        }
+
+        if pos >= len {
+            break;
+        }
+    }
+
+    // Tail: the input is exhausted and fewer than 12 bits remain.
+    while bits > 0 {
+        // The encoder pads the last byte with up to 7 one-bits, which are
+        // valid only between symbols. Anything else must decode.
+        if bits < 8 && (acc >> (64 - bits)) == (1 << bits) - 1 {
+            break;
+        }
+        decode_code_slow(&mut acc, &mut bits, dst, &mut o)?;
+    }
+
+    buf.truncate(base_len + o);
+    Ok(buf.split())
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    use bytes::BufMut;
+    use rand::rngs::StdRng;
+    use rand::{Rng, SeedableRng};
+
+    fn decode(src: &[u8]) -> Result<BytesMut, DecoderError> {
+        let mut buf = BytesMut::new();
+        super::decode(src, &mut buf)
+    }
+
+    // The simplest decoder we can write, used to double-check the real
+    // one: turn the input into a literal string of '0'/'1' characters and
+    // repeatedly strip off the first code that matches. No code is a
+    // prefix of another, so at most one can match. EOS (index 256) is left
+    // out on purpose: encoded EOS is an error per RFC 7541 §5.2.
+    fn reference_decode(src: &[u8]) -> Result<BytesMut, DecoderError> {
+        // Each symbol's code as a string of '0' and '1', e.g. "11111000".
+        let codes: Vec<String> = ENCODE_TABLE[..256]
+            .iter()
+            .map(|&(len, code)| format!("{code:0len$b}"))
+            .collect();
+        let bits: String = src.iter().map(|byte| format!("{byte:08b}")).collect();
+
+        let mut decoded = BytesMut::new();
+        let mut remaining_bits: &str = &bits;
+        while let Some(sym) = (0..256).find(|&sym| remaining_bits.starts_with(&codes[sym])) {
+            decoded.put_u8(sym as u8);
+            remaining_bits = &remaining_bits[codes[sym].len()..];
+        }
+        // We expect anything that cannot be matched to be padding. Padding is
+        // fewer than 8 bits, and all ones.
+        if remaining_bits.len() >= 8 || remaining_bits.chars().any(|bit| bit == '0') {
+            return Err(DecoderError::InvalidHuffmanCode);
+        }
+        Ok(decoded)
+    }
+
+    #[test]
+    fn decode_single_byte() {
+        assert_eq!("o", decode(&[0b00111111]).unwrap());
+        assert_eq!("0", decode(&[7]).unwrap());
+        assert_eq!("A", decode(&[(0x21 << 2) + 3]).unwrap());
+    }
+
+    #[test]
+    fn single_char_multi_byte() {
+        assert_eq!("#", decode(&[255, 160 + 15]).unwrap());
+        assert_eq!("$", decode(&[255, 200 + 7]).unwrap());
+        assert_eq!("\x0a", decode(&[255, 255, 255, 240 + 3]).unwrap());
+    }
+
+    #[test]
+    fn multi_char() {
+        assert_eq!("!0", decode(&[254, 1]).unwrap());
+        assert_eq!(" !", decode(&[0b01010011, 0b11111000]).unwrap());
+    }
+
+    // Random (mostly invalid) byte strings must produce identical results to reference impl
+    #[test]
+    fn matches_reference_on_arbitrary_bytes() {
+        let mut rng = StdRng::seed_from_u64(0xdeadbeefcafe);
+
+        for _ in 0..10_000 {
+            let len = rng.gen_range(0..40);
+            let src: Vec<u8> = (0..len).map(|_| rng.gen()).collect();
+            assert_eq!(reference_decode(&src), decode(&src), "src={:?}", src);
+        }
+
+        // Bias toward high bytes to exercise long codes and EOS prefixes.
+        for _ in 0..10_000 {
+            let len = rng.gen_range(0..40);
+            let src: Vec<u8> = (0..len).map(|_| rng.gen::<u8>() | 0xe0).collect();
+            assert_eq!(reference_decode(&src), decode(&src), "src={:?}", src);
+        }
+    }
+
+    #[test]
+    fn rejects_eos_and_invalid_padding() {
+        assert_eq!(decode(&[0xff]), Err(DecoderError::InvalidHuffmanCode));
+        assert_eq!(
+            decode(&[0xff, 0xff, 0xff, 0xff]),
+            Err(DecoderError::InvalidHuffmanCode)
+        );
+        assert_eq!(decode(&[0]), Err(DecoderError::InvalidHuffmanCode));
+    }
+}

src/hpack/huffman/encode.rs

@@ -0,0 +1,49 @@
+use crate::hpack::huffman::table::ENCODE_TABLE;
+
+use bytes::{BufMut, BytesMut};
+
+pub fn encode(src: &[u8], dst: &mut BytesMut) {
+    let mut bits: u64 = 0;
+    let mut bits_left = 40;
+
+    for &b in src {
+        let (nbits, code) = ENCODE_TABLE[b as usize];
+
+        bits |= code << (bits_left - nbits);
+        bits_left -= nbits;
+
+        while bits_left <= 32 {
+            dst.put_u8((bits >> 32) as u8);
+
+            bits <<= 8;
+            bits_left += 8;
+        }
+    }
+
+    if bits_left != 40 {
+        // This writes the EOS token
+        bits |= (1 << bits_left) - 1;
+        dst.put_u8((bits >> 32) as u8);
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn encode_single_byte() {
+        let mut dst = BytesMut::with_capacity(1);
+
+        encode(b"o", &mut dst);
+        assert_eq!(&dst[..], &[0b00111111]);
+
+        dst.clear();
+        encode(b"0", &mut dst);
+        assert_eq!(&dst[..], &[7]);
+
+        dst.clear();
+        encode(b"A", &mut dst);
+        assert_eq!(&dst[..], &[(0x21 << 2) + 3]);
+    }
+}