f1ee951e7a7aab586cbe28cb16d912926ffc9e83
[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 the queue contains
130  * (power, raw offset) pairs.  This queue must be initialized to
131  * {(0, 1), (0, 2), (0, 3), (0, 4)}.
132  *
133  * Bits from the binary range decoder must be used to disambiguate item types.
134  * The range decoder must hold two state variables: the range, which must
135  * initially be set to 0xffffffff, and the current code, which must initially be
136  * set to the first 32 bits read from the forwards bitstream.  The range must be
137  * maintained above 0xffff; when it falls below 0xffff, both the range and code
138  * must be left-shifted by 16 bits and the low 16 bits of the code must be
139  * filled in with the next 16 bits from the forwards bitstream.
140  *
141  * To decode each bit, the binary range decoder requires a probability that is
142  * logically a real number between 0 and 1.  Multiplying this probability by the
143  * current range and taking the floor gives the bound between the 0-bit region of
144  * the range and the 1-bit region of the range.  However, in LZMS, probabilities
145  * are restricted to values of n/64 where n is an integer is between 1 and 63
146  * inclusively, so the implementation may use integer operations instead.
147  * Following calculation of the bound, if the current code is in the 0-bit
148  * region, the new range becomes the current code and the decoded bit is 0;
149  * otherwise, the bound must be subtracted from both the range and the code, and
150  * the decoded bit is 1.  More information about range coding can be found at
151  * https://en.wikipedia.org/wiki/Range_encoding.  Furthermore, note that the
152  * LZMA format also uses range coding and has public domain code available for
153  * it.
154  *
155  * The probability used to range-decode each bit must be taken from a table, of
156  * which one instance must exist for each distinct context, or "binary decision
157  * class", in which a range-decoded bit is needed.  At each call of the range
158  * decoder, the appropriate probability must be obtained by indexing the
159  * appropriate probability table with the last 4 (in the context disambiguating
160  * literals from matches), 5 (in the context disambiguating LZ matches from
161  * delta matches), or 6 (in all other contexts) bits recently range-decoded in
162  * that context, ordered such that the most recently decoded bit is the
163  * low-order bit of the index.
164  *
165  * Furthermore, each probability entry itself is variable, as its value must be
166  * maintained as n/64 where n is the number of 0 bits in the most recently
167  * decoded 64 bits with that same entry.  This allows the compressed
168  * representation to adapt to the input and use fewer bits to represent the most
169  * likely data; note that LZMA uses a similar scheme.  Initially, the most
170  * recently 64 decoded bits for each probability entry are assumed to be
171  * 0x0000000055555555 (high order to low order); therefore, all probabilities
172  * are initially 48/64.  During the course of decoding, each probability may be
173  * updated to as low as 0/64 (as a result of reading many consecutive 1 bits
174  * with that entry) or as high as 64/64 (as a result of reading many consecutive
175  * 0 bits with that entry); however, probabilities of 0/64 and 64/64 cannot be
176  * used as-is but rather must be adjusted to 1/64 and 63/64, respectively,
177  * before being used for range decoding.
178  *
179  * Representations of the LZMS items themselves must be read from the backwards
180  * bitstream.  For this, there are 5 different Huffman codes used:
181  *
182  *  - The literal code, used for decoding literal bytes.  Each of the 256
183  *    symbols represents a literal byte.  This code must be rebuilt whenever
184  *    1024 symbols have been decoded with it.
185  *
186  *  - The LZ offset code, used for decoding the offsets of standard LZ77
187  *    matches.  Each symbol represents an offset slot, which corresponds to a
188  *    base value and some number of extra bits which must be read and added to
189  *    the base value to reconstitute the full offset.  The number of symbols in
190  *    this code is the number of offset slots needed to represent all possible
191  *    offsets in the uncompressed block.  This code must be rebuilt whenever
192  *    1024 symbols have been decoded with it.
193  *
194  *  - The length code, used for decoding length symbols.  Each of the 54 symbols
195  *    represents a length slot, which corresponds to a base value and some
196  *    number of extra bits which must be read and added to the base value to
197  *    reconstitute the full length.  This code must be rebuilt whenever 512
198  *    symbols have been decoded with it.
199  *
200  *  - The delta offset code, used for decoding the raw offsets of delta matches.
201  *    Each symbol corresponds to an offset slot, which corresponds to a base
202  *    value and some number of extra bits which must be read and added to the
203  *    base value to reconstitute the full raw offset.  The number of symbols in
204  *    this code is equal to the number of symbols in the LZ offset code.  This
205  *    code must be rebuilt whenever 1024 symbols have been decoded with it.
206  *
207  *  - The delta power code, used for decoding the powers of delta matches.  Each
208  *    of the 8 symbols corresponds to a power.  This code must be rebuilt
209  *    whenever 512 symbols have been decoded with it.
210  *
211  * Initially, each Huffman code must be built assuming that each symbol in that
212  * code has frequency 1.  Following that, each code must be rebuilt each time a
213  * certain number of symbols, as noted above, has been decoded with it.  The
214  * symbol frequencies for a code must be halved after each rebuild of that code;
215  * this makes the codes adapt to the more recent data.
216  *
217  * Like other compression formats such as XPRESS, LZX, and DEFLATE, the LZMS
218  * format requires that all Huffman codes be constructed in canonical form.
219  * This form requires that same-length codewords be lexicographically ordered
220  * the same way as the corresponding symbols and that all shorter codewords
221  * lexicographically precede longer codewords.  Such a code can be constructed
222  * directly from codeword lengths.
223  *
224  * Even with the canonical code restriction, the same frequencies can be used to
225  * construct multiple valid Huffman codes.  Therefore, the decompressor needs to
226  * construct the right one.  Specifically, the LZMS format requires that the
227  * Huffman code be constructed as if the well-known priority queue algorithm is
228  * used and frequency ties are always broken in favor of leaf nodes.
229  *
230  * Codewords in LZMS are guaranteed to not exceed 15 bits.  The format otherwise
231  * places no restrictions on codeword length.  Therefore, the Huffman code
232  * construction algorithm that a correct LZMS decompressor uses need not
233  * implement length-limited code construction.  But if it does (e.g. by virtue
234  * of being shared among multiple compression algorithms), the details of how it
235  * does so are unimportant, provided that the maximum codeword length parameter
236  * is set to at least 15 bits.
237  *
238  * After all LZMS items have been decoded, the data must be postprocessed to
239  * translate absolute address encoded in x86 instructions into their original
240  * relative addresses.
241  *
242  * Details omitted above can be found in the code.  Note that in the absence of
243  * an official specification there is no guarantee that this decompressor
244  * handles all possible cases.
245  */
246
247 #ifdef HAVE_CONFIG_H
248 #  include "config.h"
249 #endif
250
251 #include "wimlib/compress_common.h"
252 #include "wimlib/decompress_common.h"
253 #include "wimlib/decompressor_ops.h"
254 #include "wimlib/error.h"
255 #include "wimlib/lzms_common.h"
256 #include "wimlib/util.h"
257
258 /* The TABLEBITS values can be changed; they only affect decoding speed.  */
259 #define LZMS_LITERAL_TABLEBITS          10
260 #define LZMS_LENGTH_TABLEBITS           10
261 #define LZMS_LZ_OFFSET_TABLEBITS        10
262 #define LZMS_DELTA_OFFSET_TABLEBITS     10
263 #define LZMS_DELTA_POWER_TABLEBITS      8
264
265 struct lzms_range_decoder {
266
267         /* The relevant part of the current range.  Although the logical range
268          * for range decoding is a very large integer, only a small portion
269          * matters at any given time, and it can be normalized (shifted left)
270          * whenever it gets too small.  */
271         u32 range;
272
273         /* The current position in the range encoded by the portion of the input
274          * read so far.  */
275         u32 code;
276
277         /* Pointer to the next little-endian 16-bit integer in the compressed
278          * input data (reading forwards).  */
279         const le16 *next;
280
281         /* Pointer to the end of the compressed input data.  */
282         const le16 *end;
283 };
284
285 typedef u64 bitbuf_t;
286
287 struct lzms_input_bitstream {
288
289         /* Holding variable for bits that have been read from the compressed
290          * data.  The bit ordering is high to low.  */
291         bitbuf_t bitbuf;
292
293         /* Number of bits currently held in @bitbuf.  */
294         unsigned bitsleft;
295
296         /* Pointer to the one past the next little-endian 16-bit integer in the
297          * compressed input data (reading backwards).  */
298         const le16 *next;
299
300         /* Pointer to the beginning of the compressed input data.  */
301         const le16 *begin;
302 };
303
304 #define BITBUF_NBITS    (8 * sizeof(bitbuf_t))
305
306 /* Bookkeeping information for an adaptive Huffman code  */
307 struct lzms_huffman_rebuild_info {
308         unsigned num_syms_until_rebuild;
309         unsigned num_syms;
310         unsigned rebuild_freq;
311         u32 *codewords;
312         u32 *freqs;
313         u16 *decode_table;
314         unsigned table_bits;
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_probabilites probs;
325
326         u16 literal_decode_table[(1 << LZMS_LITERAL_TABLEBITS) +
327                                  (2 * LZMS_NUM_LITERAL_SYMS)]
328                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
329         u32 literal_freqs[LZMS_NUM_LITERAL_SYMS];
330         struct lzms_huffman_rebuild_info literal_rebuild_info;
331
332         u16 lz_offset_decode_table[(1 << LZMS_LZ_OFFSET_TABLEBITS) +
333                                    ( 2 * LZMS_MAX_NUM_OFFSET_SYMS)]
334                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
335         u32 lz_offset_freqs[LZMS_MAX_NUM_OFFSET_SYMS];
336         struct lzms_huffman_rebuild_info lz_offset_rebuild_info;
337
338         u16 length_decode_table[(1 << LZMS_LENGTH_TABLEBITS) +
339                                 (2 * LZMS_NUM_LENGTH_SYMS)]
340                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
341         u32 length_freqs[LZMS_NUM_LENGTH_SYMS];
342         struct lzms_huffman_rebuild_info length_rebuild_info;
343
344         u16 delta_offset_decode_table[(1 << LZMS_DELTA_OFFSET_TABLEBITS) +
345                                       (2 * LZMS_MAX_NUM_OFFSET_SYMS)]
346                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
347         u32 delta_offset_freqs[LZMS_MAX_NUM_OFFSET_SYMS];
348         struct lzms_huffman_rebuild_info delta_offset_rebuild_info;
349
350         u16 delta_power_decode_table[(1 << LZMS_DELTA_POWER_TABLEBITS) +
351                                      (2 * LZMS_NUM_DELTA_POWER_SYMS)]
352                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
353         u32 delta_power_freqs[LZMS_NUM_DELTA_POWER_SYMS];
354         struct lzms_huffman_rebuild_info delta_power_rebuild_info;
355
356         u32 codewords[LZMS_MAX_NUM_SYMS];
357
358         }; // struct
359
360         s32 last_target_usages[65536];
361
362         }; // union
363 };
364
365 /* Initialize the input bitstream @is to read backwards from the compressed data
366  * buffer @in that is @count 16-bit integers long.  */
367 static void
368 lzms_input_bitstream_init(struct lzms_input_bitstream *is,
369                           const le16 *in, size_t count)
370 {
371         is->bitbuf = 0;
372         is->bitsleft = 0;
373         is->next = in + count;
374         is->begin = in;
375 }
376
377 /* Ensure that at least @num_bits bits are in the bitbuffer variable.
378  * @num_bits cannot be more than 32.  */
379 static inline void
380 lzms_ensure_bits(struct lzms_input_bitstream *is, unsigned num_bits)
381 {
382         unsigned avail;
383
384         if (is->bitsleft >= num_bits)
385                 return;
386
387         avail = BITBUF_NBITS - is->bitsleft;
388
389         if (UNALIGNED_ACCESS_IS_FAST && CPU_IS_LITTLE_ENDIAN &&
390             WORDSIZE == 8 && likely((u8 *)is->next - (u8 *)is->begin >= 8))
391         {
392                 is->next -= avail >> 4;
393                 is->bitbuf |= load_u64_unaligned(is->next) << (avail & 15);
394                 is->bitsleft += avail & ~15;
395         } else {
396                 if (likely(is->next != is->begin))
397                         is->bitbuf |= (bitbuf_t)le16_to_cpu(*--is->next)
398                                         << (avail - 16);
399                 if (likely(is->next != is->begin))
400                         is->bitbuf |=(bitbuf_t)le16_to_cpu(*--is->next)
401                                         << (avail - 32);
402                 is->bitsleft += 32;
403         }
404 }
405
406 /* Get @num_bits bits from the bitbuffer variable.  */
407 static inline bitbuf_t
408 lzms_peek_bits(struct lzms_input_bitstream *is, unsigned num_bits)
409 {
410         return (is->bitbuf >> 1) >> (BITBUF_NBITS - num_bits - 1);
411 }
412
413 /* Remove @num_bits bits from the bitbuffer variable.  */
414 static inline void
415 lzms_remove_bits(struct lzms_input_bitstream *is, unsigned num_bits)
416 {
417         is->bitbuf <<= num_bits;
418         is->bitsleft -= num_bits;
419 }
420
421 /* Remove and return @num_bits bits from the bitbuffer variable.  */
422 static inline bitbuf_t
423 lzms_pop_bits(struct lzms_input_bitstream *is, unsigned num_bits)
424 {
425         bitbuf_t bits = lzms_peek_bits(is, num_bits);
426         lzms_remove_bits(is, num_bits);
427         return bits;
428 }
429
430 /* Read @num_bits bits from the input bitstream.  */
431 static inline bitbuf_t
432 lzms_read_bits(struct lzms_input_bitstream *is, unsigned num_bits)
433 {
434         lzms_ensure_bits(is, num_bits);
435         return lzms_pop_bits(is, num_bits);
436 }
437
438 /* Initialize the range decoder @rd to read forwards from the compressed data
439  * buffer @in that is @count 16-bit integers long.  */
440 static void
441 lzms_range_decoder_init(struct lzms_range_decoder *rd,
442                         const le16 *in, size_t count)
443 {
444         rd->range = 0xffffffff;
445         rd->code = ((u32)le16_to_cpu(in[0]) << 16) | le16_to_cpu(in[1]);
446         rd->next = in + 2;
447         rd->end = in + count;
448 }
449
450 /*
451  * Decode a bit using the range coder.  The current state specifies the
452  * probability entry to use.  The state and probability entry will be updated
453  * based on the decoded bit.
454  */
455 static inline int
456 lzms_decode_bit(struct lzms_range_decoder *rd, u32 *state_p, u32 num_states,
457                 struct lzms_probability_entry *probs)
458 {
459         struct lzms_probability_entry *prob_entry;
460         u32 prob;
461         u32 bound;
462
463         /* Load the probability entry corresponding to the current state.  */
464         prob_entry = &probs[*state_p];
465
466         /* Update the state early.  We'll still need to OR the state with 1
467          * later if the decoded bit is a 1.  */
468         *state_p = (*state_p << 1) & (num_states - 1);
469
470         /* Get the probability (out of LZMS_PROBABILITY_DENOMINATOR) that the
471          * next bit is 0.  */
472         prob = lzms_get_probability(prob_entry);
473
474         /* Normalize if needed.  */
475         if (!(rd->range & 0xFFFF0000)) {
476                 rd->range <<= 16;
477                 rd->code <<= 16;
478                 if (likely(rd->next != rd->end))
479                         rd->code |= le16_to_cpu(*rd->next++);
480         }
481
482         /* Based on the probability, calculate the bound between the 0-bit
483          * region and the 1-bit region of the range.  */
484         bound = (rd->range >> LZMS_PROBABILITY_BITS) * prob;
485
486         if (rd->code < bound) {
487                 /* Current code is in the 0-bit region of the range.  */
488                 rd->range = bound;
489
490                 /* Update the state and probability entry based on the decoded bit.  */
491                 lzms_update_probability_entry(prob_entry, 0);
492                 return 0;
493         } else {
494                 /* Current code is in the 1-bit region of the range.  */
495                 rd->range -= bound;
496                 rd->code -= bound;
497
498                 /* Update the state and probability entry based on the decoded bit.  */
499                 lzms_update_probability_entry(prob_entry, 1);
500                 *state_p |= 1;
501                 return 1;
502         }
503 }
504
505 static void
506 lzms_build_huffman_code(struct lzms_huffman_rebuild_info *rebuild_info)
507 {
508         make_canonical_huffman_code(rebuild_info->num_syms,
509                                     LZMS_MAX_CODEWORD_LENGTH,
510                                     rebuild_info->freqs,
511                                     (u8 *)rebuild_info->decode_table,
512                                     rebuild_info->codewords);
513
514         make_huffman_decode_table(rebuild_info->decode_table,
515                                   rebuild_info->num_syms,
516                                   rebuild_info->table_bits,
517                                   (u8 *)rebuild_info->decode_table,
518                                   LZMS_MAX_CODEWORD_LENGTH);
519
520         rebuild_info->num_syms_until_rebuild = rebuild_info->rebuild_freq;
521 }
522
523 static void
524 lzms_init_huffman_code(struct lzms_huffman_rebuild_info *rebuild_info,
525                        unsigned num_syms, unsigned rebuild_freq,
526                        u32 *codewords, u32 *freqs,
527                        u16 *decode_table, unsigned table_bits)
528 {
529         rebuild_info->num_syms = num_syms;
530         rebuild_info->rebuild_freq = rebuild_freq;
531         rebuild_info->codewords = codewords;
532         rebuild_info->freqs = freqs;
533         rebuild_info->decode_table = decode_table;
534         rebuild_info->table_bits = table_bits;
535         lzms_init_symbol_frequencies(freqs, num_syms);
536         lzms_build_huffman_code(rebuild_info);
537 }
538
539 static void
540 lzms_init_huffman_codes(struct lzms_decompressor *d, unsigned num_offset_slots)
541 {
542         lzms_init_huffman_code(&d->literal_rebuild_info,
543                                LZMS_NUM_LITERAL_SYMS,
544                                LZMS_LITERAL_CODE_REBUILD_FREQ,
545                                d->codewords,
546                                d->literal_freqs,
547                                d->literal_decode_table,
548                                LZMS_LITERAL_TABLEBITS);
549
550         lzms_init_huffman_code(&d->lz_offset_rebuild_info,
551                                num_offset_slots,
552                                LZMS_LZ_OFFSET_CODE_REBUILD_FREQ,
553                                d->codewords,
554                                d->lz_offset_freqs,
555                                d->lz_offset_decode_table,
556                                LZMS_LZ_OFFSET_TABLEBITS);
557
558         lzms_init_huffman_code(&d->length_rebuild_info,
559                                LZMS_NUM_LENGTH_SYMS,
560                                LZMS_LENGTH_CODE_REBUILD_FREQ,
561                                d->codewords,
562                                d->length_freqs,
563                                d->length_decode_table,
564                                LZMS_LENGTH_TABLEBITS);
565
566         lzms_init_huffman_code(&d->delta_offset_rebuild_info,
567                                num_offset_slots,
568                                LZMS_DELTA_OFFSET_CODE_REBUILD_FREQ,
569                                d->codewords,
570                                d->delta_offset_freqs,
571                                d->delta_offset_decode_table,
572                                LZMS_DELTA_OFFSET_TABLEBITS);
573
574         lzms_init_huffman_code(&d->delta_power_rebuild_info,
575                                LZMS_NUM_DELTA_POWER_SYMS,
576                                LZMS_DELTA_POWER_CODE_REBUILD_FREQ,
577                                d->codewords,
578                                d->delta_power_freqs,
579                                d->delta_power_decode_table,
580                                LZMS_DELTA_POWER_TABLEBITS);
581 }
582
583 static noinline void
584 lzms_rebuild_huffman_code(struct lzms_huffman_rebuild_info *rebuild_info)
585 {
586         lzms_build_huffman_code(rebuild_info);
587         lzms_dilute_symbol_frequencies(rebuild_info->freqs, rebuild_info->num_syms);
588 }
589
590 static inline unsigned
591 lzms_decode_huffman_symbol(struct lzms_input_bitstream *is, u16 decode_table[],
592                            unsigned table_bits, u32 freqs[],
593                            struct lzms_huffman_rebuild_info *rebuild_info)
594 {
595         unsigned key_bits;
596         unsigned entry;
597         unsigned sym;
598
599         lzms_ensure_bits(is, LZMS_MAX_CODEWORD_LENGTH);
600
601         /* Index the decode table by the next table_bits bits of the input.  */
602         key_bits = lzms_peek_bits(is, table_bits);
603         entry = decode_table[key_bits];
604         if (likely(entry < 0xC000)) {
605                 /* Fast case: The decode table directly provided the symbol and
606                  * codeword length.  The low 11 bits are the symbol, and the
607                  * high 5 bits are the codeword length.  */
608                 lzms_remove_bits(is, entry >> 11);
609                 sym = entry & 0x7FF;
610         } else {
611                 /* Slow case: The codeword for the symbol is longer than
612                  * table_bits, so the symbol does not have an entry directly in
613                  * the first (1 << table_bits) entries of the decode table.
614                  * Traverse the appropriate binary tree bit-by-bit in order to
615                  * decode the symbol.  */
616                 lzms_remove_bits(is, table_bits);
617                 do {
618                         key_bits = (entry & 0x3FFF) + lzms_pop_bits(is, 1);
619                 } while ((entry = decode_table[key_bits]) >= 0xC000);
620                 sym = entry;
621         }
622
623         freqs[sym]++;
624         if (--rebuild_info->num_syms_until_rebuild == 0)
625                 lzms_rebuild_huffman_code(rebuild_info);
626         return sym;
627 }
628
629 static inline unsigned
630 lzms_decode_literal(struct lzms_decompressor *d,
631                     struct lzms_input_bitstream *is)
632 {
633         return lzms_decode_huffman_symbol(is,
634                                           d->literal_decode_table,
635                                           LZMS_LITERAL_TABLEBITS,
636                                           d->literal_freqs,
637                                           &d->literal_rebuild_info);
638 }
639
640 static inline u32
641 lzms_decode_lz_offset(struct lzms_decompressor *d,
642                       struct lzms_input_bitstream *is)
643 {
644         unsigned slot = lzms_decode_huffman_symbol(is,
645                                                    d->lz_offset_decode_table,
646                                                    LZMS_LZ_OFFSET_TABLEBITS,
647                                                    d->lz_offset_freqs,
648                                                    &d->lz_offset_rebuild_info);
649         return lzms_offset_slot_base[slot] +
650                lzms_read_bits(is, lzms_extra_offset_bits[slot]);
651 }
652
653 static inline u32
654 lzms_decode_length(struct lzms_decompressor *d,
655                    struct lzms_input_bitstream *is)
656 {
657         unsigned slot = lzms_decode_huffman_symbol(is,
658                                                    d->length_decode_table,
659                                                    LZMS_LENGTH_TABLEBITS,
660                                                    d->length_freqs,
661                                                    &d->length_rebuild_info);
662         u32 length = lzms_length_slot_base[slot];
663         unsigned num_extra_bits = lzms_extra_length_bits[slot];
664         /* Usually most lengths are short and have no extra bits.  */
665         if (num_extra_bits)
666                 length += lzms_read_bits(is, num_extra_bits);
667         return length;
668 }
669
670 static inline u32
671 lzms_decode_delta_offset(struct lzms_decompressor *d,
672                          struct lzms_input_bitstream *is)
673 {
674         unsigned slot = lzms_decode_huffman_symbol(is,
675                                                    d->delta_offset_decode_table,
676                                                    LZMS_DELTA_OFFSET_TABLEBITS,
677                                                    d->delta_offset_freqs,
678                                                    &d->delta_offset_rebuild_info);
679         return lzms_offset_slot_base[slot] +
680                lzms_read_bits(is, lzms_extra_offset_bits[slot]);
681 }
682
683 static inline unsigned
684 lzms_decode_delta_power(struct lzms_decompressor *d,
685                         struct lzms_input_bitstream *is)
686 {
687         return lzms_decode_huffman_symbol(is,
688                                           d->delta_power_decode_table,
689                                           LZMS_DELTA_POWER_TABLEBITS,
690                                           d->delta_power_freqs,
691                                           &d->delta_power_rebuild_info);
692 }
693
694 static int
695 lzms_create_decompressor(size_t max_bufsize, void **d_ret)
696 {
697         struct lzms_decompressor *d;
698
699         if (max_bufsize > LZMS_MAX_BUFFER_SIZE)
700                 return WIMLIB_ERR_INVALID_PARAM;
701
702         d = ALIGNED_MALLOC(sizeof(struct lzms_decompressor),
703                            DECODE_TABLE_ALIGNMENT);
704         if (!d)
705                 return WIMLIB_ERR_NOMEM;
706
707         *d_ret = d;
708         return 0;
709 }
710
711 /*
712  * Decompress @in_nbytes bytes of LZMS-compressed data at @in and write the
713  * uncompressed data, which had original size @out_nbytes, to @out.  Return 0 if
714  * successful or -1 if the compressed data is invalid.
715  */
716 static int
717 lzms_decompress(const void * const restrict in, const size_t in_nbytes,
718                 void * const restrict out, const size_t out_nbytes,
719                 void * const restrict _d)
720 {
721         struct lzms_decompressor *d = _d;
722         u8 *out_next = out;
723         u8 * const out_end = out + out_nbytes;
724         struct lzms_range_decoder rd;
725         struct lzms_input_bitstream is;
726
727         /* LRU queues for match sources  */
728         u32 recent_lz_offsets[LZMS_NUM_LZ_REPS + 1];
729         u64 recent_delta_pairs[LZMS_NUM_DELTA_REPS + 1];
730
731         /* Previous item type: 0 = literal, 1 = LZ match, 2 = delta match.
732          * This is used to handle delayed updates of the LRU queues.  Instead of
733          * actually delaying the updates, we can check when decoding each rep
734          * match whether a delayed update needs to be taken into account, and if
735          * so get the match source from slot 'rep_idx + 1' instead of from slot
736          * 'rep_idx'.  */
737         unsigned prev_item_type = 0;
738
739         /* States and probability entries for item type disambiguation  */
740         u32 main_state = 0;
741         u32 match_state = 0;
742         u32 lz_state = 0;
743         u32 delta_state = 0;
744         u32 lz_rep_states[LZMS_NUM_LZ_REP_DECISIONS] = {};
745         u32 delta_rep_states[LZMS_NUM_DELTA_REP_DECISIONS] = {};
746
747         /*
748          * Requirements on the compressed data:
749          *
750          * 1. LZMS-compressed data is a series of 16-bit integers, so the
751          *    compressed data buffer cannot take up an odd number of bytes.
752          * 2. To prevent poor performance on some architectures, we require that
753          *    the compressed data buffer is 2-byte aligned.
754          * 3. There must be at least 4 bytes of compressed data, since otherwise
755          *    we cannot even initialize the range decoder.
756          */
757         if ((in_nbytes & 1) || ((uintptr_t)in & 1) || (in_nbytes < 4))
758                 return -1;
759
760         lzms_range_decoder_init(&rd, in, in_nbytes / sizeof(le16));
761
762         lzms_input_bitstream_init(&is, in, in_nbytes / sizeof(le16));
763
764         lzms_init_probabilities(&d->probs);
765
766         lzms_init_huffman_codes(d, lzms_get_num_offset_slots(out_nbytes));
767
768         for (int i = 0; i < LZMS_NUM_LZ_REPS + 1; i++)
769                 recent_lz_offsets[i] = i + 1;
770
771         for (int i = 0; i < LZMS_NUM_DELTA_REPS + 1; i++)
772                 recent_delta_pairs[i] = i + 1;
773
774         /* Main decode loop  */
775         while (out_next != out_end) {
776
777                 if (!lzms_decode_bit(&rd, &main_state,
778                                      LZMS_NUM_MAIN_PROBS, d->probs.main))
779                 {
780                         /* Literal  */
781                         *out_next++ = lzms_decode_literal(d, &is);
782                         prev_item_type = 0;
783
784                 } else if (!lzms_decode_bit(&rd, &match_state,
785                                             LZMS_NUM_MATCH_PROBS,
786                                             d->probs.match))
787                 {
788                         /* LZ match  */
789
790                         u32 offset;
791                         u32 length;
792
793                         STATIC_ASSERT(LZMS_NUM_LZ_REPS == 3);
794
795                         if (!lzms_decode_bit(&rd, &lz_state,
796                                              LZMS_NUM_LZ_PROBS, d->probs.lz))
797                         {
798                                 /* Explicit offset  */
799                                 offset = lzms_decode_lz_offset(d, &is);
800
801                                 recent_lz_offsets[3] = recent_lz_offsets[2];
802                                 recent_lz_offsets[2] = recent_lz_offsets[1];
803                                 recent_lz_offsets[1] = recent_lz_offsets[0];
804                         } else {
805                                 /* Repeat offset  */
806
807                                 if (!lzms_decode_bit(&rd, &lz_rep_states[0],
808                                                      LZMS_NUM_LZ_REP_PROBS,
809                                                      d->probs.lz_rep[0]))
810                                 {
811                                         offset = recent_lz_offsets[0 + (prev_item_type & 1)];
812                                         recent_lz_offsets[0 + (prev_item_type & 1)] = recent_lz_offsets[0];
813                                 } else if (!lzms_decode_bit(&rd, &lz_rep_states[1],
814                                                             LZMS_NUM_LZ_REP_PROBS,
815                                                             d->probs.lz_rep[1]))
816                                 {
817                                         offset = recent_lz_offsets[1 + (prev_item_type & 1)];
818                                         recent_lz_offsets[1 + (prev_item_type & 1)] = recent_lz_offsets[1];
819                                         recent_lz_offsets[1] = recent_lz_offsets[0];
820                                 } else {
821                                         offset = recent_lz_offsets[2 + (prev_item_type & 1)];
822                                         recent_lz_offsets[2 + (prev_item_type & 1)] = recent_lz_offsets[2];
823                                         recent_lz_offsets[2] = recent_lz_offsets[1];
824                                         recent_lz_offsets[1] = recent_lz_offsets[0];
825                                 }
826                         }
827                         recent_lz_offsets[0] = offset;
828                         prev_item_type = 1;
829
830                         length = lzms_decode_length(d, &is);
831
832                         if (unlikely(length > out_end - out_next))
833                                 return -1;
834                         if (unlikely(offset > out_next - (u8 *)out))
835                                 return -1;
836
837                         lz_copy(out_next, length, offset, out_end, LZMS_MIN_MATCH_LENGTH);
838                         out_next += length;
839                 } else {
840                         /* Delta match  */
841
842                         /* (See beginning of file for more information.)  */
843
844                         u32 power;
845                         u32 raw_offset;
846                         u32 span;
847                         u32 offset;
848                         const u8 *matchptr;
849                         u32 length;
850                         u64 pair;
851
852                         STATIC_ASSERT(LZMS_NUM_DELTA_REPS == 3);
853
854                         if (!lzms_decode_bit(&rd, &delta_state,
855                                              LZMS_NUM_DELTA_PROBS,
856                                              d->probs.delta))
857                         {
858                                 /* Explicit offset  */
859                                 power = lzms_decode_delta_power(d, &is);
860                                 raw_offset = lzms_decode_delta_offset(d, &is);
861
862                                 pair = ((u64)power << 32) | raw_offset;
863                                 recent_delta_pairs[3] = recent_delta_pairs[2];
864                                 recent_delta_pairs[2] = recent_delta_pairs[1];
865                                 recent_delta_pairs[1] = recent_delta_pairs[0];
866                         } else {
867                                 if (!lzms_decode_bit(&rd, &delta_rep_states[0],
868                                                      LZMS_NUM_DELTA_REP_PROBS,
869                                                      d->probs.delta_rep[0]))
870                                 {
871                                         pair = recent_delta_pairs[0 + (prev_item_type >> 1)];
872                                         recent_delta_pairs[0 + (prev_item_type >> 1)] = recent_delta_pairs[0];
873                                 } else if (!lzms_decode_bit(&rd, &delta_rep_states[1],
874                                                             LZMS_NUM_DELTA_REP_PROBS,
875                                                             d->probs.delta_rep[1]))
876                                 {
877                                         pair = recent_delta_pairs[1 + (prev_item_type >> 1)];
878                                         recent_delta_pairs[1 + (prev_item_type >> 1)] = recent_delta_pairs[1];
879                                         recent_delta_pairs[1] = recent_delta_pairs[0];
880                                 } else {
881                                         pair = recent_delta_pairs[2 + (prev_item_type >> 1)];
882                                         recent_delta_pairs[2 + (prev_item_type >> 1)] = recent_delta_pairs[2];
883                                         recent_delta_pairs[2] = recent_delta_pairs[1];
884                                         recent_delta_pairs[1] = recent_delta_pairs[0];
885                                 }
886
887                                 power = pair >> 32;
888                                 raw_offset = (u32)pair;
889                         }
890                         recent_delta_pairs[0] = pair;
891                         prev_item_type = 2;
892
893                         length = lzms_decode_length(d, &is);
894
895                         span = (u32)1 << power;
896                         offset = raw_offset << power;
897
898                         /* raw_offset<<power overflows?  */
899                         if (unlikely(offset >> power != raw_offset))
900                                 return -1;
901
902                         /* offset+span overflows?  */
903                         if (unlikely(offset + span < offset))
904                                 return -1;
905
906                         /* buffer underrun?  */
907                         if (unlikely(offset + span > out_next - (u8 *)out))
908                                 return -1;
909
910                         /* buffer overrun?  */
911                         if (unlikely(length > out_end - out_next))
912                                 return -1;
913
914                         matchptr = out_next - offset;
915                         do {
916                                 *out_next = *matchptr + *(out_next - span) -
917                                             *(matchptr - span);
918                                 out_next++;
919                                 matchptr++;
920                         } while (--length);
921                 }
922         }
923
924         lzms_x86_filter(out, out_nbytes, d->last_target_usages, true);
925         return 0;
926 }
927
928 static void
929 lzms_free_decompressor(void *_d)
930 {
931         struct lzms_decompressor *d = _d;
932
933         ALIGNED_FREE(d);
934 }
935
936 const struct decompressor_ops lzms_decompressor_ops = {
937         .create_decompressor  = lzms_create_decompressor,
938         .decompress           = lzms_decompress,
939         .free_decompressor    = lzms_free_decompressor,
940 };