decompress.rs: faster input reading

libbz2-rs-sys/src/bzlib.rs

@@ -244,6 +244,73 @@ mod stream {
             unsafe { Allocator::from_bz_stream(self) }
         }
 
+        /// Read up to 7 bytes into the bit buffer.
+        ///
+        /// The caller is responsible for updating `self.total_in`!
+        #[must_use]
+        #[inline(always)]
+        pub(crate) fn pull_u64(
+            &mut self,
+            mut bit_buffer: u64,
+            bits_used: i32,
+        ) -> Option<(u64, i32)> {
+            // we should only ask for more input if there are at least 8 free bits
+            debug_assert!(bits_used <= 56);
+
+            if self.avail_in < 8 {
+                return None;
+            }
+
+            // of course this uses big endian values
+            let read = u64::from_be_bytes(unsafe { self.next_in.cast::<[u8; 8]>().read() });
+
+            // because of the endianness, we can only shift in whole bytes.
+            // this calculates the number of available bits, rounded down to the nearest multiple
+            // of 8.
+            let increment_bits = (63 - bits_used) & !7;
+
+            // shift existing bits to the end, and or new bits in
+            bit_buffer = (bit_buffer << increment_bits) | (read >> (64 - increment_bits));
+
+            // we read 8 bytes above, but can only process `increment_bytes` worth of bits
+            let increment_bytes = increment_bits / 8;
+            self.next_in = unsafe { (self.next_in).add(increment_bytes as usize) };
+            self.avail_in -= increment_bytes as u32;
+
+            // skips updating `self.total_in`: the caller is responsible for keeping it updated
+
+            Some((bit_buffer, bits_used + increment_bits))
+        }
+
+        /// Read exactly 1 byte into the buffer
+        ///
+        /// The caller is responsible for updating `self.total_in`!
+        #[must_use]
+        #[inline(always)]
+        pub(crate) fn pull_u8(
+            &mut self,
+            mut bit_buffer: u64,
+            bits_used: i32,
+        ) -> Option<(u64, i32)> {
+            // we should only ask for more input if there are at least 8 free bits
+            debug_assert!(bits_used <= 56);
+
+            if self.avail_in == 0 || bits_used > 56 {
+                return None;
+            }
+
+            let read = unsafe { *(self.next_in as *mut u8) };
+            bit_buffer <<= 8;
+            bit_buffer |= u64::from(read);
+
+            self.next_in = unsafe { (self.next_in).offset(1) };
+            self.avail_in -= 1;
+
+            // skips updating `self.total_in`: the caller is responsible for keeping it updated
+
+            Some((bit_buffer, bits_used + 8))
+        }
+
         #[must_use]
         pub(crate) fn read_byte(&mut self) -> Option<u8> {
             if self.avail_in == 0 {
@@ -533,7 +600,7 @@ pub(crate) struct DState {
     pub k0: u8,
     pub rNToGo: i32,
     pub rTPos: i32,
-    pub bsBuff: u32,
+    pub bsBuff: u64,
     pub bsLive: i32,
     pub smallDecompress: DecompressMode,
     pub currBlockNo: i32,

libbz2-rs-sys/src/decompress.rs

@@ -134,6 +134,8 @@ pub(crate) fn decompress(
     let mut current_block: Block;
     let mut uc: u8;
 
+    let old_avail_in = strm.avail_in;
+
     if let State::BZ_X_MAGIC_1 = s.state {
         /*zero out the save area*/
         s.save = SaveArea::default();
@@ -180,14 +182,18 @@ pub(crate) fn decompress(
             ($strm:expr, $s:expr, $nnn:expr) => {
                 loop {
                     if $s.bsLive >= $nnn {
-                        let v: u32 = ($s.bsBuff >> ($s.bsLive - $nnn)) & ((1 << $nnn) - 1);
+                        let v: u64 = ($s.bsBuff >> ($s.bsLive - $nnn)) & ((1 << $nnn) - 1);
                         $s.bsLive -= $nnn;
                         break v;
                     }
 
-                    if let Some(next_byte) = strm.read_byte() {
-                        $s.bsBuff = $s.bsBuff << 8 | next_byte as u32;
-                        $s.bsLive += 8;
+                    if let Some((bit_buffer, bits_used)) = strm.pull_u64($s.bsBuff, $s.bsLive) {
+                        $s.bsBuff = bit_buffer;
+                        $s.bsLive = bits_used;
+                    } else if let Some((bit_buffer, bits_used)) = strm.pull_u8($s.bsBuff, $s.bsLive)
+                    {
+                        $s.bsBuff = bit_buffer;
+                        $s.bsLive = bits_used;
                     } else {
                         break 'save_state_and_return ReturnCode::BZ_OK;
                     }
@@ -1165,6 +1171,12 @@ pub(crate) fn decompress(
         gMinlen,
     };
 
+    // update total_in with how many bytes were read during this call
+    let bytes_read = old_avail_in - strm.avail_in;
+    let old_total_in_lo32 = strm.total_in_lo32;
+    strm.total_in_lo32 = strm.total_in_lo32.wrapping_add(bytes_read);
+    strm.total_in_hi32 += (strm.total_in_lo32 < old_total_in_lo32) as u32;
+
     ret_val
 }