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

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