Merge pull request #61 from lilith/fixed-array-table

perf: fixed-size decode table — zero bounds checks, +3-9% throughput

Author: 197g

src/decode.rs

@@ -82,7 +82,7 @@ trait Stateful {
     fn reset(&mut self);
 }
 
-#[derive(Clone)]
+#[derive(Clone, Copy, Default)]
 struct Link {
     prev: Code,
     byte: u8,
@@ -178,9 +178,16 @@ struct Buffer {
     write_mark: usize,
 }
 
+/// Mask for indexing into fixed-size arrays. Since MAX_ENTRIES = 4096 = 2^12,
+/// `idx & MASK` is guaranteed < MAX_ENTRIES. LLVM can prove this for
+/// `[T; MAX_ENTRIES]` arrays, eliminating bounds checks. Corrupt `prev`
+/// values wrap to a valid index instead of panicking.
+const MASK: usize = MAX_ENTRIES - 1;
+
 struct Table {
-    inner: Vec<Link>,
-    depths: Vec<u16>,
+    inner: Box<[Link; MAX_ENTRIES]>,
+    depths: Box<[u16; MAX_ENTRIES]>,
+    len: usize,
 }
 
 /// Describes the static parameters for creating a decoder.
@@ -908,7 +915,7 @@ impl<C: CodeBuffer, CgC: CodegenConstants> Stateful for DecodeState<C, CgC> {
 
                 debug_assert!(
                     // When the table is full, we have a max code above the size switch.
-                    self.table.inner.len() >= MAX_ENTRIES - usize::from(self.is_tiff)
+                    self.table.len >= MAX_ENTRIES - usize::from(self.is_tiff)
                     // When the code size is 2 we have a bit code: (0, 1, CLS, EOF). Then the
                     // computed next_code is 4 which already exceeds the bit width from the start.
                     // Then we will immediately switch code size after this code.
@@ -966,7 +973,7 @@ impl<C: CodeBuffer, CgC: CodegenConstants> Stateful for DecodeState<C, CgC> {
                     }
 
                     // Read the code length and check that we can decode directly into the out slice.
-                    let len = self.table.depths[usize::from(read_code)];
+                    let len = self.table.depths[usize::from(read_code) & MASK];
 
                     if out.len() < usize::from(len) {
                         break;
@@ -1031,9 +1038,9 @@ impl<C: CodeBuffer, CgC: CodegenConstants> Stateful for DecodeState<C, CgC> {
                 }
 
                 let required_len = if new_code == self.next_code {
-                    self.table.depths[usize::from(deriv.code)] + 1
+                    self.table.depths[usize::from(deriv.code) & MASK] + 1
                 } else {
-                    self.table.depths[usize::from(new_code)]
+                    self.table.depths[usize::from(new_code) & MASK]
                 };
 
                 // We need the decoded data of the new code if it is the `next_code`. This is the
@@ -1324,7 +1331,7 @@ impl CodeBuffer for LsbBuffer {
                 new_bits as u8
             }
         };
-        self.bit_buffer |= u64::from_be_bytes(buffer).swap_bytes() << self.bits;
+        self.bit_buffer |= u64::from_le_bytes(buffer) << self.bits;
         self.bits += new_bits;
     }
 
@@ -1397,7 +1404,7 @@ impl Buffer {
     fn fill_reconstruct(&mut self, table: &Table, code: Code) -> u8 {
         self.write_mark = 0;
         self.read_mark = 0;
-        let depth = table.depths[usize::from(code)];
+        let depth = table.depths[usize::from(code) & MASK];
         let mut memory = core::mem::replace(&mut self.bytes, Box::default());
 
         let out = &mut memory[..usize::from(depth)];
@@ -1417,89 +1424,88 @@ impl Buffer {
     }
 }
 
+fn boxed_arr<T: Clone + Default, const N: usize>() -> Box<[T; N]> {
+    use core::convert::TryInto;
+    vec![T::default(); N]
+        .into_boxed_slice()
+        .try_into()
+        .ok()
+        .unwrap()
+}
+
 impl Table {
     fn new() -> Self {
         Table {
-            inner: Vec::with_capacity(MAX_ENTRIES),
-            depths: Vec::with_capacity(MAX_ENTRIES),
+            inner: boxed_arr(),
+            depths: boxed_arr(),
+            len: 0,
         }
     }
 
     fn clear(&mut self, min_size: u8) {
-        let static_count = usize::from(1u16 << u16::from(min_size)) + 2;
-        self.inner.truncate(static_count);
-        self.depths.truncate(static_count);
+        self.len = usize::from(1u16 << u16::from(min_size)) + 2;
     }
 
     fn init(&mut self, min_size: u8) {
-        self.inner.clear();
-        self.depths.clear();
+        self.len = 0;
         for i in 0..(1u16 << u16::from(min_size)) {
-            self.inner.push(Link::base(i as u8));
-            self.depths.push(1);
+            let idx = self.len & MASK;
+            self.inner[idx] = Link::base(i as u8);
+            self.depths[idx] = 1;
+            self.len += 1;
+        }
+        // Clear code + End code: skip writing when the masked index would
+        // alias an alphabet entry (happens at min_size=12 where clear=4096
+        // wraps to index 0).
+        for _ in 0..2 {
+            if self.len < MAX_ENTRIES {
+                let idx = self.len & MASK;
+                self.inner[idx] = Link::base(0);
+                self.depths[idx] = 0;
+            }
+            self.len += 1;
         }
-        // Clear code.
-        self.inner.push(Link::base(0));
-        self.depths.push(0);
-        // End code.
-        self.inner.push(Link::base(0));
-        self.depths.push(0);
     }
 
     fn at(&self, code: Code) -> &Link {
-        &self.inner[usize::from(code)]
+        &self.inner[usize::from(code) & MASK]
     }
 
     fn is_empty(&self) -> bool {
-        self.inner.is_empty()
+        self.len == 0
     }
 
     fn is_full(&self) -> bool {
-        self.inner.len() >= MAX_ENTRIES
+        self.len >= MAX_ENTRIES
     }
 
     fn derive(&mut self, from: &DerivationBase, byte: u8) {
-        let link = from.derive(byte);
-        let depth = self.depths[usize::from(from.code)] + 1;
-        self.inner.push(link);
-        self.depths.push(depth);
+        let idx = self.len & MASK;
+        let depth = self.depths[usize::from(from.code) & MASK] + 1;
+        self.inner[idx] = from.derive(byte);
+        self.depths[idx] = depth;
+        self.len += 1;
     }
 
-    // Derive multiple codes in a row, where each base is guaranteed to already exist.
     fn derive_burst(&mut self, from: &mut DerivationBase, burst: &[Code], first: &[u8]) {
-        let mut depth_of = from.code;
-        // Note that false data dependency we want to get rid of!
-        // TODO: this pushes into a Vec, maybe we can make this cleaner.
-        for &code in burst {
-            let depth = self.depths[usize::from(depth_of)] + 1;
-            self.depths.push(depth);
-            depth_of = code;
-        }
-
-        // Llvm tends to be flaky with code layout for the case of requiring an allocation. It's
-        // not clear if that can occur in practice but it relies on iterator size hint..
-        let extensions = burst.iter().zip(first);
-        self.inner.extend(extensions.map(|(&code, &first)| {
-            let link = from.derive(first);
+        for (&code, &first_byte) in burst.iter().zip(first.iter()) {
+            self.derive(from, first_byte);
             from.code = code;
-            from.first = first;
-            link
-        }));
+            from.first = first_byte;
+        }
     }
 
     fn reconstruct(&self, code: Code, out: &mut [u8]) -> u8 {
         let mut code_iter = code;
-        let table = &self.inner[..=usize::from(code)];
-        let first = table[usize::from(code)].first;
+        let first = self.inner[usize::from(code) & MASK].first;
 
-        let len = code_iter;
+        // The `& MASK` ensures any prev value (even from corrupt data) maps
+        // to a valid array index. This replaces the previous `min(len, prev)`
+        // clamp with equivalent safety: no panic, bounded loop, wrong output
+        // on corrupt data.
         for ch in out.iter_mut().rev() {
-            //(code, cha) = self.table[k as usize];
-            // Note: This could possibly be replaced with an unchecked array access if
-            //  - value is asserted to be < self.next_code() in push
-            //  - min_size is asserted to be < MAX_CODESIZE
-            let entry = &table[usize::from(code_iter)];
-            code_iter = core::cmp::min(len, entry.prev);
+            let entry = &self.inner[usize::from(code_iter) & MASK];
+            code_iter = entry.prev;
             *ch = entry.byte;
         }