Adjust naming of (de)compression files
[wimlib] / src / lzms_decompress.c
1 /*
2  * lzms_decompress.c
3  *
4  * A decompressor for the LZMS compression format.
5  */
6
7 /*
8  * Copyright (C) 2013, 2014 Eric Biggers
9  *
10  * This file is free software; you can redistribute it and/or modify it under
11  * the terms of the GNU Lesser General Public License as published by the Free
12  * Software Foundation; either version 3 of the License, or (at your option) any
13  * later version.
14  *
15  * This file is distributed in the hope that it will be useful, but WITHOUT
16  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
17  * FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
18  * details.
19  *
20  * You should have received a copy of the GNU Lesser General Public License
21  * along with this file; if not, see http://www.gnu.org/licenses/.
22  */
23
24 /*
25  * This is a decompressor for the LZMS compression format used by Microsoft.
26  * This format is not documented, but it is one of the formats supported by the
27  * compression API available in Windows 8, and as of Windows 8 it is one of the
28  * formats that can be used in WIM files.
29  *
30  * This decompressor only implements "raw" decompression, which decompresses a
31  * single LZMS-compressed block.  This behavior is the same as that of
32  * Decompress() in the Windows 8 compression API when using a compression handle
33  * created with CreateDecompressor() with the Algorithm parameter specified as
34  * COMPRESS_ALGORITHM_LZMS | COMPRESS_RAW.  Presumably, non-raw LZMS data is a
35  * container format from which the locations and sizes (both compressed and
36  * uncompressed) of the constituent blocks can be determined.
37  *
38  * An LZMS-compressed block must be read in 16-bit little endian units from both
39  * directions.  One logical bitstream starts at the front of the block and
40  * proceeds forwards.  Another logical bitstream starts at the end of the block
41  * and proceeds backwards.  Bits read from the forwards bitstream constitute
42  * binary range-encoded data, whereas bits read from the backwards bitstream
43  * constitute Huffman-encoded symbols or verbatim bits.  For both bitstreams,
44  * the ordering of the bits within the 16-bit coding units is such that the
45  * first bit is the high-order bit and the last bit is the low-order bit.
46  *
47  * From these two logical bitstreams, an LZMS decompressor can reconstitute the
48  * series of items that make up the LZMS data representation.  Each such item
49  * may be a literal byte or a match.  Matches may be either traditional LZ77
50  * matches or "delta" matches, either of which can have its offset encoded
51  * explicitly or encoded via a reference to a recently used (repeat) offset.
52  *
53  * A traditional LZ match consists of a length and offset; it asserts that the
54  * sequence of bytes beginning at the current position and extending for the
55  * length is exactly equal to the equal-length sequence of bytes at the offset
56  * back in the data buffer.  On the other hand, a delta match consists of a
57  * length, raw offset, and power.  It asserts that the sequence of bytes
58  * beginning at the current position and extending for the length is equal to
59  * the bytewise sum of the two equal-length sequences of bytes (2**power) and
60  * (raw_offset * 2**power) bytes before the current position, minus bytewise the
61  * sequence of bytes beginning at (2**power + raw_offset * 2**power) bytes
62  * before the current position.  Although not generally as useful as traditional
63  * LZ matches, delta matches can be helpful on some types of data.  Both LZ and
64  * delta matches may overlap with the current position; in fact, the minimum
65  * offset is 1, regardless of match length.
66  *
67  * For LZ matches, up to 3 repeat offsets are allowed, similar to some other
68  * LZ-based formats such as LZX and LZMA.  They must updated in an LRU fashion,
69  * except for a quirk: inserting anything to the front of the queue must be
70  * delayed by one LZMS item.  The reason for this is presumably that there is
71  * almost no reason to code the same match offset twice in a row, since you
72  * might as well have coded a longer match at that offset.  For this same
73  * reason, it also is a requirement that when an offset in the queue is used,
74  * that offset is removed from the queue immediately (and made pending for
75  * front-insertion after the following decoded item), and everything to the
76  * right is shifted left one queue slot.  This creates a need for an "overflow"
77  * fourth entry in the queue, even though it is only possible to decode
78  * references to the first 3 entries at any given time.  The queue must be
79  * initialized to the offsets {1, 2, 3, 4}.
80  *
81  * Repeat delta matches are handled similarly, but for them there are two queues
82  * updated in lock-step: one for powers and one for raw offsets.  The power
83  * queue must be initialized to {0, 0, 0, 0}, and the raw offset queue must be
84  * initialized to {1, 2, 3, 4}.
85  *
86  * Bits from the binary range decoder must be used to disambiguate item types.
87  * The range decoder must hold two state variables: the range, which must
88  * initially be set to 0xffffffff, and the current code, which must initially be
89  * set to the first 32 bits read from the forwards bitstream.  The range must be
90  * maintained above 0xffff; when it falls below 0xffff, both the range and code
91  * must be left-shifted by 16 bits and the low 16 bits of the code must be
92  * filled in with the next 16 bits from the forwards bitstream.
93  *
94  * To decode each bit, the binary range decoder requires a probability that is
95  * logically a real number between 0 and 1.  Multiplying this probability by the
96  * current range and taking the floor gives the bound between the 0-bit region of
97  * the range and the 1-bit region of the range.  However, in LZMS, probabilities
98  * are restricted to values of n/64 where n is an integer is between 1 and 63
99  * inclusively, so the implementation may use integer operations instead.
100  * Following calculation of the bound, if the current code is in the 0-bit
101  * region, the new range becomes the current code and the decoded bit is 0;
102  * otherwise, the bound must be subtracted from both the range and the code, and
103  * the decoded bit is 1.  More information about range coding can be found at
104  * https://en.wikipedia.org/wiki/Range_encoding.  Furthermore, note that the
105  * LZMA format also uses range coding and has public domain code available for
106  * it.
107  *
108  * The probability used to range-decode each bit must be taken from a table, of
109  * which one instance must exist for each distinct context in which a
110  * range-decoded bit is needed.  At each call of the range decoder, the
111  * appropriate probability must be obtained by indexing the appropriate
112  * probability table with the last 4 (in the context disambiguating literals
113  * from matches), 5 (in the context disambiguating LZ matches from delta
114  * matches), or 6 (in all other contexts) bits recently range-decoded in that
115  * context, ordered such that the most recently decoded bit is the low-order bit
116  * of the index.
117  *
118  * Furthermore, each probability entry itself is variable, as its value must be
119  * maintained as n/64 where n is the number of 0 bits in the most recently
120  * decoded 64 bits with that same entry.  This allows the compressed
121  * representation to adapt to the input and use fewer bits to represent the most
122  * likely data; note that LZMA uses a similar scheme.  Initially, the most
123  * recently 64 decoded bits for each probability entry are assumed to be
124  * 0x0000000055555555 (high order to low order); therefore, all probabilities
125  * are initially 48/64.  During the course of decoding, each probability may be
126  * updated to as low as 0/64 (as a result of reading many consecutive 1 bits
127  * with that entry) or as high as 64/64 (as a result of reading many consecutive
128  * 0 bits with that entry); however, probabilities of 0/64 and 64/64 cannot be
129  * used as-is but rather must be adjusted to 1/64 and 63/64, respectively,
130  * before being used for range decoding.
131  *
132  * Representations of the LZMS items themselves must be read from the backwards
133  * bitstream.  For this, there are 5 different Huffman codes used:
134  *
135  *  - The literal code, used for decoding literal bytes.  Each of the 256
136  *    symbols represents a literal byte.  This code must be rebuilt whenever
137  *    1024 symbols have been decoded with it.
138  *
139  *  - The LZ offset code, used for decoding the offsets of standard LZ77
140  *    matches.  Each symbol represents an offset slot, which corresponds to a
141  *    base value and some number of extra bits which must be read and added to
142  *    the base value to reconstitute the full offset.  The number of symbols in
143  *    this code is the number of offset slots needed to represent all possible
144  *    offsets in the uncompressed block.  This code must be rebuilt whenever
145  *    1024 symbols have been decoded with it.
146  *
147  *  - The length code, used for decoding length symbols.  Each of the 54 symbols
148  *    represents a length slot, which corresponds to a base value and some
149  *    number of extra bits which must be read and added to the base value to
150  *    reconstitute the full length.  This code must be rebuilt whenever 512
151  *    symbols have been decoded with it.
152  *
153  *  - The delta offset code, used for decoding the offsets of delta matches.
154  *    Each symbol corresponds to an offset slot, which corresponds to a base
155  *    value and some number of extra bits which must be read and added to the
156  *    base value to reconstitute the full offset.  The number of symbols in this
157  *    code is equal to the number of symbols in the LZ offset code.  This code
158  *    must be rebuilt whenever 1024 symbols have been decoded with it.
159  *
160  *  - The delta power code, used for decoding the powers of delta matches.  Each
161  *    of the 8 symbols corresponds to a power.  This code must be rebuilt
162  *    whenever 512 symbols have been decoded with it.
163  *
164  * Initially, each Huffman code must be built assuming that each symbol in that
165  * code has frequency 1.  Following that, each code must be rebuilt each time a
166  * certain number of symbols, as noted above, has been decoded with it.  The
167  * symbol frequencies for a code must be halved after each rebuild of that code;
168  * this makes the codes adapt to the more recent data.
169  *
170  * Like other compression formats such as XPRESS, LZX, and DEFLATE, the LZMS
171  * format requires that all Huffman codes be constructed in canonical form.
172  * This form requires that same-length codewords be lexicographically ordered
173  * the same way as the corresponding symbols and that all shorter codewords
174  * lexicographically precede longer codewords.  Such a code can be constructed
175  * directly from codeword lengths.
176  *
177  * Even with the canonical code restriction, the same frequencies can be used to
178  * construct multiple valid Huffman codes.  Therefore, the decompressor needs to
179  * construct the right one.  Specifically, the LZMS format requires that the
180  * Huffman code be constructed as if the well-known priority queue algorithm is
181  * used and frequency ties are always broken in favor of leaf nodes.
182  *
183  * Codewords in LZMS are guaranteed to not exceed 15 bits.  The format otherwise
184  * places no restrictions on codeword length.  Therefore, the Huffman code
185  * construction algorithm that a correct LZMS decompressor uses need not
186  * implement length-limited code construction.  But if it does (e.g. by virtue
187  * of being shared among multiple compression algorithms), the details of how it
188  * does so are unimportant, provided that the maximum codeword length parameter
189  * is set to at least 15 bits.
190  *
191  * After all LZMS items have been decoded, the data must be postprocessed to
192  * translate absolute address encoded in x86 instructions into their original
193  * relative addresses.
194  *
195  * Details omitted above can be found in the code.  Note that in the absence of
196  * an official specification there is no guarantee that this decompressor
197  * handles all possible cases.
198  */
199
200 #ifdef HAVE_CONFIG_H
201 #  include "config.h"
202 #endif
203
204 #include "wimlib/compress_common.h"
205 #include "wimlib/decompressor_ops.h"
206 #include "wimlib/decompress_common.h"
207 #include "wimlib/error.h"
208 #include "wimlib/lzms_common.h"
209 #include "wimlib/util.h"
210
211 /* The TABLEBITS values can be changed; they only affect decoding speed.  */
212 #define LZMS_LITERAL_TABLEBITS          10
213 #define LZMS_LENGTH_TABLEBITS           10
214 #define LZMS_LZ_OFFSET_TABLEBITS        10
215 #define LZMS_DELTA_OFFSET_TABLEBITS     10
216 #define LZMS_DELTA_POWER_TABLEBITS      8
217
218 struct lzms_range_decoder {
219
220         /* The relevant part of the current range.  Although the logical range
221          * for range decoding is a very large integer, only a small portion
222          * matters at any given time, and it can be normalized (shifted left)
223          * whenever it gets too small.  */
224         u32 range;
225
226         /* The current position in the range encoded by the portion of the input
227          * read so far.  */
228         u32 code;
229
230         /* Pointer to the next little-endian 16-bit integer in the compressed
231          * input data (reading forwards).  */
232         const le16 *next;
233
234         /* Pointer to the end of the compressed input data.  */
235         const le16 *end;
236 };
237
238 typedef u64 bitbuf_t;
239
240 struct lzms_input_bitstream {
241
242         /* Holding variable for bits that have been read from the compressed
243          * data.  The bit ordering is high to low.  */
244         bitbuf_t bitbuf;
245
246         /* Number of bits currently held in @bitbuf.  */
247         unsigned bitsleft;
248
249         /* Pointer to the one past the next little-endian 16-bit integer in the
250          * compressed input data (reading backwards).  */
251         const le16 *next;
252
253         /* Pointer to the beginning of the compressed input data.  */
254         const le16 *begin;
255 };
256
257 /* Bookkeeping information for an adaptive Huffman code  */
258 struct lzms_huffman_rebuild_info {
259         unsigned num_syms_until_rebuild;
260         unsigned rebuild_freq;
261         u16 *decode_table;
262         unsigned table_bits;
263         u32 *freqs;
264         u32 *codewords;
265         u8 *lens;
266         unsigned num_syms;
267 };
268
269 struct lzms_decompressor {
270
271         /* 'last_target_usages' is in union with everything else because it is
272          * only used for postprocessing.  */
273         union {
274         struct {
275
276         struct lzms_range_decoder rd;
277         struct lzms_input_bitstream is;
278
279         /* Match offset LRU queues  */
280         u32 recent_lz_offsets[LZMS_NUM_RECENT_OFFSETS + 1];
281         u64 recent_delta_offsets[LZMS_NUM_RECENT_OFFSETS + 1];
282         u32 pending_lz_offset;
283         u64 pending_delta_offset;
284         const u8 *lz_offset_still_pending;
285         const u8 *delta_offset_still_pending;
286
287         /* States and probabilities for range decoding  */
288
289         u32 main_state;
290         struct lzms_probability_entry main_prob_entries[
291                         LZMS_NUM_MAIN_STATES];
292
293         u32 match_state;
294         struct lzms_probability_entry match_prob_entries[
295                         LZMS_NUM_MATCH_STATES];
296
297         u32 lz_match_state;
298         struct lzms_probability_entry lz_match_prob_entries[
299                         LZMS_NUM_LZ_MATCH_STATES];
300
301         u32 delta_match_state;
302         struct lzms_probability_entry delta_match_prob_entries[
303                         LZMS_NUM_DELTA_MATCH_STATES];
304
305         u32 lz_repeat_match_states[LZMS_NUM_RECENT_OFFSETS - 1];
306         struct lzms_probability_entry lz_repeat_match_prob_entries[
307                         LZMS_NUM_RECENT_OFFSETS - 1][LZMS_NUM_LZ_REPEAT_MATCH_STATES];
308
309         u32 delta_repeat_match_states[LZMS_NUM_RECENT_OFFSETS - 1];
310         struct lzms_probability_entry delta_repeat_match_prob_entries[
311                         LZMS_NUM_RECENT_OFFSETS - 1][LZMS_NUM_DELTA_REPEAT_MATCH_STATES];
312
313         /* Huffman decoding  */
314
315         u16 literal_decode_table[(1 << LZMS_LITERAL_TABLEBITS) +
316                                  (2 * LZMS_NUM_LITERAL_SYMS)]
317                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
318         u32 literal_freqs[LZMS_NUM_LITERAL_SYMS];
319         struct lzms_huffman_rebuild_info literal_rebuild_info;
320
321         u16 length_decode_table[(1 << LZMS_LENGTH_TABLEBITS) +
322                                 (2 * LZMS_NUM_LENGTH_SYMS)]
323                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
324         u32 length_freqs[LZMS_NUM_LENGTH_SYMS];
325         struct lzms_huffman_rebuild_info length_rebuild_info;
326
327         u16 lz_offset_decode_table[(1 << LZMS_LZ_OFFSET_TABLEBITS) +
328                                    ( 2 * LZMS_MAX_NUM_OFFSET_SYMS)]
329                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
330         u32 lz_offset_freqs[LZMS_MAX_NUM_OFFSET_SYMS];
331         struct lzms_huffman_rebuild_info lz_offset_rebuild_info;
332
333         u16 delta_offset_decode_table[(1 << LZMS_DELTA_OFFSET_TABLEBITS) +
334                                       (2 * LZMS_MAX_NUM_OFFSET_SYMS)]
335                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
336         u32 delta_offset_freqs[LZMS_MAX_NUM_OFFSET_SYMS];
337         struct lzms_huffman_rebuild_info delta_offset_rebuild_info;
338
339         u16 delta_power_decode_table[(1 << LZMS_DELTA_POWER_TABLEBITS) +
340                                      (2 * LZMS_NUM_DELTA_POWER_SYMS)]
341                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
342         u32 delta_power_freqs[LZMS_NUM_DELTA_POWER_SYMS];
343         struct lzms_huffman_rebuild_info delta_power_rebuild_info;
344
345         u32 codewords[LZMS_MAX_NUM_SYMS];
346         u8 lens[LZMS_MAX_NUM_SYMS];
347
348         }; // struct
349
350         s32 last_target_usages[65536];
351
352         }; // union
353 };
354
355 /* Initialize the input bitstream @is to read backwards from the compressed data
356  * buffer @in that is @count 16-bit integers long.  */
357 static void
358 lzms_input_bitstream_init(struct lzms_input_bitstream *is,
359                           const le16 *in, size_t count)
360 {
361         is->bitbuf = 0;
362         is->bitsleft = 0;
363         is->next = in + count;
364         is->begin = in;
365 }
366
367 /* Ensure that at least @num_bits bits are in the bitbuffer variable.
368  * @num_bits cannot be more than 32.  */
369 static inline void
370 lzms_ensure_bits(struct lzms_input_bitstream *is, unsigned num_bits)
371 {
372         if (is->bitsleft >= num_bits)
373                 return;
374
375         if (likely(is->next != is->begin))
376                 is->bitbuf |= (bitbuf_t)le16_to_cpu(*--is->next)
377                                 << (sizeof(is->bitbuf) * 8 - is->bitsleft - 16);
378         is->bitsleft += 16;
379
380         if (likely(is->next != is->begin))
381                 is->bitbuf |= (bitbuf_t)le16_to_cpu(*--is->next)
382                                 << (sizeof(is->bitbuf) * 8 - is->bitsleft - 16);
383         is->bitsleft += 16;
384 }
385
386 /* Get @num_bits bits from the bitbuffer variable.  */
387 static inline bitbuf_t
388 lzms_peek_bits(struct lzms_input_bitstream *is, unsigned num_bits)
389 {
390         if (unlikely(num_bits == 0))
391                 return 0;
392         return is->bitbuf >> (sizeof(is->bitbuf) * 8 - num_bits);
393 }
394
395 /* Remove @num_bits bits from the bitbuffer variable.  */
396 static inline void
397 lzms_remove_bits(struct lzms_input_bitstream *is, unsigned num_bits)
398 {
399         is->bitbuf <<= num_bits;
400         is->bitsleft -= num_bits;
401 }
402
403 /* Remove and return @num_bits bits from the bitbuffer variable.  */
404 static inline bitbuf_t
405 lzms_pop_bits(struct lzms_input_bitstream *is, unsigned num_bits)
406 {
407         bitbuf_t bits = lzms_peek_bits(is, num_bits);
408         lzms_remove_bits(is, num_bits);
409         return bits;
410 }
411
412 /* Read @num_bits bits from the input bitstream.  */
413 static inline bitbuf_t
414 lzms_read_bits(struct lzms_input_bitstream *is, unsigned num_bits)
415 {
416         lzms_ensure_bits(is, num_bits);
417         return lzms_pop_bits(is, num_bits);
418 }
419
420 /* Initialize the range decoder @rd to read forwards from the compressed data
421  * buffer @in that is @count 16-bit integers long.  */
422 static void
423 lzms_range_decoder_init(struct lzms_range_decoder *rd,
424                         const le16 *in, size_t count)
425 {
426         rd->range = 0xffffffff;
427         rd->code = ((u32)le16_to_cpu(in[0]) << 16) | le16_to_cpu(in[1]);
428         rd->next = in + 2;
429         rd->end = in + count;
430 }
431
432 /* Decode and return the next bit from the range decoder.
433  *
434  * @prob is the chance out of LZMS_PROBABILITY_MAX that the next bit is 0.
435  */
436 static inline int
437 lzms_range_decoder_decode_bit(struct lzms_range_decoder *rd, u32 prob)
438 {
439         u32 bound;
440
441         /* Normalize if needed.  */
442         if (rd->range <= 0xffff) {
443                 rd->range <<= 16;
444                 rd->code <<= 16;
445                 if (likely(rd->next != rd->end))
446                         rd->code |= le16_to_cpu(*rd->next++);
447         }
448
449         /* Based on the probability, calculate the bound between the 0-bit
450          * region and the 1-bit region of the range.  */
451         bound = (rd->range >> LZMS_PROBABILITY_BITS) * prob;
452
453         if (rd->code < bound) {
454                 /* Current code is in the 0-bit region of the range.  */
455                 rd->range = bound;
456                 return 0;
457         } else {
458                 /* Current code is in the 1-bit region of the range.  */
459                 rd->range -= bound;
460                 rd->code -= bound;
461                 return 1;
462         }
463 }
464
465 /* Decode and return the next bit from the range decoder.  This wraps around
466  * lzms_range_decoder_decode_bit() to handle using and updating the appropriate
467  * state and probability entry.  */
468 static inline int
469 lzms_range_decode_bit(struct lzms_range_decoder *rd,
470                       u32 *state_p, u32 num_states,
471                       struct lzms_probability_entry prob_entries[])
472 {
473         struct lzms_probability_entry *prob_entry;
474         u32 prob;
475         int bit;
476
477         /* Load the probability entry corresponding to the current state.  */
478         prob_entry = &prob_entries[*state_p];
479
480         /* Get the probability that the next bit is 0.  */
481         prob = lzms_get_probability(prob_entry);
482
483         /* Decode the next bit.  */
484         bit = lzms_range_decoder_decode_bit(rd, prob);
485
486         /* Update the state and probability entry based on the decoded bit.  */
487         *state_p = ((*state_p << 1) | bit) & (num_states - 1);
488         lzms_update_probability_entry(prob_entry, bit);
489
490         /* Return the decoded bit.  */
491         return bit;
492 }
493
494 static int
495 lzms_decode_main_bit(struct lzms_decompressor *d)
496 {
497         return lzms_range_decode_bit(&d->rd, &d->main_state,
498                                      LZMS_NUM_MAIN_STATES,
499                                      d->main_prob_entries);
500 }
501
502 static int
503 lzms_decode_match_bit(struct lzms_decompressor *d)
504 {
505         return lzms_range_decode_bit(&d->rd, &d->match_state,
506                                      LZMS_NUM_MATCH_STATES,
507                                      d->match_prob_entries);
508 }
509
510 static int
511 lzms_decode_lz_match_bit(struct lzms_decompressor *d)
512 {
513         return lzms_range_decode_bit(&d->rd, &d->lz_match_state,
514                                      LZMS_NUM_LZ_MATCH_STATES,
515                                      d->lz_match_prob_entries);
516 }
517
518 static int
519 lzms_decode_delta_match_bit(struct lzms_decompressor *d)
520 {
521         return lzms_range_decode_bit(&d->rd, &d->delta_match_state,
522                                      LZMS_NUM_DELTA_MATCH_STATES,
523                                      d->delta_match_prob_entries);
524 }
525
526 static noinline int
527 lzms_decode_lz_repeat_match_bit(struct lzms_decompressor *d, int idx)
528 {
529         return lzms_range_decode_bit(&d->rd, &d->lz_repeat_match_states[idx],
530                                      LZMS_NUM_LZ_REPEAT_MATCH_STATES,
531                                      d->lz_repeat_match_prob_entries[idx]);
532 }
533
534 static noinline int
535 lzms_decode_delta_repeat_match_bit(struct lzms_decompressor *d, int idx)
536 {
537         return lzms_range_decode_bit(&d->rd, &d->delta_repeat_match_states[idx],
538                                      LZMS_NUM_DELTA_REPEAT_MATCH_STATES,
539                                      d->delta_repeat_match_prob_entries[idx]);
540 }
541
542 static void
543 lzms_init_huffman_rebuild_info(struct lzms_huffman_rebuild_info *info,
544                                unsigned rebuild_freq,
545                                u16 *decode_table, unsigned table_bits,
546                                u32 *freqs, u32 *codewords, u8 *lens,
547                                unsigned num_syms)
548 {
549         info->num_syms_until_rebuild = 1;
550         info->rebuild_freq = rebuild_freq;
551         info->decode_table = decode_table;
552         info->table_bits = table_bits;
553         info->freqs = freqs;
554         info->codewords = codewords;
555         info->lens = lens;
556         info->num_syms = num_syms;
557         lzms_init_symbol_frequencies(freqs, num_syms);
558 }
559
560 static noinline void
561 lzms_rebuild_huffman_code(struct lzms_huffman_rebuild_info *info)
562 {
563         make_canonical_huffman_code(info->num_syms, LZMS_MAX_CODEWORD_LEN,
564                                     info->freqs, info->lens, info->codewords);
565         make_huffman_decode_table(info->decode_table, info->num_syms,
566                                   info->table_bits, info->lens,
567                                   LZMS_MAX_CODEWORD_LEN);
568         for (unsigned i = 0; i < info->num_syms; i++)
569                 info->freqs[i] = (info->freqs[i] >> 1) + 1;
570         info->num_syms_until_rebuild = info->rebuild_freq;
571 }
572
573 static inline unsigned
574 lzms_decode_huffman_symbol(struct lzms_input_bitstream *is,
575                            u16 decode_table[], unsigned table_bits,
576                            struct lzms_huffman_rebuild_info *rebuild_info)
577 {
578         unsigned key_bits;
579         unsigned entry;
580         unsigned sym;
581
582         if (unlikely(--rebuild_info->num_syms_until_rebuild == 0))
583                 lzms_rebuild_huffman_code(rebuild_info);
584
585         lzms_ensure_bits(is, LZMS_MAX_CODEWORD_LEN);
586
587         /* Index the decode table by the next table_bits bits of the input.  */
588         key_bits = lzms_peek_bits(is, table_bits);
589         entry = decode_table[key_bits];
590         if (likely(entry < 0xC000)) {
591                 /* Fast case: The decode table directly provided the symbol and
592                  * codeword length.  The low 11 bits are the symbol, and the
593                  * high 5 bits are the codeword length.  */
594                 lzms_remove_bits(is, entry >> 11);
595                 sym = entry & 0x7FF;
596         } else {
597                 /* Slow case: The codeword for the symbol is longer than
598                  * table_bits, so the symbol does not have an entry directly in
599                  * the first (1 << table_bits) entries of the decode table.
600                  * Traverse the appropriate binary tree bit-by-bit in order to
601                  * decode the symbol.  */
602                 lzms_remove_bits(is, table_bits);
603                 do {
604                         key_bits = (entry & 0x3FFF) + lzms_pop_bits(is, 1);
605                 } while ((entry = decode_table[key_bits]) >= 0xC000);
606                 sym = entry;
607         }
608
609         /* Tally and return the decoded symbol.  */
610         rebuild_info->freqs[sym]++;
611         return sym;
612 }
613
614 static unsigned
615 lzms_decode_literal(struct lzms_decompressor *d)
616 {
617         return lzms_decode_huffman_symbol(&d->is,
618                                           d->literal_decode_table,
619                                           LZMS_LITERAL_TABLEBITS,
620                                           &d->literal_rebuild_info);
621 }
622
623 static u32
624 lzms_decode_length(struct lzms_decompressor *d)
625 {
626         unsigned slot = lzms_decode_huffman_symbol(&d->is,
627                                                    d->length_decode_table,
628                                                    LZMS_LENGTH_TABLEBITS,
629                                                    &d->length_rebuild_info);
630         u32 length = lzms_length_slot_base[slot];
631         unsigned num_extra_bits = lzms_extra_length_bits[slot];
632         /* Usually most lengths are short and have no extra bits.  */
633         if (num_extra_bits)
634                 length += lzms_read_bits(&d->is, num_extra_bits);
635         return length;
636 }
637
638 static u32
639 lzms_decode_lz_offset(struct lzms_decompressor *d)
640 {
641         unsigned slot = lzms_decode_huffman_symbol(&d->is,
642                                                    d->lz_offset_decode_table,
643                                                    LZMS_LZ_OFFSET_TABLEBITS,
644                                                    &d->lz_offset_rebuild_info);
645         return lzms_offset_slot_base[slot] +
646                lzms_read_bits(&d->is, lzms_extra_offset_bits[slot]);
647 }
648
649 static u32
650 lzms_decode_delta_offset(struct lzms_decompressor *d)
651 {
652         unsigned slot = lzms_decode_huffman_symbol(&d->is,
653                                                    d->delta_offset_decode_table,
654                                                    LZMS_DELTA_OFFSET_TABLEBITS,
655                                                    &d->delta_offset_rebuild_info);
656         return lzms_offset_slot_base[slot] +
657                lzms_read_bits(&d->is, lzms_extra_offset_bits[slot]);
658 }
659
660 static unsigned
661 lzms_decode_delta_power(struct lzms_decompressor *d)
662 {
663         return lzms_decode_huffman_symbol(&d->is,
664                                           d->delta_power_decode_table,
665                                           LZMS_DELTA_POWER_TABLEBITS,
666                                           &d->delta_power_rebuild_info);
667 }
668
669 /* Decode the series of literals and matches from the LZMS-compressed data.
670  * Return 0 if successful or -1 if the compressed data is invalid.  */
671 static int
672 lzms_decode_items(struct lzms_decompressor * const restrict d,
673                   u8 * const restrict out, const size_t out_nbytes)
674 {
675         u8 *out_next = out;
676         u8 * const out_end = out + out_nbytes;
677
678         while (out_next != out_end) {
679
680                 if (!lzms_decode_main_bit(d)) {
681
682                         /* Literal  */
683                         *out_next++ = lzms_decode_literal(d);
684
685                 } else if (!lzms_decode_match_bit(d)) {
686
687                         /* LZ match  */
688
689                         u32 offset;
690                         u32 length;
691
692                         if (d->pending_lz_offset != 0 &&
693                             out_next != d->lz_offset_still_pending)
694                         {
695                                 BUILD_BUG_ON(LZMS_NUM_RECENT_OFFSETS != 3);
696                                 d->recent_lz_offsets[3] = d->recent_lz_offsets[2];
697                                 d->recent_lz_offsets[2] = d->recent_lz_offsets[1];
698                                 d->recent_lz_offsets[1] = d->recent_lz_offsets[0];
699                                 d->recent_lz_offsets[0] = d->pending_lz_offset;
700                                 d->pending_lz_offset = 0;
701                         }
702
703                         if (!lzms_decode_lz_match_bit(d)) {
704                                 /* Explicit offset  */
705                                 offset = lzms_decode_lz_offset(d);
706                         } else {
707                                 /* Repeat offset  */
708
709                                 BUILD_BUG_ON(LZMS_NUM_RECENT_OFFSETS != 3);
710                                 if (!lzms_decode_lz_repeat_match_bit(d, 0)) {
711                                         offset = d->recent_lz_offsets[0];
712                                         d->recent_lz_offsets[0] = d->recent_lz_offsets[1];
713                                         d->recent_lz_offsets[1] = d->recent_lz_offsets[2];
714                                         d->recent_lz_offsets[2] = d->recent_lz_offsets[3];
715                                 } else if (!lzms_decode_lz_repeat_match_bit(d, 1)) {
716                                         offset = d->recent_lz_offsets[1];
717                                         d->recent_lz_offsets[1] = d->recent_lz_offsets[2];
718                                         d->recent_lz_offsets[2] = d->recent_lz_offsets[3];
719                                 } else {
720                                         offset = d->recent_lz_offsets[2];
721                                         d->recent_lz_offsets[2] = d->recent_lz_offsets[3];
722                                 }
723                         }
724
725                         if (d->pending_lz_offset != 0) {
726                                 BUILD_BUG_ON(LZMS_NUM_RECENT_OFFSETS != 3);
727                                 d->recent_lz_offsets[3] = d->recent_lz_offsets[2];
728                                 d->recent_lz_offsets[2] = d->recent_lz_offsets[1];
729                                 d->recent_lz_offsets[1] = d->recent_lz_offsets[0];
730                                 d->recent_lz_offsets[0] = d->pending_lz_offset;
731                         }
732                         d->pending_lz_offset = offset;
733
734                         length = lzms_decode_length(d);
735
736                         if (unlikely(length > out_end - out_next))
737                                 return -1;
738                         if (unlikely(offset > out_next - out))
739                                 return -1;
740
741                         lz_copy(out_next, length, offset, out_end, LZMS_MIN_MATCH_LEN);
742                         out_next += length;
743
744                         d->lz_offset_still_pending = out_next;
745                 } else {
746                         /* Delta match  */
747
748                         u32 power;
749                         u32 raw_offset, offset1, offset2, offset;
750                         const u8 *matchptr1, *matchptr2, *matchptr;
751                         u32 length;
752
753                         if (d->pending_delta_offset != 0 &&
754                             out_next != d->delta_offset_still_pending)
755                         {
756                                 BUILD_BUG_ON(LZMS_NUM_RECENT_OFFSETS != 3);
757                                 d->recent_delta_offsets[3] = d->recent_delta_offsets[2];
758                                 d->recent_delta_offsets[2] = d->recent_delta_offsets[1];
759                                 d->recent_delta_offsets[1] = d->recent_delta_offsets[0];
760                                 d->recent_delta_offsets[0] = d->pending_delta_offset;
761                                 d->pending_delta_offset = 0;
762                         }
763
764                         if (!lzms_decode_delta_match_bit(d)) {
765                                 /* Explicit offset  */
766                                 power = lzms_decode_delta_power(d);
767                                 raw_offset = lzms_decode_delta_offset(d);
768                         } else {
769                                 /* Repeat offset  */
770                                 u64 val;
771
772                                 BUILD_BUG_ON(LZMS_NUM_RECENT_OFFSETS != 3);
773                                 if (!lzms_decode_delta_repeat_match_bit(d, 0)) {
774                                         val = d->recent_delta_offsets[0];
775                                         d->recent_delta_offsets[0] = d->recent_delta_offsets[1];
776                                         d->recent_delta_offsets[1] = d->recent_delta_offsets[2];
777                                         d->recent_delta_offsets[2] = d->recent_delta_offsets[3];
778                                 } else if (!lzms_decode_delta_repeat_match_bit(d, 1)) {
779                                         val = d->recent_delta_offsets[1];
780                                         d->recent_delta_offsets[1] = d->recent_delta_offsets[2];
781                                         d->recent_delta_offsets[2] = d->recent_delta_offsets[3];
782                                 } else {
783                                         val = d->recent_delta_offsets[2];
784                                         d->recent_delta_offsets[2] = d->recent_delta_offsets[3];
785                                 }
786                                 power = val >> 32;
787                                 raw_offset = (u32)val;
788                         }
789
790                         if (d->pending_delta_offset != 0) {
791                                 BUILD_BUG_ON(LZMS_NUM_RECENT_OFFSETS != 3);
792                                 d->recent_delta_offsets[3] = d->recent_delta_offsets[2];
793                                 d->recent_delta_offsets[2] = d->recent_delta_offsets[1];
794                                 d->recent_delta_offsets[1] = d->recent_delta_offsets[0];
795                                 d->recent_delta_offsets[0] = d->pending_delta_offset;
796                                 d->pending_delta_offset = 0;
797                         }
798                         d->pending_delta_offset = raw_offset | ((u64)power << 32);
799
800                         length = lzms_decode_length(d);
801
802                         offset1 = (u32)1 << power;
803                         offset2 = raw_offset << power;
804                         offset = offset1 + offset2;
805
806                         /* raw_offset<<power overflowed?  */
807                         if (unlikely((offset2 >> power) != raw_offset))
808                                 return -1;
809
810                         /* offset1+offset2 overflowed?  */
811                         if (unlikely(offset < offset2))
812                                 return -1;
813
814                         if (unlikely(length > out_end - out_next))
815                                 return -1;
816
817                         if (unlikely(offset > out_next - out))
818                                 return -1;
819
820                         matchptr1 = out_next - offset1;
821                         matchptr2 = out_next - offset2;
822                         matchptr = out_next - offset;
823
824                         do {
825                                 *out_next++ = *matchptr1++ + *matchptr2++ - *matchptr++;
826                         } while (--length);
827
828                         d->delta_offset_still_pending = out_next;
829                 }
830         }
831         return 0;
832 }
833
834 static void
835 lzms_init_decompressor(struct lzms_decompressor *d, const void *in,
836                        size_t in_nbytes, unsigned num_offset_slots)
837 {
838         /* Match offset LRU queues  */
839         for (int i = 0; i < LZMS_NUM_RECENT_OFFSETS + 1; i++) {
840                 d->recent_lz_offsets[i] = i + 1;
841                 d->recent_delta_offsets[i] = i + 1;
842         }
843         d->pending_lz_offset = 0;
844         d->pending_delta_offset = 0;
845
846         /* Range decoding  */
847
848         lzms_range_decoder_init(&d->rd, in, in_nbytes / sizeof(le16));
849
850         d->main_state = 0;
851         lzms_init_probability_entries(d->main_prob_entries, LZMS_NUM_MAIN_STATES);
852
853         d->match_state = 0;
854         lzms_init_probability_entries(d->match_prob_entries, LZMS_NUM_MATCH_STATES);
855
856         d->lz_match_state = 0;
857         lzms_init_probability_entries(d->lz_match_prob_entries, LZMS_NUM_LZ_MATCH_STATES);
858
859         d->delta_match_state = 0;
860         lzms_init_probability_entries(d->delta_match_prob_entries, LZMS_NUM_DELTA_MATCH_STATES);
861
862         for (int i = 0; i < LZMS_NUM_RECENT_OFFSETS - 1; i++) {
863                 d->lz_repeat_match_states[i] = 0;
864                 lzms_init_probability_entries(d->lz_repeat_match_prob_entries[i],
865                                               LZMS_NUM_LZ_REPEAT_MATCH_STATES);
866
867                 d->delta_repeat_match_states[i] = 0;
868                 lzms_init_probability_entries(d->delta_repeat_match_prob_entries[i],
869                                               LZMS_NUM_DELTA_REPEAT_MATCH_STATES);
870         }
871
872         /* Huffman decoding  */
873
874         lzms_input_bitstream_init(&d->is, in, in_nbytes / sizeof(le16));
875
876         lzms_init_huffman_rebuild_info(&d->literal_rebuild_info,
877                                        LZMS_LITERAL_CODE_REBUILD_FREQ,
878                                        d->literal_decode_table,
879                                        LZMS_LITERAL_TABLEBITS,
880                                        d->literal_freqs,
881                                        d->codewords,
882                                        d->lens,
883                                        LZMS_NUM_LITERAL_SYMS);
884
885         lzms_init_huffman_rebuild_info(&d->length_rebuild_info,
886                                        LZMS_LENGTH_CODE_REBUILD_FREQ,
887                                        d->length_decode_table,
888                                        LZMS_LENGTH_TABLEBITS,
889                                        d->length_freqs,
890                                        d->codewords,
891                                        d->lens,
892                                        LZMS_NUM_LENGTH_SYMS);
893
894         lzms_init_huffman_rebuild_info(&d->lz_offset_rebuild_info,
895                                        LZMS_LZ_OFFSET_CODE_REBUILD_FREQ,
896                                        d->lz_offset_decode_table,
897                                        LZMS_LZ_OFFSET_TABLEBITS,
898                                        d->lz_offset_freqs,
899                                        d->codewords,
900                                        d->lens,
901                                        num_offset_slots);
902
903         lzms_init_huffman_rebuild_info(&d->delta_offset_rebuild_info,
904                                        LZMS_DELTA_OFFSET_CODE_REBUILD_FREQ,
905                                        d->delta_offset_decode_table,
906                                        LZMS_DELTA_OFFSET_TABLEBITS,
907                                        d->delta_offset_freqs,
908                                        d->codewords,
909                                        d->lens,
910                                        num_offset_slots);
911
912         lzms_init_huffman_rebuild_info(&d->delta_power_rebuild_info,
913                                        LZMS_DELTA_POWER_CODE_REBUILD_FREQ,
914                                        d->delta_power_decode_table,
915                                        LZMS_DELTA_POWER_TABLEBITS,
916                                        d->delta_power_freqs,
917                                        d->codewords,
918                                        d->lens,
919                                        LZMS_NUM_DELTA_POWER_SYMS);
920 }
921
922 static int
923 lzms_create_decompressor(size_t max_bufsize, void **d_ret)
924 {
925         struct lzms_decompressor *d;
926
927         if (max_bufsize > LZMS_MAX_BUFFER_SIZE)
928                 return WIMLIB_ERR_INVALID_PARAM;
929
930         d = ALIGNED_MALLOC(sizeof(struct lzms_decompressor),
931                            DECODE_TABLE_ALIGNMENT);
932         if (!d)
933                 return WIMLIB_ERR_NOMEM;
934
935         *d_ret = d;
936         return 0;
937 }
938
939 /* Decompress @in_nbytes bytes of LZMS-compressed data at @in and write the
940  * uncompressed data, which had original size @out_nbytes, to @out.  Return 0 if
941  * successful or -1 if the compressed data is invalid.  */
942 static int
943 lzms_decompress(const void *in, size_t in_nbytes, void *out, size_t out_nbytes,
944                 void *_d)
945 {
946         struct lzms_decompressor *d = _d;
947
948         /*
949          * Requirements on the compressed data:
950          *
951          * 1. LZMS-compressed data is a series of 16-bit integers, so the
952          *    compressed data buffer cannot take up an odd number of bytes.
953          * 2. To prevent poor performance on some architectures, we require that
954          *    the compressed data buffer is 2-byte aligned.
955          * 3. There must be at least 4 bytes of compressed data, since otherwise
956          *    we cannot even initialize the range decoder.
957          */
958         if ((in_nbytes & 1) || ((uintptr_t)in & 1) || (in_nbytes < 4))
959                 return -1;
960
961         lzms_init_decompressor(d, in, in_nbytes,
962                                lzms_get_num_offset_slots(out_nbytes));
963
964         if (lzms_decode_items(d, out, out_nbytes))
965                 return -1;
966
967         lzms_x86_filter(out, out_nbytes, d->last_target_usages, true);
968         return 0;
969 }
970
971 static void
972 lzms_free_decompressor(void *_d)
973 {
974         struct lzms_decompressor *d = _d;
975
976         ALIGNED_FREE(d);
977 }
978
979 const struct decompressor_ops lzms_decompressor_ops = {
980         .create_decompressor  = lzms_create_decompressor,
981         .decompress           = lzms_decompress,
982         .free_decompressor    = lzms_free_decompressor,
983 };