6cc6dab326e795484ed678a3716cf06bff77efe0
[wimlib] / src / lzms-decompress.c
1 /*
2  * lzms-decompress.c
3  */
4
5 /*
6  * Copyright (C) 2013 Eric Biggers
7  *
8  * This file is part of wimlib, a library for working with WIM files.
9  *
10  * wimlib is free software; you can redistribute it and/or modify it under the
11  * terms of the GNU General Public License as published by the Free
12  * Software Foundation; either version 3 of the License, or (at your option)
13  * any later version.
14  *
15  * wimlib is distributed in the hope that it will be useful, but WITHOUT ANY
16  * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
17  * A PARTICULAR PURPOSE. See the GNU General Public License for more
18  * details.
19  *
20  * You should have received a copy of the GNU General Public License
21  * along with wimlib; 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
35  * is a 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  * range-encoded data, whereas bits read from the backwards bitstream constitute
43  * Huffman-encoded symbols or verbatim bits.  For both bitstreams, the ordering
44  * of the bits within the 16-bit coding units is such that the first bit is the
45  * high-order bit and the last bit is the low-order bit.
46  *
47  * From these two logical bitstreams, an LZMS decompressor can reconstitute the
48  * series of items that make up the LZMS data representation.  Each such item
49  * may be a literal byte or a match.  Matches may be either traditional LZ77
50  * matches or "delta" matches, either of which can have its offset encoded
51  * explicitly or encoded via a reference to a recently used (repeat) offset.
52  *
53  * A traditional LZ match consists of a length and offset; it asserts that the
54  * sequence of bytes beginning at the current position and extending for the
55  * length is exactly equal to the equal-length sequence of bytes at the offset
56  * back in the window.  On the other hand, a delta match consists of a length,
57  * raw offset, and power.  It asserts that the sequence of bytes beginning at
58  * the current position and extending for the length is equal to the bytewise
59  * sum of the two equal-length sequences of bytes (2**power) and (raw_offset *
60  * 2**power) bytes before the current position, minus bytewise the sequence of
61  * bytes beginning at (2**power + raw_offset * 2**power) bytes before the
62  * current position.  Although not generally as useful as traditional LZ
63  * matches, delta matches can be helpful on some types of data.  Both LZ and
64  * delta matches may overlap with the current position; in fact, the minimum
65  * offset is 1, regardless of match length.
66  *
67  * For LZ matches, up to 3 repeat offsets are allowed, similar to some other
68  * LZ-based formats such as LZX and LZMA.  They must updated in an LRU fashion,
69  * except for a quirk: updates to the queue must be delayed by one LZMS item,
70  * except for the removal of a repeat match.  As a result, 4 entries are
71  * actually needed in the queue, even though it is only possible to decode
72  * references to the first 3 at any given time.  The queue must be initialized
73  * to the offsets {1, 2, 3, 4}.
74  *
75  * Repeat delta matches are handled similarly, but for them there are two queues
76  * updated in lock-step: one for powers and one for raw offsets.  The power
77  * queue must be initialized to {0, 0, 0, 0}, and the raw offset queue must be
78  * initialized to {1, 2, 3, 4}.
79  *
80  * Bits from the range decoder must be used to disambiguate item types.  The
81  * range decoder must hold two state variables: the range, which must initially
82  * be set to 0xffffffff, and the current code, which must initially be set to
83  * the first 32 bits read from the forwards bitstream.  The range must be
84  * maintained above 0xffff; when it falls below 0xffff, both the range and code
85  * must be left-shifted by 16 bits and the low 16 bits of the code must be
86  * filled in with the next 16 bits from the forwards bitstream.
87  *
88  * To decode each bit, the range decoder requires a probability that is
89  * logically a real number between 0 and 1.  Multiplying this probability by the
90  * current range and taking the floor gives the bound between the 0-bit region
91  * of the range and the 1-bit region of the range.  However, in LZMS,
92  * probabilities are restricted to values of n/64 where n is an integer is
93  * between 1 and 63 inclusively, so the implementation may use integer
94  * operations instead.  Following calculation of the bound, if the current code
95  * is in the 0-bit region, the new range becomes the current code and the
96  * decoded bit is 0; otherwise, the bound must be subtracted from both the range
97  * and the code, and the decoded bit is 1.  More information about range coding
98  * can be found at https://en.wikipedia.org/wiki/Range_encoding.  Furthermore,
99  * note that the LZMA format also uses range coding and has public domain code
100  * available for it.
101  *
102  * The probability used to range-decode each bit must be taken from a table, of
103  * which one instance must exist for each distinct context in which a
104  * range-decoded bit is needed.  At each call of the range decoder, the
105  * appropriate probability must be obtained by indexing the appropriate
106  * probability table with the last 4 (in the context disambiguating literals
107  * from matches), 5 (in the context disambiguating LZ matches from delta
108  * matches), or 6 (in all other contexts) bits recently range-decoded in that
109  * context, ordered such that the most recently decoded bit is the low-order bit
110  * of the index.
111  *
112  * Furthermore, each probability entry itself is variable, as its value must be
113  * maintained as n/64 where n is the number of 0 bits in the most recently
114  * decoded 64 bits with that same entry.  This allows the compressed
115  * representation to adapt to the input and use fewer bits to represent the most
116  * likely data; note that LZMA uses a similar scheme.  Initially, the most
117  * recently 64 decoded bits for each probability entry are assumed to be
118  * 0x0000000055555555 (high order to low order); therefore, all probabilities
119  * are initially 48/64.  During the course of decoding, each probability may be
120  * updated to as low as 0/64 (as a result of reading many consecutive 1 bits
121  * with that entry) or as high as 64/64 (as a result of reading many consecutive
122  * 0 bits with that entry); however, probabilities of 0/64 and 64/64 cannot be
123  * used as-is but rather must be adjusted to 1/64 and 63/64, respectively,
124  * before being used for range decoding.
125  *
126  * Representations of the LZMS items themselves must be read from the backwards
127  * bitstream.  For this, there are 5 different Huffman codes used:
128  *
129  *  - The literal code, used for decoding literal bytes.  Each of the 256
130  *    symbols represents a literal byte.  This code must be rebuilt whenever
131  *    1024 symbols have been decoded with it.
132  *
133  *  - The LZ offset code, used for decoding the offsets of standard LZ77
134  *    matches.  Each symbol represents a position slot, which corresponds to a
135  *    base value and some number of extra bits which must be read and added to
136  *    the base value to reconstitute the full offset.  The number of symbols in
137  *    this code is the number of position slots needed to represent all possible
138  *    offsets in the uncompressed block.  This code must be rebuilt whenever
139  *    1024 symbols have been decoded with it.
140  *
141  *  - The length code, used for decoding length symbols.  Each of the 54 symbols
142  *    represents a length slot, which corresponds to a base value and some
143  *    number of extra bits which must be read and added to the base value to
144  *    reconstitute the full length.  This code must be rebuilt whenever 512
145  *    symbols have been decoded with it.
146  *
147  *  - The delta offset code, used for decoding the offsets of delta matches.
148  *    Each symbol corresponds to a position slot, which corresponds to a base
149  *    value and some number of extra bits which must be read and added to the
150  *    base value to reconstitute the full offset.  The number of symbols in this
151  *    code is equal to the number of symbols in the LZ offset code.  This code
152  *    must be rebuilt whenever 1024 symbols have been decoded with it.
153  *
154  *  - The delta power code, used for decoding the powers of delta matches.  Each
155  *    of the 8 symbols corresponds to a power.  This code must be rebuilt
156  *    whenever 512 symbols have been decoded with it.
157  *
158  * All the LZMS Huffman codes must be built adaptively based on symbol
159  * frequencies.  Initially, each code must be built assuming that all symbols
160  * have equal frequency.  Following that, each code must be rebuilt whenever a
161  * certain number of symbols has been decoded with it.
162  *
163  * In general, multiple valid Huffman codes can be constructed from a set of
164  * symbol frequencies.  Like other compression formats such as XPRESS, LZX, and
165  * DEFLATE, the LZMS format solves this ambiguity by requiring that all Huffman
166  * codes be constructed in canonical form.  This form requires that same-length
167  * codewords be lexicographically ordered the same way as the corresponding
168  * symbols and that all shorter codewords lexicographically precede longer
169  * codewords.
170  *
171  * Codewords in all the LZMS Huffman codes are limited to 15 bits.  If the
172  * canonical code for a given set of symbol frequencies has any codewords longer
173  * than 15 bits, then all frequencies must be divided by 2, rounding up, and the
174  * code construction must be attempted again.
175  *
176  * An LZMS-compressed block seemingly cannot have a compressed size greater than
177  * or equal to the uncompressed size.  In such cases the block must be stored
178  * uncompressed.
179  *
180  * After all LZMS items have been decoded, the data must be postprocessed to
181  * translate absolute address encoded in x86 instructions into their original
182  * relative addresses.
183  *
184  * Details omitted above can be found in the code.  Note that in the absence of
185  * an official specification there is no guarantee that this decompressor
186  * handles all possible cases.
187  */
188
189 #ifdef HAVE_CONFIG_H
190 #  include "config.h"
191 #endif
192
193 #include "wimlib.h"
194 #include "wimlib/compress_common.h"
195 #include "wimlib/decompressor_ops.h"
196 #include "wimlib/decompress_common.h"
197 #include "wimlib/lzms.h"
198 #include "wimlib/util.h"
199
200 #include <limits.h>
201
202 #define LZMS_DECODE_TABLE_BITS  10
203
204 /* Structure used for range decoding, reading bits forwards.  This is the first
205  * logical bitstream mentioned above.  */
206 struct lzms_range_decoder_raw {
207         /* The relevant part of the current range.  Although the logical range
208          * for range decoding is a very large integer, only a small portion
209          * matters at any given time, and it can be normalized (shifted left)
210          * whenever it gets too small.  */
211         u32 range;
212
213         /* The current position in the range encoded by the portion of the input
214          * read so far.  */
215         u32 code;
216
217         /* Pointer to the next little-endian 16-bit integer in the compressed
218          * input data (reading forwards).  */
219         const le16 *in;
220
221         /* Number of 16-bit integers remaining in the compressed input data
222          * (reading forwards).  */
223         size_t num_le16_remaining;
224 };
225
226 /* Structure used for reading raw bits backwards.  This is the second logical
227  * bitstream mentioned above.  */
228 struct lzms_input_bitstream {
229         /* Holding variable for bits that have been read from the compressed
230          * data.  The bits are ordered from high-order to low-order.  */
231         /* XXX:  Without special-case code to handle reading more than 17 bits
232          * at a time, this needs to be 64 bits rather than 32 bits.  */
233         u64 bitbuf;
234
235         /* Number of bits in @bitbuf that are used.  */
236         unsigned num_filled_bits;
237
238         /* Pointer to the one past the next little-endian 16-bit integer in the
239          * compressed input data (reading backwards).  */
240         const le16 *in;
241
242         /* Number of 16-bit integers remaining in the compressed input data
243          * (reading backwards).  */
244         size_t num_le16_remaining;
245 };
246
247 /* Structure used for range decoding.  This wraps around `struct
248  * lzms_range_decoder_raw' to use and maintain probability entries.  */
249 struct lzms_range_decoder {
250         /* Pointer to the raw range decoder, which has no persistent knowledge
251          * of probabilities.  Multiple lzms_range_decoder's share the same
252          * lzms_range_decoder_raw.  */
253         struct lzms_range_decoder_raw *rd;
254
255         /* Bits recently decoded by this range decoder.  This are used as in
256          * index into @prob_entries.  */
257         u32 state;
258
259         /* Bitmask for @state to prevent its value from exceeding the number of
260          * probability entries.  */
261         u32 mask;
262
263         /* Probability entries being used for this range decoder.  */
264         struct lzms_probability_entry prob_entries[LZMS_MAX_NUM_STATES];
265 };
266
267 /* Structure used for Huffman decoding, optionally using the decoded symbols as
268  * slots into a base table to determine how many extra bits need to be read to
269  * reconstitute the full value.  */
270 struct lzms_huffman_decoder {
271
272         /* Bitstream to read Huffman-encoded symbols and verbatim bits from.
273          * Multiple lzms_huffman_decoder's share the same lzms_input_bitstream.
274          */
275         struct lzms_input_bitstream *is;
276
277         /* Pointer to the slot base table to use.  It is indexed by the decoded
278          * Huffman symbol that specifies the slot.  The entry specifies the base
279          * value to use, and the position of its high bit is the number of
280          * additional bits that must be read to reconstitute the full value.
281          *
282          * This member need not be set if only raw Huffman symbols are being
283          * read using this decoder.  */
284         const u32 *slot_base_tab;
285
286         const u8 *extra_bits_tab;
287
288         /* Number of symbols that have been read using this code far.  Reset to
289          * 0 whenever the code is rebuilt.  */
290         u32 num_syms_read;
291
292         /* When @num_syms_read reaches this number, the Huffman code must be
293          * rebuilt.  */
294         u32 rebuild_freq;
295
296         /* Number of symbols in the represented Huffman code.  */
297         unsigned num_syms;
298
299         /* Running totals of symbol frequencies.  These are diluted slightly
300          * whenever the code is rebuilt.  */
301         u32 sym_freqs[LZMS_MAX_NUM_SYMS];
302
303         /* The length, in bits, of each symbol in the Huffman code.  */
304         u8 lens[LZMS_MAX_NUM_SYMS];
305
306         /* The codeword of each symbol in the Huffman code.  */
307         u16 codewords[LZMS_MAX_NUM_SYMS];
308
309         /* A table for quickly decoding symbols encoded using the Huffman code.
310          */
311         u16 decode_table[(1U << LZMS_DECODE_TABLE_BITS) + 2 * LZMS_MAX_NUM_SYMS]
312                                 _aligned_attribute(DECODE_TABLE_ALIGNMENT);
313 };
314
315 /* State of the LZMS decompressor.  */
316 struct lzms_decompressor {
317
318         /* Pointer to the beginning of the uncompressed data buffer.  */
319         u8 *out_begin;
320
321         /* Pointer to the next position in the uncompressed data buffer.  */
322         u8 *out_next;
323
324         /* Pointer to one past the end of the uncompressed data buffer.  */
325         u8 *out_end;
326
327         /* Range decoder, which reads bits from the beginning of the compressed
328          * block, going forwards.  */
329         struct lzms_range_decoder_raw rd;
330
331         /* Input bitstream, which reads from the end of the compressed block,
332          * going backwards.  */
333         struct lzms_input_bitstream is;
334
335         /* Range decoders.  */
336         struct lzms_range_decoder main_range_decoder;
337         struct lzms_range_decoder match_range_decoder;
338         struct lzms_range_decoder lz_match_range_decoder;
339         struct lzms_range_decoder lz_repeat_match_range_decoders[LZMS_NUM_RECENT_OFFSETS - 1];
340         struct lzms_range_decoder delta_match_range_decoder;
341         struct lzms_range_decoder delta_repeat_match_range_decoders[LZMS_NUM_RECENT_OFFSETS - 1];
342
343         /* Huffman decoders.  */
344         struct lzms_huffman_decoder literal_decoder;
345         struct lzms_huffman_decoder lz_offset_decoder;
346         struct lzms_huffman_decoder length_decoder;
347         struct lzms_huffman_decoder delta_power_decoder;
348         struct lzms_huffman_decoder delta_offset_decoder;
349
350         /* LRU (least-recently-used) queues for match information.  */
351         struct lzms_lru_queues lru;
352
353         /* Used for postprocessing.  */
354         s32 last_target_usages[65536];
355 };
356
357 /* Initialize the input bitstream @is to read forwards from the specified
358  * compressed data buffer @in that is @in_limit 16-bit integers long.  */
359 static void
360 lzms_input_bitstream_init(struct lzms_input_bitstream *is,
361                           const le16 *in, size_t in_limit)
362 {
363         is->bitbuf = 0;
364         is->num_filled_bits = 0;
365         is->in = in + in_limit;
366         is->num_le16_remaining = in_limit;
367 }
368
369 /* Ensures that @num_bits bits are buffered in the input bitstream.  */
370 static int
371 lzms_input_bitstream_ensure_bits(struct lzms_input_bitstream *is,
372                                  unsigned num_bits)
373 {
374         while (is->num_filled_bits < num_bits) {
375                 u64 next;
376
377                 LZMS_ASSERT(is->num_filled_bits + 16 <= sizeof(is->bitbuf) * 8);
378
379                 if (unlikely(is->num_le16_remaining == 0))
380                         return -1;
381
382                 next = le16_to_cpu(*--is->in);
383                 is->num_le16_remaining--;
384
385                 is->bitbuf |= next << (sizeof(is->bitbuf) * 8 - is->num_filled_bits - 16);
386                 is->num_filled_bits += 16;
387         }
388         return 0;
389
390 }
391
392 /* Returns the next @num_bits bits that are buffered in the input bitstream.  */
393 static u32
394 lzms_input_bitstream_peek_bits(struct lzms_input_bitstream *is,
395                                unsigned num_bits)
396 {
397         LZMS_ASSERT(is->num_filled_bits >= num_bits);
398         return is->bitbuf >> (sizeof(is->bitbuf) * 8 - num_bits);
399 }
400
401 /* Removes the next @num_bits bits that are buffered in the input bitstream.  */
402 static void
403 lzms_input_bitstream_remove_bits(struct lzms_input_bitstream *is,
404                                  unsigned num_bits)
405 {
406         LZMS_ASSERT(is->num_filled_bits >= num_bits);
407         is->bitbuf <<= num_bits;
408         is->num_filled_bits -= num_bits;
409 }
410
411 /* Removes and returns the next @num_bits bits that are buffered in the input
412  * bitstream.  */
413 static u32
414 lzms_input_bitstream_pop_bits(struct lzms_input_bitstream *is,
415                               unsigned num_bits)
416 {
417         u32 bits = lzms_input_bitstream_peek_bits(is, num_bits);
418         lzms_input_bitstream_remove_bits(is, num_bits);
419         return bits;
420 }
421
422 /* Reads the next @num_bits from the input bitstream.  */
423 static u32
424 lzms_input_bitstream_read_bits(struct lzms_input_bitstream *is,
425                                unsigned num_bits)
426 {
427         if (unlikely(lzms_input_bitstream_ensure_bits(is, num_bits)))
428                 return 0;
429         return lzms_input_bitstream_pop_bits(is, num_bits);
430 }
431
432 /* Initialize the range decoder @rd to read forwards from the specified
433  * compressed data buffer @in that is @in_limit 16-bit integers long.  */
434 static void
435 lzms_range_decoder_raw_init(struct lzms_range_decoder_raw *rd,
436                             const le16 *in, size_t in_limit)
437 {
438         rd->range = 0xffffffff;
439         rd->code = ((u32)le16_to_cpu(in[0]) << 16) |
440                    ((u32)le16_to_cpu(in[1]) <<  0);
441         rd->in = in + 2;
442         rd->num_le16_remaining = in_limit - 2;
443 }
444
445 /* Ensures the current range of the range decoder has at least 16 bits of
446  * precision.  */
447 static int
448 lzms_range_decoder_raw_normalize(struct lzms_range_decoder_raw *rd)
449 {
450         if (rd->range <= 0xffff) {
451                 rd->range <<= 16;
452                 if (unlikely(rd->num_le16_remaining == 0))
453                         return -1;
454                 rd->code = (rd->code << 16) | le16_to_cpu(*rd->in++);
455                 rd->num_le16_remaining--;
456         }
457         return 0;
458 }
459
460 /* Decode and return the next bit from the range decoder (raw version).
461  *
462  * @prob is the chance out of LZMS_PROBABILITY_MAX that the next bit is 0.
463  */
464 static int
465 lzms_range_decoder_raw_decode_bit(struct lzms_range_decoder_raw *rd, u32 prob)
466 {
467         u32 bound;
468
469         /* Ensure the range has at least 16 bits of precision.  */
470         lzms_range_decoder_raw_normalize(rd);
471
472         /* Based on the probability, calculate the bound between the 0-bit
473          * region and the 1-bit region of the range.  */
474         bound = (rd->range >> LZMS_PROBABILITY_BITS) * prob;
475
476         if (rd->code < bound) {
477                 /* Current code is in the 0-bit region of the range.  */
478                 rd->range = bound;
479                 return 0;
480         } else {
481                 /* Current code is in the 1-bit region of the range.  */
482                 rd->range -= bound;
483                 rd->code -= bound;
484                 return 1;
485         }
486 }
487
488 /* Decode and return the next bit from the range decoder.  This wraps around
489  * lzms_range_decoder_raw_decode_bit() to handle using and updating the
490  * appropriate probability table.  */
491 static int
492 lzms_range_decode_bit(struct lzms_range_decoder *dec)
493 {
494         struct lzms_probability_entry *prob_entry;
495         u32 prob;
496         int bit;
497
498         /* Load the probability entry corresponding to the current state.  */
499         prob_entry = &dec->prob_entries[dec->state];
500
501         /* Treat the number of zero bits in the most recently decoded
502          * LZMS_PROBABILITY_MAX bits with this probability entry as the chance,
503          * out of LZMS_PROBABILITY_MAX, that the next bit will be a 0.  However,
504          * don't allow 0% or 100% probabilities.  */
505         prob = prob_entry->num_recent_zero_bits;
506         if (prob == LZMS_PROBABILITY_MAX)
507                 prob = LZMS_PROBABILITY_MAX - 1;
508         else if (prob == 0)
509                 prob = 1;
510
511         /* Decode the next bit.  */
512         bit = lzms_range_decoder_raw_decode_bit(dec->rd, prob);
513
514         /* Update the state based on the newly decoded bit.  */
515         dec->state = (((dec->state << 1) | bit) & dec->mask);
516
517         /* Update the recent bits, including the cached count of 0's.  */
518         BUILD_BUG_ON(LZMS_PROBABILITY_MAX > sizeof(prob_entry->recent_bits) * 8);
519         if (bit == 0) {
520                 if (prob_entry->recent_bits & (1ULL << (LZMS_PROBABILITY_MAX - 1))) {
521                         /* Replacing 1 bit with 0 bit; increment the zero count.
522                          */
523                         prob_entry->num_recent_zero_bits++;
524                 }
525         } else {
526                 if (!(prob_entry->recent_bits & (1ULL << (LZMS_PROBABILITY_MAX - 1)))) {
527                         /* Replacing 0 bit with 1 bit; decrement the zero count.
528                          */
529                         prob_entry->num_recent_zero_bits--;
530                 }
531         }
532         prob_entry->recent_bits = (prob_entry->recent_bits << 1) | bit;
533
534         /* Return the decoded bit.  */
535         return bit;
536 }
537
538
539 /* Build the decoding table for a new adaptive Huffman code using the alphabet
540  * used in the specified Huffman decoder, with the symbol frequencies
541  * dec->sym_freqs.  */
542 static void
543 lzms_rebuild_adaptive_huffman_code(struct lzms_huffman_decoder *dec)
544 {
545
546         /* XXX:  This implementation makes use of code already implemented for
547          * the XPRESS and LZX compression formats.  However, since for the
548          * adaptive codes used in LZMS we don't actually need the explicit codes
549          * themselves, only the decode tables, it may be possible to optimize
550          * this by somehow directly building or updating the Huffman decode
551          * table.  This may be a worthwhile optimization because the adaptive
552          * codes change many times throughout a decompression run.  */
553         LZMS_DEBUG("Rebuilding adaptive Huffman code (num_syms=%u)",
554                    dec->num_syms);
555         make_canonical_huffman_code(dec->num_syms, LZMS_MAX_CODEWORD_LEN,
556                                     dec->sym_freqs, dec->lens, dec->codewords);
557 #if defined(ENABLE_LZMS_DEBUG)
558         int ret =
559 #endif
560         make_huffman_decode_table(dec->decode_table, dec->num_syms,
561                                   LZMS_DECODE_TABLE_BITS, dec->lens,
562                                   LZMS_MAX_CODEWORD_LEN);
563         LZMS_ASSERT(ret == 0);
564 }
565
566 /* Decode and return the next Huffman-encoded symbol from the LZMS-compressed
567  * block using the specified Huffman decoder.  */
568 static u32
569 lzms_huffman_decode_symbol(struct lzms_huffman_decoder *dec)
570 {
571         const u16 *decode_table = dec->decode_table;
572         struct lzms_input_bitstream *is = dec->is;
573         u16 entry;
574         u16 key_bits;
575         u16 sym;
576
577         /* The Huffman codes used in LZMS are adaptive and must be rebuilt
578          * whenever a certain number of symbols have been read.  Each such
579          * rebuild uses the current symbol frequencies, but the format also
580          * requires that the symbol frequencies be halved after each code
581          * rebuild.  This diminishes the effect of old symbols on the current
582          * Huffman codes, thereby causing the Huffman codes to be more locally
583          * adaptable.  */
584         if (dec->num_syms_read == dec->rebuild_freq) {
585                 lzms_rebuild_adaptive_huffman_code(dec);
586                 for (unsigned i = 0; i < dec->num_syms; i++) {
587                         dec->sym_freqs[i] >>= 1;
588                         dec->sym_freqs[i] += 1;
589                 }
590                 dec->num_syms_read = 0;
591         }
592
593         /* XXX: Copied from read_huffsym() (decompress_common.h), since this
594          * uses a different input bitstream type.  Should unify the
595          * implementations.  */
596         lzms_input_bitstream_ensure_bits(is, LZMS_MAX_CODEWORD_LEN);
597
598         /* Index the decode table by the next table_bits bits of the input.  */
599         key_bits = lzms_input_bitstream_peek_bits(is, LZMS_DECODE_TABLE_BITS);
600         entry = decode_table[key_bits];
601         if (likely(entry < 0xC000)) {
602                 /* Fast case: The decode table directly provided the symbol and
603                  * codeword length.  The low 11 bits are the symbol, and the
604                  * high 5 bits are the codeword length.  */
605                 lzms_input_bitstream_remove_bits(is, entry >> 11);
606                 sym = entry & 0x7FF;
607         } else {
608                 /* Slow case: The codeword for the symbol is longer than
609                  * table_bits, so the symbol does not have an entry directly in
610                  * the first (1 << table_bits) entries of the decode table.
611                  * Traverse the appropriate binary tree bit-by-bit in order to
612                  * decode the symbol.  */
613                 lzms_input_bitstream_remove_bits(is, LZMS_DECODE_TABLE_BITS);
614                 do {
615                         key_bits = (entry & 0x3FFF) + lzms_input_bitstream_pop_bits(is, 1);
616                 } while ((entry = decode_table[key_bits]) >= 0xC000);
617                 sym = entry;
618         }
619
620         /* Tally and return the decoded symbol.  */
621         ++dec->sym_freqs[sym];
622         ++dec->num_syms_read;
623         return sym;
624 }
625
626 /* Decode a number from the LZMS bitstream, encoded as a Huffman-encoded symbol
627  * specifying a "slot" (whose corresponding value is looked up in a static
628  * table) plus the number specified by a number of extra bits depending on the
629  * slot.  */
630 static u32
631 lzms_decode_value(struct lzms_huffman_decoder *dec)
632 {
633         unsigned slot;
634         unsigned num_extra_bits;
635         u32 extra_bits;
636
637         LZMS_ASSERT(dec->slot_base_tab != NULL);
638         LZMS_ASSERT(dec->extra_bits_tab != NULL);
639
640         /* Read the slot (position slot, length slot, etc.), which is encoded as
641          * a Huffman symbol.  */
642         slot = lzms_huffman_decode_symbol(dec);
643
644         /* Get the number of extra bits needed to represent the range of values
645          * that share the slot.  */
646         num_extra_bits = dec->extra_bits_tab[slot];
647
648         /* Read the number of extra bits and add them to the slot base to form
649          * the final decoded value.  */
650         extra_bits = lzms_input_bitstream_read_bits(dec->is, num_extra_bits);
651         return dec->slot_base_tab[slot] + extra_bits;
652 }
653
654 /* Copy a literal to the output buffer.  */
655 static int
656 lzms_copy_literal(struct lzms_decompressor *ctx, u8 literal)
657 {
658         *ctx->out_next++ = literal;
659         return 0;
660 }
661
662 /* Validate an LZ match and copy it to the output buffer.  */
663 static int
664 lzms_copy_lz_match(struct lzms_decompressor *ctx, u32 length, u32 offset)
665 {
666         u8 *out_next;
667         u8 *matchptr;
668
669         if (length > ctx->out_end - ctx->out_next) {
670                 LZMS_DEBUG("Match overrun!");
671                 return -1;
672         }
673         if (offset > ctx->out_next - ctx->out_begin) {
674                 LZMS_DEBUG("Match underrun!");
675                 return -1;
676         }
677
678         out_next = ctx->out_next;
679         matchptr = out_next - offset;
680         while (length--)
681                 *out_next++ = *matchptr++;
682
683         ctx->out_next = out_next;
684         return 0;
685 }
686
687 /* Validate a delta match and copy it to the output buffer.  */
688 static int
689 lzms_copy_delta_match(struct lzms_decompressor *ctx, u32 length,
690                       u32 power, u32 raw_offset)
691 {
692         u32 offset1 = 1U << power;
693         u32 offset2 = raw_offset << power;
694         u32 offset = offset1 + offset2;
695         u8 *out_next;
696         u8 *matchptr1;
697         u8 *matchptr2;
698         u8 *matchptr;
699
700         if (length > ctx->out_end - ctx->out_next) {
701                 LZMS_DEBUG("Match overrun!");
702                 return -1;
703         }
704         if (offset > ctx->out_next - ctx->out_begin) {
705                 LZMS_DEBUG("Match underrun!");
706                 return -1;
707         }
708
709         out_next = ctx->out_next;
710         matchptr1 = out_next - offset1;
711         matchptr2 = out_next - offset2;
712         matchptr = out_next - offset;
713
714         while (length--)
715                 *out_next++ = *matchptr1++ + *matchptr2++ - *matchptr++;
716
717         ctx->out_next = out_next;
718         return 0;
719 }
720
721 /* Decode a (length, offset) pair from the input.  */
722 static int
723 lzms_decode_lz_match(struct lzms_decompressor *ctx)
724 {
725         int bit;
726         u32 length, offset;
727
728         /* Decode the match offset.  The next range-encoded bit indicates
729          * whether it's a repeat offset or an explicit offset.  */
730
731         bit = lzms_range_decode_bit(&ctx->lz_match_range_decoder);
732         if (bit == 0) {
733                 /* Explicit offset.  */
734                 offset = lzms_decode_value(&ctx->lz_offset_decoder);
735         } else {
736                 /* Repeat offset.  */
737                 int i;
738
739                 for (i = 0; i < LZMS_NUM_RECENT_OFFSETS - 1; i++)
740                         if (!lzms_range_decode_bit(&ctx->lz_repeat_match_range_decoders[i]))
741                                 break;
742
743                 offset = ctx->lru.lz.recent_offsets[i];
744
745                 for (; i < LZMS_NUM_RECENT_OFFSETS; i++)
746                         ctx->lru.lz.recent_offsets[i] = ctx->lru.lz.recent_offsets[i + 1];
747         }
748
749         /* Decode match length, which is always given explicitly (there is no
750          * LRU queue for repeat lengths).  */
751         length = lzms_decode_value(&ctx->length_decoder);
752
753         ctx->lru.lz.upcoming_offset = offset;
754
755         LZMS_DEBUG("Decoded %s LZ match: length=%u, offset=%u",
756                    (bit ? "repeat" : "explicit"), length, offset);
757
758         /* Validate the match and copy it to the output.  */
759         return lzms_copy_lz_match(ctx, length, offset);
760 }
761
762 /* Decodes a "delta" match from the input.  */
763 static int
764 lzms_decode_delta_match(struct lzms_decompressor *ctx)
765 {
766         int bit;
767         u32 length, power, raw_offset;
768
769         /* Decode the match power and raw offset.  The next range-encoded bit
770          * indicates whether these data are a repeat, or given explicitly.  */
771
772         bit = lzms_range_decode_bit(&ctx->delta_match_range_decoder);
773         if (bit == 0) {
774                 power = lzms_huffman_decode_symbol(&ctx->delta_power_decoder);
775                 raw_offset = lzms_decode_value(&ctx->delta_offset_decoder);
776         } else {
777                 int i;
778
779                 for (i = 0; i < LZMS_NUM_RECENT_OFFSETS - 1; i++)
780                         if (!lzms_range_decode_bit(&ctx->delta_repeat_match_range_decoders[i]))
781                                 break;
782
783                 power = ctx->lru.delta.recent_powers[i];
784                 raw_offset = ctx->lru.delta.recent_offsets[i];
785
786                 for (; i < LZMS_NUM_RECENT_OFFSETS; i++) {
787                         ctx->lru.delta.recent_powers[i] = ctx->lru.delta.recent_powers[i + 1];
788                         ctx->lru.delta.recent_offsets[i] = ctx->lru.delta.recent_offsets[i + 1];
789                 }
790         }
791
792         length = lzms_decode_value(&ctx->length_decoder);
793
794         ctx->lru.delta.upcoming_power = power;
795         ctx->lru.delta.upcoming_offset = raw_offset;
796
797         LZMS_DEBUG("Decoded %s delta match: length=%u, power=%u, raw_offset=%u",
798                    (bit ? "repeat" : "explicit"), length, power, raw_offset);
799
800         /* Validate the match and copy it to the output.  */
801         return lzms_copy_delta_match(ctx, length, power, raw_offset);
802 }
803
804 /* Decode an LZ or delta match.  */
805 static int
806 lzms_decode_match(struct lzms_decompressor *ctx)
807 {
808         if (!lzms_range_decode_bit(&ctx->match_range_decoder))
809                 return lzms_decode_lz_match(ctx);
810         else
811                 return lzms_decode_delta_match(ctx);
812 }
813
814 /* Decode a literal byte encoded using the literal Huffman code.  */
815 static int
816 lzms_decode_literal(struct lzms_decompressor *ctx)
817 {
818         u8 literal = lzms_huffman_decode_symbol(&ctx->literal_decoder);
819         LZMS_DEBUG("Decoded literal: 0x%02x", literal);
820         return lzms_copy_literal(ctx, literal);
821 }
822
823 /* Decode the next LZMS match or literal.  */
824 static int
825 lzms_decode_item(struct lzms_decompressor *ctx)
826 {
827         int ret;
828
829         ctx->lru.lz.upcoming_offset = 0;
830         ctx->lru.delta.upcoming_power = 0;
831         ctx->lru.delta.upcoming_offset = 0;
832
833         if (lzms_range_decode_bit(&ctx->main_range_decoder))
834                 ret = lzms_decode_match(ctx);
835         else
836                 ret = lzms_decode_literal(ctx);
837
838         if (ret)
839                 return ret;
840
841         lzms_update_lru_queues(&ctx->lru);
842         return 0;
843 }
844
845 static void
846 lzms_init_range_decoder(struct lzms_range_decoder *dec,
847                         struct lzms_range_decoder_raw *rd, u32 num_states)
848 {
849         dec->rd = rd;
850         dec->state = 0;
851         dec->mask = num_states - 1;
852         for (u32 i = 0; i < num_states; i++) {
853                 dec->prob_entries[i].num_recent_zero_bits = LZMS_INITIAL_PROBABILITY;
854                 dec->prob_entries[i].recent_bits = LZMS_INITIAL_RECENT_BITS;
855         }
856 }
857
858 static void
859 lzms_init_huffman_decoder(struct lzms_huffman_decoder *dec,
860                           struct lzms_input_bitstream *is,
861                           const u32 *slot_base_tab,
862                           const u8 *extra_bits_tab,
863                           unsigned num_syms,
864                           unsigned rebuild_freq)
865 {
866         dec->is = is;
867         dec->slot_base_tab = slot_base_tab;
868         dec->extra_bits_tab = extra_bits_tab;
869         dec->num_syms = num_syms;
870         dec->num_syms_read = rebuild_freq;
871         dec->rebuild_freq = rebuild_freq;
872         for (unsigned i = 0; i < num_syms; i++)
873                 dec->sym_freqs[i] = 1;
874 }
875
876 /* Prepare to decode items from an LZMS-compressed block.  */
877 static void
878 lzms_init_decompressor(struct lzms_decompressor *ctx,
879                        const void *cdata, unsigned clen,
880                        void *ubuf, unsigned ulen)
881 {
882         unsigned num_position_slots;
883
884         LZMS_DEBUG("Initializing decompressor (clen=%u, ulen=%u)", clen, ulen);
885
886         /* Initialize output pointers.  */
887         ctx->out_begin = ubuf;
888         ctx->out_next = ubuf;
889         ctx->out_end = (u8*)ubuf + ulen;
890
891         /* Initialize the raw range decoder (reading forwards).  */
892         lzms_range_decoder_raw_init(&ctx->rd, cdata, clen / 2);
893
894         /* Initialize the input bitstream for Huffman symbols (reading
895          * backwards)  */
896         lzms_input_bitstream_init(&ctx->is, cdata, clen / 2);
897
898         /* Calculate the number of position slots needed for this compressed
899          * block.  */
900         num_position_slots = lzms_get_position_slot(ulen - 1) + 1;
901
902         LZMS_DEBUG("Using %u position slots", num_position_slots);
903
904         /* Initialize Huffman decoders for each alphabet used in the compressed
905          * representation.  */
906         lzms_init_huffman_decoder(&ctx->literal_decoder, &ctx->is,
907                                   NULL, NULL, LZMS_NUM_LITERAL_SYMS,
908                                   LZMS_LITERAL_CODE_REBUILD_FREQ);
909
910         lzms_init_huffman_decoder(&ctx->lz_offset_decoder, &ctx->is,
911                                   lzms_position_slot_base,
912                                   lzms_extra_position_bits,
913                                   num_position_slots,
914                                   LZMS_LZ_OFFSET_CODE_REBUILD_FREQ);
915
916         lzms_init_huffman_decoder(&ctx->length_decoder, &ctx->is,
917                                   lzms_length_slot_base,
918                                   lzms_extra_length_bits,
919                                   LZMS_NUM_LEN_SYMS,
920                                   LZMS_LENGTH_CODE_REBUILD_FREQ);
921
922         lzms_init_huffman_decoder(&ctx->delta_offset_decoder, &ctx->is,
923                                   lzms_position_slot_base,
924                                   lzms_extra_position_bits,
925                                   num_position_slots,
926                                   LZMS_DELTA_OFFSET_CODE_REBUILD_FREQ);
927
928         lzms_init_huffman_decoder(&ctx->delta_power_decoder, &ctx->is,
929                                   NULL, NULL, LZMS_NUM_DELTA_POWER_SYMS,
930                                   LZMS_DELTA_POWER_CODE_REBUILD_FREQ);
931
932
933         /* Initialize range decoders, all of which wrap around the same
934          * lzms_range_decoder_raw.  */
935         lzms_init_range_decoder(&ctx->main_range_decoder,
936                                 &ctx->rd, LZMS_NUM_MAIN_STATES);
937
938         lzms_init_range_decoder(&ctx->match_range_decoder,
939                                 &ctx->rd, LZMS_NUM_MATCH_STATES);
940
941         lzms_init_range_decoder(&ctx->lz_match_range_decoder,
942                                 &ctx->rd, LZMS_NUM_LZ_MATCH_STATES);
943
944         for (size_t i = 0; i < ARRAY_LEN(ctx->lz_repeat_match_range_decoders); i++)
945                 lzms_init_range_decoder(&ctx->lz_repeat_match_range_decoders[i],
946                                         &ctx->rd, LZMS_NUM_LZ_REPEAT_MATCH_STATES);
947
948         lzms_init_range_decoder(&ctx->delta_match_range_decoder,
949                                 &ctx->rd, LZMS_NUM_DELTA_MATCH_STATES);
950
951         for (size_t i = 0; i < ARRAY_LEN(ctx->delta_repeat_match_range_decoders); i++)
952                 lzms_init_range_decoder(&ctx->delta_repeat_match_range_decoders[i],
953                                         &ctx->rd, LZMS_NUM_DELTA_REPEAT_MATCH_STATES);
954
955         /* Initialize LRU match information.  */
956         lzms_init_lru_queues(&ctx->lru);
957
958         LZMS_DEBUG("Decompressor successfully initialized");
959 }
960
961 /* Decode the series of literals and matches from the LZMS-compressed data.
962  * Returns 0 on success; nonzero if the compressed data is invalid.  */
963 static int
964 lzms_decode_items(const u8 *cdata, size_t clen, u8 *ubuf, size_t ulen,
965                   struct lzms_decompressor *ctx)
966 {
967         /* Initialize the LZMS decompressor.  */
968         lzms_init_decompressor(ctx, cdata, clen, ubuf, ulen);
969
970         /* Decode the sequence of items.  */
971         while (ctx->out_next != ctx->out_end) {
972                 LZMS_DEBUG("Position %u", ctx->out_next - ctx->out_begin);
973                 if (lzms_decode_item(ctx))
974                         return -1;
975         }
976         return 0;
977 }
978
979 static int
980 lzms_decompress(const void *compressed_data, size_t compressed_size,
981                 void *uncompressed_data, size_t uncompressed_size, void *_ctx)
982 {
983         struct lzms_decompressor *ctx = _ctx;
984
985         /* The range decoder requires that a minimum of 4 bytes of compressed
986          * data be initially available.  */
987         if (compressed_size < 4) {
988                 LZMS_DEBUG("Compressed size too small (got %zu, expected >= 4)",
989                            compressed_size);
990                 return -1;
991         }
992
993         /* An LZMS-compressed data block should be evenly divisible into 16-bit
994          * integers.  */
995         if (compressed_size % 2 != 0) {
996                 LZMS_DEBUG("Compressed size not divisible by 2 (got %zu)",
997                            compressed_size);
998                 return -1;
999         }
1000
1001         /* Handle the trivial case where nothing needs to be decompressed.
1002          * (Necessary because a window of size 0 does not have a valid position
1003          * slot.)  */
1004         if (uncompressed_size == 0)
1005                 return 0;
1006
1007         /* The x86 post-processor requires that the uncompressed length fit into
1008          * a signed 32-bit integer.  Also, the position slot table cannot be
1009          * searched for a position of INT32_MAX or greater.  */
1010         if (uncompressed_size >= INT32_MAX) {
1011                 LZMS_DEBUG("Uncompressed length too large "
1012                            "(got %zu, expected < INT32_MAX)",
1013                            uncompressed_size);
1014                 return -1;
1015         }
1016
1017         /* Decode the literals and matches.  */
1018         if (lzms_decode_items(compressed_data, compressed_size,
1019                               uncompressed_data, uncompressed_size, ctx))
1020                 return -1;
1021
1022         /* Postprocess the data.  */
1023         lzms_x86_filter(uncompressed_data, uncompressed_size,
1024                         ctx->last_target_usages, true);
1025
1026         LZMS_DEBUG("Decompression successful.");
1027         return 0;
1028 }
1029
1030 static void
1031 lzms_free_decompressor(void *_ctx)
1032 {
1033         struct lzms_decompressor *ctx = _ctx;
1034
1035         FREE(ctx);
1036 }
1037
1038 static int
1039 lzms_create_decompressor(size_t max_block_size,
1040                          const struct wimlib_decompressor_params_header *params,
1041                          void **ctx_ret)
1042 {
1043         struct lzms_decompressor *ctx;
1044
1045         ctx = MALLOC(sizeof(struct lzms_decompressor));
1046         if (ctx == NULL)
1047                 return WIMLIB_ERR_NOMEM;
1048
1049         /* Initialize position and length slot data if not done already.  */
1050         lzms_init_slots();
1051
1052         *ctx_ret = ctx;
1053         return 0;
1054 }
1055
1056 const struct decompressor_ops lzms_decompressor_ops = {
1057         .create_decompressor  = lzms_create_decompressor,
1058         .decompress           = lzms_decompress,
1059         .free_decompressor    = lzms_free_decompressor,
1060 };