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