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 provide a bitstream. */
  17 struct input_bitstream {
  18
  19         /* A variable of length at least 32 bits that is used to hold bits that
  20          * have been read from the stream.  The bits are ordered from high-order
  21          * to low-order; the next bit is always the high-order bit. */
  22         input_bitbuf_t  bitbuf;
  23
  24         /* Pointer to the next byte to be retrieved from the input. */
  25         const u8 *data;
  26
  27         /* Number of bits in @bitbuf that are valid. */
  28         unsigned bitsleft;
  29
  30         /* Number of words of data that are left. */
  31         unsigned data_bytes_left;
  32 };
  33
  34 /* Initializes a bitstream to receive its input from @data. */
  35 static inline void init_input_bitstream(struct input_bitstream *istream,
  36                                         const void *data, unsigned num_data_bytes)
  37 {
  38         istream->bitbuf          = 0;
  39         istream->bitsleft        = 0;
  40         istream->data            = data;
  41         istream->data_bytes_left = num_data_bytes;
  42 }
  43
  44 /* Ensures that the bit buffer contains @num_bits bits. */
  45 static inline int bitstream_ensure_bits(struct input_bitstream *istream,
  46                                         unsigned num_bits)
  47 {
  48         wimlib_assert2(num_bits <= 16);
  49
  50         int ret = 0;
  51
  52         /* Unfortunately this needs to be different for the different
  53          * compression types.  LZX requires reading no more than the number of
  54          * bits needed, otherwise the end of the compressed data may be overrun.
  55          * XPRESS, on the other hand, requires that we always return with at
  56          * least 16 bits in the buffer, even if fewer are requested.  This is
  57          * important because this may change the location of a literal byte
  58          * read with bitstream_read_byte(). */
  59 #ifdef XPRESS_DECOMP
  60         if (istream->bitsleft < 16) {
  61 #else
  62         if (istream->bitsleft < num_bits) {
  63 #endif
  64                 if (istream->data_bytes_left >= 2) {
  65                         unsigned shift = sizeof(input_bitbuf_t) * 8 - 16 -
  66                                          istream->bitsleft;
  67                         istream->bitbuf |= (input_bitbuf_t)le16_to_cpu(
  68                                                 *(u16*)istream->data) << shift;
  69                         istream->data += 2;
  70                         istream->bitsleft += 16;
  71                         istream->data_bytes_left -= 2;
  72                 } else {
  73                         ret = 1;
  74                 }
  75         }
  76         wimlib_assert2(ret != 0 || istream->bitsleft >= num_bits);
  77         return ret;
  78 }
  79
  80 /* Returns the next @num_bits bits in the bit buffer.  It must contain at least
  81  * @num_bits bits to call this function. */
  82 static inline unsigned
  83 bitstream_peek_bits(const struct input_bitstream *istream, unsigned num_bits)
  84 {
  85         wimlib_assert2(istream->bitsleft >= num_bits);
  86         int ret;
  87         if (num_bits == 0)
  88                 ret = 0;
  89         else
  90                 ret = istream->bitbuf >> (sizeof(input_bitbuf_t) * 8 - num_bits);
  91         return ret;
  92 }
  93
  94 /* Removes @num_bits bits from the bit buffer.  It must contain at least
  95  * @num_bits bits to call this function. */
  96 static inline void bitstream_remove_bits(struct input_bitstream *istream,
  97                                          unsigned num_bits)
  98 {
  99         wimlib_assert2(istream->bitsleft >= num_bits);
 100         istream->bitbuf <<= num_bits;
 101         istream->bitsleft -= num_bits;
 102 }
 103
 104 /* Reads and returns @num_bits bits from the input bitstream. */
 105 static inline int bitstream_read_bits(struct input_bitstream *istream,
 106                                       unsigned num_bits, unsigned *n)
 107 {
 108         int ret = bitstream_ensure_bits(istream, num_bits);
 109         if (ret == 0) {
 110                 *n = bitstream_peek_bits(istream, num_bits);
 111                 bitstream_remove_bits(istream, num_bits);
 112         } else {
 113                 ERROR("bitstream_read_bits(): Input buffer exhausted");
 114         }
 115         return ret;
 116 }
 117
 118 /* In the XPRESS format there can be literal length bytes embedded in the
 119  * compressed bitstream.  These bytes are basically separate from the bitstream,
 120  * as they come AFTER the bits that are currently in the buffer variable (based
 121  * on reading 16 bits at a time), even though the buffer variable may not be
 122  * empty.
 123  *
 124  * This function returns the next such literal length byte in the input
 125  * bitstream.  Returns -1 if we are at the end of the bitstream. */
 126 static inline int bitstream_read_byte(struct input_bitstream *istream)
 127 {
 128         wimlib_assert2(istream->bitsleft < 32);
 129         int ret;
 130
 131         if (istream->data_bytes_left == 0) {
 132                 ERROR("bitstream_read_byte(): Input buffer exhausted");
 133                 ret = -1;
 134         } else {
 135                 istream->data_bytes_left--;
 136                 ret = *istream->data++;
 137         }
 138         return ret;
 139 }
 140
 141 /* Reads @num_bits bits from the bit buffer without checking to see if that many
 142  * bits are in the buffer or not. */
 143 static inline unsigned
 144 bitstream_read_bits_nocheck(struct input_bitstream *istream, unsigned num_bits)
 145 {
 146         unsigned n = bitstream_peek_bits(istream, num_bits);
 147         bitstream_remove_bits(istream, num_bits);
 148         return n;
 149 }
 150
 151 /* Removes the bits that have been read into the bit buffer variable. */
 152 static inline void flush_input_bitstream(struct input_bitstream *istream)
 153 {
 154         bitstream_remove_bits(istream, istream->bitsleft);
 155         istream->bitbuf = 0;
 156         wimlib_assert2(istream->bitsleft == 0);
 157 }
 158
 159 extern int bitstream_read_bytes(struct input_bitstream *istream, size_t n,
 160                                 void *dest);
 161
 162 /* Aligns the bitstream on a 16-bit boundary. */
 163 static inline void align_input_bitstream(struct input_bitstream *istream)
 164 {
 165         bitstream_remove_bits(istream, istream->bitsleft & 15);
 166 }
 167
 168 extern int read_huffsym_near_end_of_input(struct input_bitstream *istream,
 169                                           const u16 decode_table[],
 170                                           const u8 lens[],
 171                                           unsigned num_syms,
 172                                           unsigned table_bits,
 173                                           unsigned *n);
 174
 175 /*
 176  * Reads a Huffman-encoded symbol from a bitstream.
 177  *
 178  * This function may be called hundreds of millions of times when extracting a
 179  * large WIM file.  I'm not sure it could be made much faster that it is,
 180  * especially since there isn't enough time to make a big table that allows
 181  * decoding multiple symbols per lookup.  But if extracting files to a hard
 182  * disk, the IO will be the bottleneck anyway.
 183  *
 184  * @buf:        The input buffer from which the symbol will be read.
 185  * @decode_table:       The fast Huffman decoding table for the Huffman tree.
 186  * @lengths:            The table that gives the length of the code for each
 187  *                              symbol.
 188  * @num_symbols:        The number of symbols in the Huffman code.
 189  * @table_bits:         Huffman codes this length or less can be looked up
 190  *                              directory in the decode_table, as the
 191  *                              decode_table contains 2**table_bits entries.
 192  */
 193 static inline int 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 make_huffman_decode_table(u16 decode_table[], unsigned num_syms,
 241                                      unsigned num_bits, const u8 lengths[],
 242                                      unsigned max_codeword_len);
 243
 244 #endif /* _WIMLIB_DECOMP_H */