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