wimlib.net Git - wimlib/blob - src/decompress.h

   1 /*
   2  * decompress.h
   3  *
   4  * Functions useful for decompression, mainly bitstreams.
   5  */
   6
   7 #ifndef _WIMLIB_DECOMP_H
   8 #define _WIMLIB_DECOMP_H
   9
  10 #include "util.h"
  11 #include "endianness.h"
  12
  13 /* Must be at least 32 bits. */
  14 typedef unsigned long input_bitbuf_t;
  15
  16 /* Structure to encapsulate a block of in-memory data that is being interpreted
  17  * as a stream of bits.
  18  *
  19  * This is geared specifically towards the XPRESS and LZX compression formats
  20  * with regards to the actual ordering the bits within the byte sequence. */
  21 struct input_bitstream {
  22
  23         /* A variable of length at least 32 bits that is used to hold bits that
  24          * have been read from the stream.  The bits are ordered from high-order
  25          * to low-order, and the next bit is always the high-order bit. */
  26         input_bitbuf_t  bitbuf;
  27
  28         /* Pointer to the next byte to be retrieved from the input. */
  29         const u8 *data;
  30
  31         /* Number of bits in @bitbuf that are valid. */
  32         unsigned bitsleft;
  33
  34         /* Number of words of data that are left. */
  35         unsigned data_bytes_left;
  36 };
  37
  38 /* Initializes a bitstream to receive its input from @data. */
  39 static inline void init_input_bitstream(struct input_bitstream *istream,
  40                                         const void *data, unsigned num_data_bytes)
  41 {
  42         istream->bitbuf          = 0;
  43         istream->bitsleft        = 0;
  44         istream->data            = data;
  45         istream->data_bytes_left = num_data_bytes;
  46 }
  47
  48 /* Ensures that the bit buffer variable for the bitstream contains @num_bits
  49  * bits.
  50  *
  51  * If there are at least @num_bits bits remaining in the bitstream, 0 is
  52  * returned.  Otherwise, -1 is returned.  */
  53 static inline int bitstream_ensure_bits(struct input_bitstream *istream,
  54                                         unsigned num_bits)
  55 {
  56         wimlib_assert2(num_bits <= 16);
  57
  58         int ret = 0;
  59
  60         /* Unfortunately this needs to be different for the different
  61          * compression types.  LZX requires reading no more than the number of
  62          * bits needed, otherwise the end of the compressed data may be overrun.
  63          * XPRESS, on the other hand, requires that we always return with at
  64          * least 16 bits in the buffer, even if fewer are requested.  This is
  65          * important because this may change the location of a literal byte
  66          * read with bitstream_read_byte(). */
  67 #ifdef XPRESS_DECOMP
  68         if (istream->bitsleft < 16) {
  69 #else
  70         if (istream->bitsleft < num_bits) {
  71 #endif
  72                 if (istream->data_bytes_left >= 2) {
  73                         unsigned shift = sizeof(input_bitbuf_t) * 8 - 16 -
  74                                          istream->bitsleft;
  75                         istream->bitbuf |= (input_bitbuf_t)le16_to_cpu(
  76                                                 *(u16*)istream->data) << shift;
  77                         istream->data += 2;
  78                         istream->bitsleft += 16;
  79                         istream->data_bytes_left -= 2;
  80                 } else {
  81                         ret = -1;
  82                 }
  83         }
  84         wimlib_assert2(ret != 0 || istream->bitsleft >= num_bits);
  85         return ret;
  86 }
  87
  88 /* Returns the next @num_bits bits in the buffer variable, which must contain at
  89  * least @num_bits bits, for the bitstream. */
  90 static inline unsigned
  91 bitstream_peek_bits(const struct input_bitstream *istream, unsigned num_bits)
  92 {
  93         wimlib_assert2(istream->bitsleft >= num_bits);
  94         int ret;
  95         if (num_bits == 0)
  96                 ret = 0;
  97         else
  98                 ret = istream->bitbuf >> (sizeof(input_bitbuf_t) * 8 - num_bits);
  99         return ret;
 100 }
 101
 102 /* Removes @num_bits bits from the buffer variable, which must contain at least
 103  * @num_bits bits, for the bitstream. */
 104 static inline void bitstream_remove_bits(struct input_bitstream *istream,
 105                                          unsigned num_bits)
 106 {
 107         wimlib_assert2(istream->bitsleft >= num_bits);
 108         istream->bitbuf <<= num_bits;
 109         istream->bitsleft -= num_bits;
 110 }
 111
 112 /* Reads @num_bits bits from the input bitstream.  @num_bits must be 16 or fewer.
 113  * On success, returns 0 and returns the requested bits in @n.  If there are
 114  * fewer than @num_bits remaining in the bitstream, -1 is returned. */
 115 static inline int bitstream_read_bits(struct input_bitstream *istream,
 116                                       unsigned num_bits, unsigned *n)
 117 {
 118         wimlib_assert2(num_bits <= 16);
 119         int ret = bitstream_ensure_bits(istream, num_bits);
 120         if (ret == 0) {
 121                 *n = bitstream_peek_bits(istream, num_bits);
 122                 bitstream_remove_bits(istream, num_bits);
 123                 wimlib_assert2(istream->bitsleft < 16);
 124         } else {
 125                 ERROR("bitstream_read_bits(): Input buffer exhausted");
 126         }
 127         return ret;
 128 }
 129
 130 /* In the XPRESS format there can be literal bytes embedded in the bitstream.
 131  * These bytes are basically separate from the bitstream, as they come AFTER the
 132  * bits that are currently in the buffer variable (based on reading 16 bits at a
 133  * time), even though the buffer variable may not be empty.
 134  *
 135  * This function returns the next such literal byte, or -1 if there are no more.
 136  */
 137 static inline int bitstream_read_byte(struct input_bitstream *istream)
 138 {
 139         wimlib_assert2(istream->bitsleft < 32);
 140         int ret;
 141
 142         if (istream->data_bytes_left == 0) {
 143                 ERROR("bitstream_read_byte(): Input buffer exhausted");
 144                 ret = -1;
 145         } else {
 146                 istream->data_bytes_left--;
 147                 ret = *istream->data++;
 148         }
 149         return ret;
 150 }
 151
 152 /* Reads @num_bits bits from the buffer variable for a bistream without checking
 153  * to see if that many bits are in the buffer or not. */
 154 static inline unsigned
 155 bitstream_read_bits_nocheck(struct input_bitstream *istream, unsigned num_bits)
 156 {
 157         unsigned n = bitstream_peek_bits(istream, num_bits);
 158         bitstream_remove_bits(istream, num_bits);
 159         return n;
 160 }
 161
 162 extern int read_huffsym_near_end_of_input(struct input_bitstream *istream,
 163                                           const u16 decode_table[],
 164                                           const u8 lens[],
 165                                           unsigned num_syms,
 166                                           unsigned table_bits,
 167                                           unsigned *n);
 168
 169 /*
 170  * Reads a Huffman-encoded symbol from a bitstream.
 171  *
 172  * This function may be called hundreds of millions of times when extracting a
 173  * large WIM file.  I'm not sure it could be made much faster that it is,
 174  * especially since there isn't enough time to make a big table that allows
 175  * decoding multiple symbols per lookup.  But if extracting files to a hard
 176  * disk, the I/O will be the bottleneck anyway.
 177  *
 178  * @buf:        The input buffer from which the symbol will be read.
 179  * @decode_table:       The fast Huffman decoding table for the Huffman tree.
 180  * @lengths:            The table that gives the length of the code for each
 181  *                              symbol.
 182  * @num_symbols:        The number of symbols in the Huffman code.
 183  * @table_bits:         Huffman codes this length or less can be looked up
 184  *                              directory in the decode_table, as the
 185  *                              decode_table contains 2**table_bits entries.
 186  */
 187 static inline int read_huffsym(struct input_bitstream *istream,
 188                                const u16 decode_table[],
 189                                const u8 lens[],
 190                                unsigned num_syms,
 191                                unsigned table_bits,
 192                                unsigned *n,
 193                                unsigned max_codeword_len)
 194 {
 195         int ret;
 196
 197         /* In the most common case, there are at least max_codeword_len bits
 198          * remaining in the stream. */
 199         if (bitstream_ensure_bits(istream, max_codeword_len) == 0) {
 200
 201                 /* Use the next table_bits of the input as an index into the
 202                  * decode_table. */
 203                 u16 key_bits = bitstream_peek_bits(istream, table_bits);
 204
 205                 u16 sym = decode_table[key_bits];
 206
 207                 /* If the entry in the decode table is not a valid symbol, it is
 208                  * the offset of the root of its Huffman subtree. */
 209                 if (sym >= num_syms) {
 210                         bitstream_remove_bits(istream, table_bits);
 211                         do {
 212                                 key_bits = sym + bitstream_peek_bits(istream, 1);
 213                                 bitstream_remove_bits(istream, 1);
 214
 215                                 wimlib_assert2(key_bits < num_syms * 2 +
 216                                                (1 << table_bits));
 217                         } while ((sym = decode_table[key_bits]) >= num_syms);
 218                 } else {
 219                         wimlib_assert2(lens[sym] <= table_bits);
 220                         bitstream_remove_bits(istream, lens[sym]);
 221                 }
 222                 *n = sym;
 223                 ret = 0;
 224         } else {
 225                 /* Otherwise, we must be careful to use only the bits that are
 226                  * actually remaining.  */
 227                 ret = read_huffsym_near_end_of_input(istream, decode_table,
 228                                                      lens, num_syms,
 229                                                      table_bits, n);
 230         }
 231         return ret;
 232 }
 233
 234 extern int make_huffman_decode_table(u16 decode_table[], unsigned num_syms,
 235                                      unsigned num_bits, const u8 lengths[],
 236                                      unsigned max_codeword_len);
 237
 238 #endif /* _WIMLIB_DECOMP_H */