Decompression optimizations
[wimlib] / src / decompress.h
index b4c4579543bbe59b6b57f58704689fc28216582d..ae1c99becda88f81e77146d3957683d65743b575 100644 (file)
@@ -45,7 +45,9 @@ static inline void init_input_bitstream(struct input_bitstream *istream,
 static inline int bitstream_ensure_bits(struct input_bitstream *istream,
                                        unsigned num_bits)
 {
-       wimlib_assert(num_bits <= 16);
+       wimlib_assert2(num_bits <= 16);
+
+       int ret = 0;
 
        /* Unfortunately this needs to be different for the different
         * compression types.  LZX requires reading no more than the number of
@@ -55,22 +57,24 @@ static inline int bitstream_ensure_bits(struct input_bitstream *istream,
         * important because this may change the location of a literal byte
         * read with bitstream_read_byte(). */
 #ifdef XPRESS_DECOMP
-       while (istream->bitsleft < 16) {
+       if (istream->bitsleft < 16) {
 #else
-       while (istream->bitsleft < num_bits) {
+       if (istream->bitsleft < num_bits) {
 #endif
-               if (istream->data_bytes_left < 2)
-                       return 1;
-
-               unsigned shift = sizeof(input_bitbuf_t) * 8 - 16 -
-                                istream->bitsleft;
-               istream->bitbuf |= (input_bitbuf_t)le16_to_cpu(
-                                       *(u16*)istream->data) << shift;
-               istream->data += 2;
-               istream->bitsleft += 16;
-               istream->data_bytes_left -= 2;
+               if (istream->data_bytes_left >= 2) {
+                       unsigned shift = sizeof(input_bitbuf_t) * 8 - 16 -
+                                        istream->bitsleft;
+                       istream->bitbuf |= (input_bitbuf_t)le16_to_cpu(
+                                               *(u16*)istream->data) << shift;
+                       istream->data += 2;
+                       istream->bitsleft += 16;
+                       istream->data_bytes_left -= 2;
+               } else {
+                       ret = 1;
+               }
        }
-       return 0;
+       wimlib_assert2(ret != 0 || istream->bitsleft >= num_bits);
+       return ret;
 }
 
 /* Returns the next @num_bits bits in the bit buffer.  It must contain at least
@@ -78,9 +82,13 @@ static inline int bitstream_ensure_bits(struct input_bitstream *istream,
 static inline unsigned
 bitstream_peek_bits(const struct input_bitstream *istream, unsigned num_bits)
 {
+       wimlib_assert2(istream->bitsleft >= num_bits);
+       int ret;
        if (num_bits == 0)
-               return 0;
-       return istream->bitbuf >> (sizeof(input_bitbuf_t) * 8 - num_bits);
+               ret = 0;
+       else
+               ret = istream->bitbuf >> (sizeof(input_bitbuf_t) * 8 - num_bits);
+       return ret;
 }
 
 /* Removes @num_bits bits from the bit buffer.  It must contain at least
@@ -88,6 +96,7 @@ bitstream_peek_bits(const struct input_bitstream *istream, unsigned num_bits)
 static inline void bitstream_remove_bits(struct input_bitstream *istream,
                                         unsigned num_bits)
 {
+       wimlib_assert2(istream->bitsleft >= num_bits);
        istream->bitbuf <<= num_bits;
        istream->bitsleft -= num_bits;
 }
@@ -96,15 +105,14 @@ static inline void bitstream_remove_bits(struct input_bitstream *istream,
 static inline int bitstream_read_bits(struct input_bitstream *istream,
                                      unsigned num_bits, unsigned *n)
 {
-       int ret;
-       ret = bitstream_ensure_bits(istream, num_bits);
-       if (ret != 0) {
+       int ret = bitstream_ensure_bits(istream, num_bits);
+       if (ret == 0) {
+               *n = bitstream_peek_bits(istream, num_bits);
+               bitstream_remove_bits(istream, num_bits);
+       } else {
                ERROR("bitstream_read_bits(): Input buffer exhausted");
-               return ret;
        }
-       *n = bitstream_peek_bits(istream, num_bits);
-       bitstream_remove_bits(istream, num_bits);
-       return 0;
+       return ret;
 }
 
 /* In the XPRESS format there can be literal length bytes embedded in the
@@ -117,14 +125,17 @@ static inline int bitstream_read_bits(struct input_bitstream *istream,
  * bitstream.  Returns -1 if we are at the end of the bitstream. */
 static inline int bitstream_read_byte(struct input_bitstream *istream)
 {
-       wimlib_assert(istream->bitsleft < 32);
+       wimlib_assert2(istream->bitsleft < 32);
+       int ret;
 
        if (istream->data_bytes_left == 0) {
                ERROR("bitstream_read_byte(): Input buffer exhausted");
-               return -1;
+               ret = -1;
+       } else {
+               istream->data_bytes_left--;
+               ret = *istream->data++;
        }
-       istream->data_bytes_left--;
-       return *istream->data++;
+       return ret;
 }
 
 /* Reads @num_bits bits from the bit buffer without checking to see if that many
@@ -141,23 +152,90 @@ bitstream_read_bits_nocheck(struct input_bitstream *istream, unsigned num_bits)
 static inline void flush_input_bitstream(struct input_bitstream *istream)
 {
        bitstream_remove_bits(istream, istream->bitsleft);
-       istream->bitsleft = 0;
        istream->bitbuf = 0;
+       wimlib_assert2(istream->bitsleft == 0);
 }
 
 extern int bitstream_read_bytes(struct input_bitstream *istream, size_t n,
                                void *dest);
 
-extern int align_input_bitstream(struct input_bitstream *istream,
-                                bool skip_word_if_aligned);
+/* Aligns the bitstream on a 16-bit boundary. */
+static inline void align_input_bitstream(struct input_bitstream *istream)
+{
+       bitstream_remove_bits(istream, istream->bitsleft & 15);
+}
 
-extern int read_huffsym(struct input_bitstream *stream,
-                       const u16 decode_table[],
-                       const u8 lengths[],
-                       unsigned num_symbols,
-                       unsigned table_bits,
-                       unsigned *n,
-                       unsigned max_codeword_len);
+extern int read_huffsym_near_end_of_input(struct input_bitstream *istream,
+                                         const u16 decode_table[],
+                                         const u8 lens[],
+                                         unsigned num_syms,
+                                         unsigned table_bits,
+                                         unsigned *n);
+
+/*
+ * Reads a Huffman-encoded symbol from a bitstream.
+ *
+ * This function may be called hundreds of millions of times when extracting a
+ * large WIM file.  I'm not sure it could be made much faster that it is,
+ * especially since there isn't enough time to make a big table that allows
+ * decoding multiple symbols per lookup.  But if extracting files to a hard
+ * disk, the IO will be the bottleneck anyway.
+ *
+ * @buf:       The input buffer from which the symbol will be read.
+ * @decode_table:      The fast Huffman decoding table for the Huffman tree.
+ * @lengths:           The table that gives the length of the code for each
+ *                             symbol.
+ * @num_symbols:       The number of symbols in the Huffman code.
+ * @table_bits:                Huffman codes this length or less can be looked up
+ *                             directory in the decode_table, as the
+ *                             decode_table contains 2**table_bits entries.
+ */
+static inline int read_huffsym(struct input_bitstream *istream,
+                              const u16 decode_table[],
+                              const u8 lens[],
+                              unsigned num_syms,
+                              unsigned table_bits,
+                              unsigned *n,
+                              unsigned max_codeword_len)
+{
+       int ret;
+
+       /* In the most common case, there are at least max_codeword_len bits
+        * remaining in the stream. */
+       if (bitstream_ensure_bits(istream, max_codeword_len) == 0) {
+
+               /* Use the next table_bits of the input as an index into the
+                * decode_table. */
+               u16 key_bits = bitstream_peek_bits(istream, table_bits);
+
+               u16 sym = decode_table[key_bits];
+
+               /* If the entry in the decode table is not a valid symbol, it is
+                * the offset of the root of its Huffman subtree. */
+               if (sym >= num_syms) {
+                       bitstream_remove_bits(istream, table_bits);
+                       do {
+                               key_bits = sym + bitstream_peek_bits(istream, 1);
+                               bitstream_remove_bits(istream, 1);
+
+                               wimlib_assert2(key_bits < num_syms * 2 +
+                                              (1 << table_bits));
+                       } while ((sym = decode_table[key_bits]) >= num_syms);
+               } else {
+                       wimlib_assert2(lens[sym] <= table_bits);
+                       bitstream_remove_bits(istream, lens[sym]);
+               }
+               *n = sym;
+               ret = 0;
+       } else {
+               /* Otherwise, we must be careful to use only the bits that are
+                * actually remaining.  */
+               ret = read_huffsym_near_end_of_input(istream, decode_table,
+                                                    lens, num_syms,
+                                                    table_bits, n);
+       }
+       return ret;
+}
 
 extern int make_huffman_decode_table(u16 decode_table[], unsigned num_syms,
                                     unsigned num_bits, const u8 lengths[],