Add WIMLIB_OPEN_FLAG_WRITE_ACCESS flag
[wimlib] / 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 (likely(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 (unlikely(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 (likely(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 (unlikely(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 inline 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 (likely(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 (likely(sym < num_syms)) {
216                         wimlib_assert2(lens[sym] <= table_bits);
217                         bitstream_remove_bits(istream, lens[sym]);
218                 } else {
219                         bitstream_remove_bits(istream, table_bits);
220                         do {
221                                 key_bits = sym + bitstream_peek_bits(istream, 1);
222                                 bitstream_remove_bits(istream, 1);
223
224                                 wimlib_assert2(key_bits < num_syms * 2 +
225                                                (1 << table_bits));
226                         } while ((sym = decode_table[key_bits]) >= num_syms);
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 /* Minimum alignment for the decode_table parameter to
246  * make_huffman_decode_table(). */
247 #define DECODE_TABLE_ALIGNMENT 16
248
249 #endif /* _WIMLIB_DECOMPRESS_H */