X-Git-Url: https://wimlib.net/git/?p=wimlib;a=blobdiff_plain;f=src%2Flzms-decompress.c;h=32ce082c29b1ebb6a5940a60619d74c7bbfe3d74;hp=c73a3fdd44afc3b2bb53917119ac248693d97d85;hb=abe65b1f2f12bcedec1103cc7924897720f919be;hpb=60523d25f34692d6f3a7c8bbda88eead17f23b12 diff --git a/src/lzms-decompress.c b/src/lzms-decompress.c index c73a3fdd..32ce082c 100644 --- a/src/lzms-decompress.c +++ b/src/lzms-decompress.c @@ -1,7 +1,5 @@ /* * lzms-decompress.c - * - * LZMS decompression routines. */ /* @@ -23,120 +21,1035 @@ * along with wimlib; if not, see http://www.gnu.org/licenses/. */ -#include "wimlib/win32_common.h" +/* + * This is a decompressor for the LZMS compression format used by Microsoft. + * This format is not documented, but it is one of the formats supported by the + * compression API available in Windows 8, and as of Windows 8 it is one of the + * formats that can be used in WIM files. + * + * This decompressor only implements "raw" decompression, which decompresses a + * single LZMS-compressed block. This behavior is the same as that of + * Decompress() in the Windows 8 compression API when using a compression handle + * created with CreateDecompressor() with the Algorithm parameter specified as + * COMPRESS_ALGORITHM_LZMS | COMPRESS_RAW. Presumably, non-raw LZMS data + * is a container format from which the locations and sizes (both compressed and + * uncompressed) of the constituent blocks can be determined. + * + * A LZMS-compressed block must be read in 16-bit little endian units from both + * directions. One logical bitstream starts at the front of the block and + * proceeds forwards. Another logical bitstream starts at the end of the block + * and proceeds backwards. Bits read from the forwards bitstream constitute + * range-encoded data, whereas bits read from the backwards bitstream constitute + * Huffman-encoded symbols or verbatim bits. For both bitstreams, the ordering + * of the bits within the 16-bit coding units is such that the first bit is the + * high-order bit and the last bit is the low-order bit. + * + * From these two logical bitstreams, an LZMS decompressor can reconstitute the + * series of items that make up the LZMS data representation. Each such item + * may be a literal byte or a match. Matches may be either traditional LZ77 + * matches or "delta" matches, either of which can have its offset encoded + * explicitly or encoded via a reference to a recently used (repeat) offset. + * + * A traditional LZ match consists of a length and offset; it asserts that the + * sequence of bytes beginning at the current position and extending for the + * length is exactly equal to the equal-length sequence of bytes at the offset + * back in the window. On the other hand, a delta match consists of a length, + * raw offset, and power. It asserts that the sequence of bytes beginning at + * the current position and extending for the length is equal to the bytewise + * sum of the two equal-length sequences of bytes (2**power) and (raw_offset * + * 2**power) bytes before the current position, minus bytewise the sequence of + * bytes beginning at (2**power + raw_offset * 2**power) bytes before the + * current position. Although not generally as useful as traditional LZ + * matches, delta matches can be helpful on some types of data. Both LZ and + * delta matches may overlap with the current position; in fact, the minimum + * offset is 1, regardless of match length. + * + * For LZ matches, up to 3 repeat offsets are allowed, similar to some other + * LZ-based formats such as LZX and LZMA. They must updated in a LRU fashion, + * except for a quirk: updates to the queue must be delayed by one LZMS item, + * except for the removal of a repeat match. As a result, 4 entries are + * actually needed in the queue, even though it is only possible to decode + * references to the first 3 at any given time. The queue must be initialized + * to the offsets {1, 2, 3, 4}. + * + * Repeat delta matches are handled similarly, but for them there are two queues + * updated in lock-step: one for powers and one for raw offsets. The power + * queue must be initialized to {0, 0, 0, 0}, and the raw offset queue must be + * initialized to {1, 2, 3, 4}. + * + * Bits from the range decoder must be used to disambiguate item types. The + * range decoder must hold two state variables: the range, which must initially + * be set to 0xffffffff, and the current code, which must initially be set to + * the first 32 bits read from the forwards bitstream. The range must be + * maintained above 0xffff; when it falls below 0xffff, both the range and code + * must be left-shifted by 16 bits and the low 16 bits of the code must be + * filled in with the next 16 bits from the forwards bitstream. + * + * To decode each bit, the range decoder requires a probability that is + * logically a real number between 0 and 1. Multiplying this probability by the + * current range and taking the floor gives the bound between the 0-bit region + * of the range and the 1-bit region of the range. However, in LZMS, + * probabilities are restricted to values of n/64 where n is an integer is + * between 1 and 63 inclusively, so the implementation may use integer + * operations instead. Following calculation of the bound, if the current code + * is in the 0-bit region, the new range becomes the current code and the + * decoded bit is 0; otherwise, the bound must be subtracted from both the range + * and the code, and the decoded bit is 1. More information about range coding + * can be found at https://en.wikipedia.org/wiki/Range_encoding. Furthermore, + * note that the LZMA format also uses range coding and has public domain code + * available for it. + * + * The probability used to range-decode each bit must be taken from a table, of + * which one instance must exist for each distinct context in which a + * range-decoded bit is needed. At each call of the range decoder, the + * appropriate probability must be obtained by indexing the appropriate + * probability table with the last 4 (in the context disambiguating literals + * from matches), 5 (in the context disambiguating LZ matches from delta + * matches), or 6 (in all other contexts) bits recently range-decoded in that + * context, ordered such that the most recently decoded bit is the low-order bit + * of the index. + * + * Furthermore, each probability entry itself is variable, as its value must be + * maintained as n/64 where n is the number of 0 bits in the most recently + * decoded 64 bits with that same entry. This allows the compressed + * representation to adapt to the input and use fewer bits to represent the most + * likely data; note that LZMA uses a similar scheme. Initially, the most + * recently 64 decoded bits for each probability entry are assumed to be + * 0x0000000055555555 (high order to low order); therefore, all probabilities + * are initially 48/64. During the course of decoding, each probability may be + * updated to as low as 0/64 (as a result of reading many consecutive 1 bits + * with that entry) or as high as 64/64 (as a result of reading many consecutive + * 0 bits with that entry); however, probabilities of 0/64 and 64/64 cannot be + * used as-is but rather must be adjusted to 1/64 and 63/64, respectively, + * before being used for range decoding. + * + * Representations of the LZMS items themselves must be read from the backwards + * bitstream. For this, there are 5 different Huffman codes used: + * + * - The literal code, used for decoding literal bytes. Each of the 256 + * symbols represents a literal byte. This code must be rebuilt whenever + * 1024 symbols have been decoded with it. + * + * - The LZ offset code, used for decoding the offsets of standard LZ77 + * matches. Each symbol represents a position slot, which corresponds to a + * base value and some number of extra bits which must be read and added to + * the base value to reconstitute the full offset. The number of symbols in + * this code is the number of position slots needed to represent all possible + * offsets in the uncompressed block. This code must be rebuilt whenever + * 1024 symbols have been decoded with it. + * + * - The length code, used for decoding length symbols. Each of the 54 symbols + * represents a length slot, which corresponds to a base value and some + * number of extra bits which must be read and added to the base value to + * reconstitute the full length. This code must be rebuilt whenever 512 + * symbols have been decoded with it. + * + * - The delta offset code, used for decoding the offsets of delta matches. + * Each symbol corresponds to a position slot, which corresponds to a base + * value and some number of extra bits which must be read and added to the + * base value to reconstitute the full offset. The number of symbols in this + * code is equal to the number of symbols in the LZ offset code. This code + * must be rebuilt whenever 1024 symbols have been decoded with it. + * + * - The delta power code, used for decoding the powers of delta matches. Each + * of the 8 symbols corresponds to a power. This code must be rebuilt + * whenever 512 symbols have been decoded with it. + * + * All the LZMS Huffman codes must be built adaptively based on symbol + * frequencies. Initially, each code must be built assuming that all symbols + * have equal frequency. Following that, each code must be rebuilt whenever a + * certain number of symbols has been decoded with it. + * + * In general, multiple valid Huffman codes can be constructed from a set of + * symbol frequencies. Like other compression formats such as XPRESS, LZX, and + * DEFLATE, the LZMS format solves this ambiguity by requiring that all Huffman + * codes be constructed in canonical form. This form requires that same-length + * codewords be lexicographically ordered the same way as the corresponding + * symbols and that all shorter codewords lexicographically precede longer + * codewords. + * + * Codewords in all the LZMS Huffman codes are limited to 15 bits. If the + * canonical code for a given set of symbol frequencies has any codewords longer + * than 15 bits, then all frequencies must be divided by 2, rounding up, and the + * code construction must be attempted again. + * + * A LZMS-compressed block seemingly cannot have a compressed size greater than + * or equal to the uncompressed size. In such cases the block must be stored + * uncompressed. + * + * After all LZMS items have been decoded, the data must be postprocessed to + * translate absolute address encoded in x86 instructions into their original + * relative addresses. + * + * Details omitted above can be found in the code. Note that in the absence of + * an official specification there is no guarantee that this decompressor + * handles all possible cases. + */ + +#ifdef HAVE_CONFIG_H +# include "config.h" +#endif + +#include "wimlib.h" +#include "wimlib/compress_common.h" +#include "wimlib/decompressor_ops.h" +#include "wimlib/decompress_common.h" #include "wimlib/lzms.h" -#include "wimlib/error.h" -#include +#include "wimlib/util.h" + +#include + +#define LZMS_DECODE_TABLE_BITS 10 + +/* Structure used for range decoding, reading bits forwards. This is the first + * logical bitstream mentioned above. */ +struct lzms_range_decoder_raw { + /* The relevant part of the current range. Although the logical range + * for range decoding is a very large integer, only a small portion + * matters at any given time, and it can be normalized (shifted left) + * whenever it gets too small. */ + u32 range; + + /* The current position in the range encoded by the portion of the input + * read so far. */ + u32 code; -typedef HANDLE DECOMPRESSOR_HANDLE; -typedef PVOID PCOMPRESS_ALLOCATION_ROUTINES; -typedef DECOMPRESSOR_HANDLE *PDECOMPRESSOR_HANDLE; + /* Pointer to the next little-endian 16-bit integer in the compressed + * input data (reading forwards). */ + const le16 *in; -typedef enum { - COMPRESS_INFORMATION_CLASS_INVALID = 0, - COMPRESS_INFORMATION_CLASS_LEVEL, - COMPRESS_INFORMATION_CLASS_BLOCK_SIZE -} COMPRESS_INFORMATION_CLASS; + /* Number of 16-bit integers remaining in the compressed input data + * (reading forwards). */ + size_t num_le16_remaining; +}; -#define COMPRESS_ALGORITHM_LZMS 0x00000005 -#define COMPRESS_RAW 0x20000000 /* Not documented */ +/* Structure used for reading raw bits backwards. This is the second logical + * bitstream mentioned above. */ +struct lzms_input_bitstream { + /* Holding variable for bits that have been read from the compressed + * data. The bits are ordered from high-order to low-order. */ + /* XXX: Without special-case code to handle reading more than 17 bits + * at a time, this needs to be 64 bits rather than 32 bits. */ + u64 bitbuf; -static HMODULE hCabinetDll; -static pthread_mutex_t cabinetDllMutex = PTHREAD_MUTEX_INITIALIZER; + /* Number of bits in @bitbuf that are are used. */ + unsigned num_filled_bits; -static BOOL (WINAPI *CreateDecompressor) - (DWORD Algorithm, - PCOMPRESS_ALLOCATION_ROUTINES AllocationRoutines, - PDECOMPRESSOR_HANDLE DecompressorHandle); + /* Pointer to the one past the next little-endian 16-bit integer in the + * compressed input data (reading backwards). */ + const le16 *in; -static BOOL (WINAPI *CloseDecompressor) - (DECOMPRESSOR_HANDLE DecompressorHandle); + /* Number of 16-bit integers remaining in the compressed input data + * (reading backwards). */ + size_t num_le16_remaining; +}; -static BOOL (WINAPI *Decompress) - (DECOMPRESSOR_HANDLE DecompressorHandle, - PVOID CompressedData, - SIZE_T CompressedDataSize, - PVOID UncompressedBuffer, - SIZE_T UncompressedBufferSize, - PSIZE_T UncompressedDataSize); +/* Structure used for range decoding. This wraps around `struct + * lzms_range_decoder_raw' to use and maintain probability entries. */ +struct lzms_range_decoder { + /* Pointer to the raw range decoder, which has no persistent knowledge + * of probabilities. Multiple lzms_range_decoder's share the same + * lzms_range_decoder_raw. */ + struct lzms_range_decoder_raw *rd; -int -lzms_decompress(const void *cbuf, unsigned clen, void *ubuf, unsigned ulen, - unsigned window_size) + /* Bits recently decoded by this range decoder. This are used as in + * index into @prob_entries. */ + u32 state; + + /* Bitmask for @state to prevent its value from exceeding the number of + * probability entries. */ + u32 mask; + + /* Probability entries being used for this range decoder. */ + struct lzms_probability_entry prob_entries[LZMS_MAX_NUM_STATES]; +}; + +/* Structure used for Huffman decoding, optionally using the decoded symbols as + * slots into a base table to determine how many extra bits need to be read to + * reconstitute the full value. */ +struct lzms_huffman_decoder { + + /* Bitstream to read Huffman-encoded symbols and verbatim bits from. + * Multiple lzms_huffman_decoder's share the same lzms_input_bitstream. + */ + struct lzms_input_bitstream *is; + + /* Pointer to the slot base table to use. It is indexed by the decoded + * Huffman symbol that specifies the slot. The entry specifies the base + * value to use, and the position of its high bit is the number of + * additional bits that must be read to reconstitute the full value. + * + * This member need not be set if only raw Huffman symbols are being + * read using this decoder. */ + const u32 *slot_base_tab; + + const u8 *extra_bits_tab; + + /* Number of symbols that have been read using this code far. Reset to + * 0 whenever the code is rebuilt. */ + u32 num_syms_read; + + /* When @num_syms_read reaches this number, the Huffman code must be + * rebuilt. */ + u32 rebuild_freq; + + /* Number of symbols in the represented Huffman code. */ + unsigned num_syms; + + /* Running totals of symbol frequencies. These are diluted slightly + * whenever the code is rebuilt. */ + u32 sym_freqs[LZMS_MAX_NUM_SYMS]; + + /* The length, in bits, of each symbol in the Huffman code. */ + u8 lens[LZMS_MAX_NUM_SYMS]; + + /* The codeword of each symbol in the Huffman code. */ + u16 codewords[LZMS_MAX_NUM_SYMS]; + + /* A table for quickly decoding symbols encoded using the Huffman code. + */ + u16 decode_table[(1U << LZMS_DECODE_TABLE_BITS) + 2 * LZMS_MAX_NUM_SYMS] + _aligned_attribute(DECODE_TABLE_ALIGNMENT); +}; + +/* State of the LZMS decompressor. */ +struct lzms_decompressor { + + /* Pointer to the beginning of the uncompressed data buffer. */ + u8 *out_begin; + + /* Pointer to the next position in the uncompressed data buffer. */ + u8 *out_next; + + /* Pointer to one past the end of the uncompressed data buffer. */ + u8 *out_end; + + /* Range decoder, which reads bits from the beginning of the compressed + * block, going forwards. */ + struct lzms_range_decoder_raw rd; + + /* Input bitstream, which reads from the end of the compressed block, + * going backwards. */ + struct lzms_input_bitstream is; + + /* Range decoders. */ + struct lzms_range_decoder main_range_decoder; + struct lzms_range_decoder match_range_decoder; + struct lzms_range_decoder lz_match_range_decoder; + struct lzms_range_decoder lz_repeat_match_range_decoders[LZMS_NUM_RECENT_OFFSETS - 1]; + struct lzms_range_decoder delta_match_range_decoder; + struct lzms_range_decoder delta_repeat_match_range_decoders[LZMS_NUM_RECENT_OFFSETS - 1]; + + /* Huffman decoders. */ + struct lzms_huffman_decoder literal_decoder; + struct lzms_huffman_decoder lz_offset_decoder; + struct lzms_huffman_decoder length_decoder; + struct lzms_huffman_decoder delta_power_decoder; + struct lzms_huffman_decoder delta_offset_decoder; + + /* LRU (least-recently-used) queues for match information. */ + struct lzms_lru_queues lru; + + /* Used for postprocessing. */ + s32 last_target_usages[65536]; +}; + +/* Initialize the input bitstream @is to read forwards from the specified + * compressed data buffer @in that is @in_limit 16-bit integers long. */ +static void +lzms_input_bitstream_init(struct lzms_input_bitstream *is, + const le16 *in, size_t in_limit) { - int ret; - DECOMPRESSOR_HANDLE h; - - ERROR("clen=%u, ulen=%u, window_size=%u", clen, ulen, window_size); - - if (hCabinetDll == NULL) { - pthread_mutex_lock(&cabinetDllMutex); - - if (hCabinetDll == NULL) { - hCabinetDll = LoadLibrary(L"Cabinet.dll"); - - if (hCabinetDll == NULL) { - ERROR("Can't load Cabinet.dll"); - ret = -1; - goto unlock; - } - - CreateDecompressor = (void*)GetProcAddress(hCabinetDll, "CreateDecompressor"); - Decompress = (void*)GetProcAddress(hCabinetDll, "Decompress"); - CloseDecompressor = (void*)GetProcAddress(hCabinetDll, "CloseDecompressor"); - - if (CreateDecompressor == NULL || - Decompress == NULL || - CloseDecompressor == NULL) - { - ERROR("Can't find LZMS compression routines in Cabinet.dll"); - ret = -1; - goto unlock; - } + is->bitbuf = 0; + is->num_filled_bits = 0; + is->in = in + in_limit; + is->num_le16_remaining = in_limit; +} + +/* Ensures that @num_bits bits are buffered in the input bitstream. */ +static int +lzms_input_bitstream_ensure_bits(struct lzms_input_bitstream *is, + unsigned num_bits) +{ + while (is->num_filled_bits < num_bits) { + u64 next; + + LZMS_ASSERT(is->num_filled_bits + 16 <= sizeof(is->bitbuf) * 8); + + if (unlikely(is->num_le16_remaining == 0)) + return -1; + + next = le16_to_cpu(*--is->in); + is->num_le16_remaining--; + + is->bitbuf |= next << (sizeof(is->bitbuf) * 8 - is->num_filled_bits - 16); + is->num_filled_bits += 16; + } + return 0; + +} + +/* Returns the next @num_bits bits that are buffered in the input bitstream. */ +static u32 +lzms_input_bitstream_peek_bits(struct lzms_input_bitstream *is, + unsigned num_bits) +{ + LZMS_ASSERT(is->num_filled_bits >= num_bits); + return is->bitbuf >> (sizeof(is->bitbuf) * 8 - num_bits); +} + +/* Removes the next @num_bits bits that are buffered in the input bitstream. */ +static void +lzms_input_bitstream_remove_bits(struct lzms_input_bitstream *is, + unsigned num_bits) +{ + LZMS_ASSERT(is->num_filled_bits >= num_bits); + is->bitbuf <<= num_bits; + is->num_filled_bits -= num_bits; +} + +/* Removes and returns the next @num_bits bits that are buffered in the input + * bitstream. */ +static u32 +lzms_input_bitstream_pop_bits(struct lzms_input_bitstream *is, + unsigned num_bits) +{ + u32 bits = lzms_input_bitstream_peek_bits(is, num_bits); + lzms_input_bitstream_remove_bits(is, num_bits); + return bits; +} + +/* Reads the next @num_bits from the input bitstream. */ +static u32 +lzms_input_bitstream_read_bits(struct lzms_input_bitstream *is, + unsigned num_bits) +{ + if (unlikely(lzms_input_bitstream_ensure_bits(is, num_bits))) + return 0; + return lzms_input_bitstream_pop_bits(is, num_bits); +} + +/* Initialize the range decoder @rd to read forwards from the specified + * compressed data buffer @in that is @in_limit 16-bit integers long. */ +static void +lzms_range_decoder_raw_init(struct lzms_range_decoder_raw *rd, + const le16 *in, size_t in_limit) +{ + rd->range = 0xffffffff; + rd->code = ((u32)le16_to_cpu(in[0]) << 16) | + ((u32)le16_to_cpu(in[1]) << 0); + rd->in = in + 2; + rd->num_le16_remaining = in_limit - 2; +} + +/* Ensures the current range of the range decoder has at least 16 bits of + * precision. */ +static int +lzms_range_decoder_raw_normalize(struct lzms_range_decoder_raw *rd) +{ + if (rd->range <= 0xffff) { + rd->range <<= 16; + if (unlikely(rd->num_le16_remaining == 0)) + return -1; + rd->code = (rd->code << 16) | le16_to_cpu(*rd->in++); + rd->num_le16_remaining--; + } + return 0; +} + +/* Decode and return the next bit from the range decoder (raw version). + * + * @prob is the chance out of LZMS_PROBABILITY_MAX that the next bit is 0. + */ +static int +lzms_range_decoder_raw_decode_bit(struct lzms_range_decoder_raw *rd, u32 prob) +{ + u32 bound; + + /* Ensure the range has at least 16 bits of precision. */ + lzms_range_decoder_raw_normalize(rd); + + /* Based on the probability, calculate the bound between the 0-bit + * region and the 1-bit region of the range. */ + bound = (rd->range >> LZMS_PROBABILITY_BITS) * prob; + + if (rd->code < bound) { + /* Current code is in the 0-bit region of the range. */ + rd->range = bound; + return 0; + } else { + /* Current code is in the 1-bit region of the range. */ + rd->range -= bound; + rd->code -= bound; + return 1; + } +} + +/* Decode and return the next bit from the range decoder. This wraps around + * lzms_range_decoder_raw_decode_bit() to handle using and updating the + * appropriate probability table. */ +static int +lzms_range_decode_bit(struct lzms_range_decoder *dec) +{ + struct lzms_probability_entry *prob_entry; + u32 prob; + int bit; + + /* Load the probability entry corresponding to the current state. */ + prob_entry = &dec->prob_entries[dec->state]; + + /* Treat the number of zero bits in the most recently decoded + * LZMS_PROBABILITY_MAX bits with this probability entry as the chance, + * out of LZMS_PROBABILITY_MAX, that the next bit will be a 0. However, + * don't allow 0% or 100% probabilities. */ + prob = prob_entry->num_recent_zero_bits; + if (prob == LZMS_PROBABILITY_MAX) + prob = LZMS_PROBABILITY_MAX - 1; + else if (prob == 0) + prob = 1; + + /* Decode the next bit. */ + bit = lzms_range_decoder_raw_decode_bit(dec->rd, prob); + + /* Update the state based on the newly decoded bit. */ + dec->state = (((dec->state << 1) | bit) & dec->mask); + + /* Update the recent bits, including the cached count of 0's. */ + BUILD_BUG_ON(LZMS_PROBABILITY_MAX > sizeof(prob_entry->recent_bits) * 8); + if (bit == 0) { + if (prob_entry->recent_bits & (1ULL << (LZMS_PROBABILITY_MAX - 1))) { + /* Replacing 1 bit with 0 bit; increment the zero count. + */ + prob_entry->num_recent_zero_bits++; + } + } else { + if (!(prob_entry->recent_bits & (1ULL << (LZMS_PROBABILITY_MAX - 1)))) { + /* Replacing 0 bit with 1 bit; decrement the zero count. + */ + prob_entry->num_recent_zero_bits--; } - ret = 0; - unlock: - pthread_mutex_unlock(&cabinetDllMutex); - if (ret) - goto out; } + prob_entry->recent_bits = (prob_entry->recent_bits << 1) | bit; + + /* Return the decoded bit. */ + return bit; +} - if (!CreateDecompressor(COMPRESS_ALGORITHM_LZMS | COMPRESS_RAW, NULL, &h)) { - ERROR("Failed to create LZMS decompressor (err %d)!", GetLastError()); - ret = -1; - goto out; +/* Build the decoding table for a new adaptive Huffman code using the alphabet + * used in the specified Huffman decoder, with the symbol frequencies + * dec->sym_freqs. */ +static void +lzms_rebuild_adaptive_huffman_code(struct lzms_huffman_decoder *dec) +{ + + /* XXX: This implementation makes use of code already implemented for + * the XPRESS and LZX compression formats. However, since for the + * adaptive codes used in LZMS we don't actually need the explicit codes + * themselves, only the decode tables, it may be possible to optimize + * this by somehow directly building or updating the Huffman decode + * table. This may be a worthwhile optimization because the adaptive + * codes change many times throughout a decompression run. */ + LZMS_DEBUG("Rebuilding adaptive Huffman code (num_syms=%u)", + dec->num_syms); + make_canonical_huffman_code(dec->num_syms, LZMS_MAX_CODEWORD_LEN, + dec->sym_freqs, dec->lens, dec->codewords); +#if defined(ENABLE_LZMS_DEBUG) + int ret = +#endif + make_huffman_decode_table(dec->decode_table, dec->num_syms, + LZMS_DECODE_TABLE_BITS, dec->lens, + LZMS_MAX_CODEWORD_LEN); + LZMS_ASSERT(ret == 0); +} + +/* Decode and return the next Huffman-encoded symbol from the LZMS-compressed + * block using the specified Huffman decoder. */ +static u32 +lzms_huffman_decode_symbol(struct lzms_huffman_decoder *dec) +{ + const u8 *lens = dec->lens; + const u16 *decode_table = dec->decode_table; + struct lzms_input_bitstream *is = dec->is; + + /* The Huffman codes used in LZMS are adaptive and must be rebuilt + * whenever a certain number of symbols have been read. Each such + * rebuild uses the current symbol frequencies, but the format also + * requires that the symbol frequencies be halved after each code + * rebuild. This diminishes the effect of old symbols on the current + * Huffman codes, thereby causing the Huffman codes to be more locally + * adaptable. */ + if (dec->num_syms_read == dec->rebuild_freq) { + lzms_rebuild_adaptive_huffman_code(dec); + for (unsigned i = 0; i < dec->num_syms; i++) { + dec->sym_freqs[i] >>= 1; + dec->sym_freqs[i] += 1; + } + dec->num_syms_read = 0; } + /* In the following Huffman decoding implementation, the first + * LZMS_DECODE_TABLE_BITS of the input are used as an offset into a + * decode table. The entry will either provide the decoded symbol + * directly, or else a "real" Huffman binary tree will be searched to + * decode the symbol. */ + + lzms_input_bitstream_ensure_bits(is, LZMS_MAX_CODEWORD_LEN); - /* TODO: Some sort of chunk header? */ - unsigned offset; - if (clen <= window_size) { - offset = 0; + u16 key_bits = lzms_input_bitstream_peek_bits(is, LZMS_DECODE_TABLE_BITS); + u16 sym = decode_table[key_bits]; + + if (sym < dec->num_syms) { + /* Fast case: The decode table directly provided the symbol. */ + lzms_input_bitstream_remove_bits(is, lens[sym]); } else { - const unsigned *p = cbuf; - ERROR("%08x(%u) %08x %08x %08x %08x(%u)", - p[0], p[0], p[1], p[2], p[3], p[4], p[4]); - offset = 20; + /* Slow case: The symbol took too many bits to include directly + * in the decode table, so search for it in a binary tree at the + * end of the decode table. */ + lzms_input_bitstream_remove_bits(is, LZMS_DECODE_TABLE_BITS); + do { + key_bits = sym + lzms_input_bitstream_pop_bits(is, 1); + } while ((sym = decode_table[key_bits]) >= dec->num_syms); } - SIZE_T actual_ulen = -1; - if (!Decompress(h, (void*)cbuf + offset, clen - offset, ubuf, ulen, &actual_ulen)) { - ERROR("Failed to decompress LZMS-compressed data (err %d)!", GetLastError()); - ret = -1; - goto out_close_decompressor; + + /* Tally and return the decoded symbol. */ + ++dec->sym_freqs[sym]; + ++dec->num_syms_read; + return sym; +} + +/* Decode a number from the LZMS bitstream, encoded as a Huffman-encoded symbol + * specifying a "slot" (whose corresponding value is looked up in a static + * table) plus the number specified by a number of extra bits depending on the + * slot. */ +static u32 +lzms_decode_value(struct lzms_huffman_decoder *dec) +{ + unsigned slot; + unsigned num_extra_bits; + u32 extra_bits; + + LZMS_ASSERT(dec->slot_base_tab != NULL); + LZMS_ASSERT(dec->extra_bits_tab != NULL); + + /* Read the slot (position slot, length slot, etc.), which is encoded as + * a Huffman symbol. */ + slot = lzms_huffman_decode_symbol(dec); + + /* Get the number of extra bits needed to represent the range of values + * that share the slot. */ + num_extra_bits = dec->extra_bits_tab[slot]; + + /* Read the number of extra bits and add them to the slot base to form + * the final decoded value. */ + extra_bits = lzms_input_bitstream_read_bits(dec->is, num_extra_bits); + return dec->slot_base_tab[slot] + extra_bits; +} + +/* Copy a literal to the output buffer. */ +static int +lzms_copy_literal(struct lzms_decompressor *ctx, u8 literal) +{ + *ctx->out_next++ = literal; + return 0; +} + +/* Validate an LZ match and copy it to the output buffer. */ +static int +lzms_copy_lz_match(struct lzms_decompressor *ctx, u32 length, u32 offset) +{ + u8 *out_next; + u8 *matchptr; + + if (length > ctx->out_end - ctx->out_next) { + LZMS_DEBUG("Match overrun!"); + return -1; + } + if (offset > ctx->out_next - ctx->out_begin) { + LZMS_DEBUG("Match underrun!"); + return -1; + } + + out_next = ctx->out_next; + matchptr = out_next - offset; + while (length--) + *out_next++ = *matchptr++; + + ctx->out_next = out_next; + return 0; +} + +/* Validate a delta match and copy it to the output buffer. */ +static int +lzms_copy_delta_match(struct lzms_decompressor *ctx, u32 length, + u32 power, u32 raw_offset) +{ + u32 offset1 = 1U << power; + u32 offset2 = raw_offset << power; + u32 offset = offset1 + offset2; + u8 *out_next; + u8 *matchptr1; + u8 *matchptr2; + u8 *matchptr; + + if (length > ctx->out_end - ctx->out_next) { + LZMS_DEBUG("Match overrun!"); + return -1; + } + if (offset > ctx->out_next - ctx->out_begin) { + LZMS_DEBUG("Match underrun!"); + return -1; + } + + out_next = ctx->out_next; + matchptr1 = out_next - offset1; + matchptr2 = out_next - offset2; + matchptr = out_next - offset; + + while (length--) + *out_next++ = *matchptr1++ + *matchptr2++ - *matchptr++; + + ctx->out_next = out_next; + return 0; +} + +/* Decode a (length, offset) pair from the input. */ +static int +lzms_decode_lz_match(struct lzms_decompressor *ctx) +{ + int bit; + u32 length, offset; + + /* Decode the match offset. The next range-encoded bit indicates + * whether it's a repeat offset or an explicit offset. */ + + bit = lzms_range_decode_bit(&ctx->lz_match_range_decoder); + if (bit == 0) { + /* Explicit offset. */ + offset = lzms_decode_value(&ctx->lz_offset_decoder); + } else { + /* Repeat offset. */ + int i; + + for (i = 0; i < LZMS_NUM_RECENT_OFFSETS - 1; i++) + if (!lzms_range_decode_bit(&ctx->lz_repeat_match_range_decoders[i])) + break; + + offset = ctx->lru.lz.recent_offsets[i]; + + for (; i < LZMS_NUM_RECENT_OFFSETS; i++) + ctx->lru.lz.recent_offsets[i] = ctx->lru.lz.recent_offsets[i + 1]; + } + + /* Decode match length, which is always given explicitly (there is no + * LRU queue for repeat lengths). */ + length = lzms_decode_value(&ctx->length_decoder); + + ctx->lru.lz.upcoming_offset = offset; + + LZMS_DEBUG("Decoded %s LZ match: length=%u, offset=%u", + (bit ? "repeat" : "explicit"), length, offset); + + /* Validate the match and copy it to the output. */ + return lzms_copy_lz_match(ctx, length, offset); +} + +/* Decodes a "delta" match from the input. */ +static int +lzms_decode_delta_match(struct lzms_decompressor *ctx) +{ + int bit; + u32 length, power, raw_offset; + + /* Decode the match power and raw offset. The next range-encoded bit + * indicates whether these data are a repeat, or given explicitly. */ + + bit = lzms_range_decode_bit(&ctx->delta_match_range_decoder); + if (bit == 0) { + power = lzms_huffman_decode_symbol(&ctx->delta_power_decoder); + raw_offset = lzms_decode_value(&ctx->delta_offset_decoder); + } else { + int i; + + for (i = 0; i < LZMS_NUM_RECENT_OFFSETS - 1; i++) + if (!lzms_range_decode_bit(&ctx->delta_repeat_match_range_decoders[i])) + break; + + power = ctx->lru.delta.recent_powers[i]; + raw_offset = ctx->lru.delta.recent_offsets[i]; + + for (; i < LZMS_NUM_RECENT_OFFSETS; i++) { + ctx->lru.delta.recent_powers[i] = ctx->lru.delta.recent_powers[i + 1]; + ctx->lru.delta.recent_offsets[i] = ctx->lru.delta.recent_offsets[i + 1]; + } + } + + length = lzms_decode_value(&ctx->length_decoder); + + ctx->lru.delta.upcoming_power = power; + ctx->lru.delta.upcoming_offset = raw_offset; + + LZMS_DEBUG("Decoded %s delta match: length=%u, power=%u, raw_offset=%u", + (bit ? "repeat" : "explicit"), length, power, raw_offset); + + /* Validate the match and copy it to the output. */ + return lzms_copy_delta_match(ctx, length, power, raw_offset); +} + +/* Decode a LZ or delta match. */ +static int +lzms_decode_match(struct lzms_decompressor *ctx) +{ + if (!lzms_range_decode_bit(&ctx->match_range_decoder)) + return lzms_decode_lz_match(ctx); + else + return lzms_decode_delta_match(ctx); +} + +/* Decode a literal byte encoded using the literal Huffman code. */ +static int +lzms_decode_literal(struct lzms_decompressor *ctx) +{ + u8 literal = lzms_huffman_decode_symbol(&ctx->literal_decoder); + LZMS_DEBUG("Decoded literal: 0x%02x", literal); + return lzms_copy_literal(ctx, literal); +} + +/* Decode the next LZMS match or literal. */ +static int +lzms_decode_item(struct lzms_decompressor *ctx) +{ + int ret; + + ctx->lru.lz.upcoming_offset = 0; + ctx->lru.delta.upcoming_power = 0; + ctx->lru.delta.upcoming_offset = 0; + + if (lzms_range_decode_bit(&ctx->main_range_decoder)) + ret = lzms_decode_match(ctx); + else + ret = lzms_decode_literal(ctx); + + if (ret) + return ret; + + lzms_update_lru_queues(&ctx->lru); + return 0; +} + +static void +lzms_init_range_decoder(struct lzms_range_decoder *dec, + struct lzms_range_decoder_raw *rd, u32 num_states) +{ + dec->rd = rd; + dec->state = 0; + dec->mask = num_states - 1; + for (u32 i = 0; i < num_states; i++) { + dec->prob_entries[i].num_recent_zero_bits = LZMS_INITIAL_PROBABILITY; + dec->prob_entries[i].recent_bits = LZMS_INITIAL_RECENT_BITS; + } +} + +static void +lzms_init_huffman_decoder(struct lzms_huffman_decoder *dec, + struct lzms_input_bitstream *is, + const u32 *slot_base_tab, + const u8 *extra_bits_tab, + unsigned num_syms, + unsigned rebuild_freq) +{ + dec->is = is; + dec->slot_base_tab = slot_base_tab; + dec->extra_bits_tab = extra_bits_tab; + dec->num_syms = num_syms; + dec->num_syms_read = rebuild_freq; + dec->rebuild_freq = rebuild_freq; + for (unsigned i = 0; i < num_syms; i++) + dec->sym_freqs[i] = 1; +} + +/* Prepare to decode items from an LZMS-compressed block. */ +static void +lzms_init_decompressor(struct lzms_decompressor *ctx, + const void *cdata, unsigned clen, + void *ubuf, unsigned ulen) +{ + unsigned num_position_slots; + + LZMS_DEBUG("Initializing decompressor (clen=%u, ulen=%u)", clen, ulen); + + /* Initialize output pointers. */ + ctx->out_begin = ubuf; + ctx->out_next = ubuf; + ctx->out_end = (u8*)ubuf + ulen; + + /* Initialize the raw range decoder (reading forwards). */ + lzms_range_decoder_raw_init(&ctx->rd, cdata, clen / 2); + + /* Initialize the input bitstream for Huffman symbols (reading + * backwards) */ + lzms_input_bitstream_init(&ctx->is, cdata, clen / 2); + + /* Calculate the number of position slots needed for this compressed + * block. */ + num_position_slots = lzms_get_position_slot(ulen - 1) + 1; + + LZMS_DEBUG("Using %u position slots", num_position_slots); + + /* Initialize Huffman decoders for each alphabet used in the compressed + * representation. */ + lzms_init_huffman_decoder(&ctx->literal_decoder, &ctx->is, + NULL, NULL, LZMS_NUM_LITERAL_SYMS, + LZMS_LITERAL_CODE_REBUILD_FREQ); + + lzms_init_huffman_decoder(&ctx->lz_offset_decoder, &ctx->is, + lzms_position_slot_base, + lzms_extra_position_bits, + num_position_slots, + LZMS_LZ_OFFSET_CODE_REBUILD_FREQ); + + lzms_init_huffman_decoder(&ctx->length_decoder, &ctx->is, + lzms_length_slot_base, + lzms_extra_length_bits, + LZMS_NUM_LEN_SYMS, + LZMS_LENGTH_CODE_REBUILD_FREQ); + + lzms_init_huffman_decoder(&ctx->delta_offset_decoder, &ctx->is, + lzms_position_slot_base, + lzms_extra_position_bits, + num_position_slots, + LZMS_DELTA_OFFSET_CODE_REBUILD_FREQ); + + lzms_init_huffman_decoder(&ctx->delta_power_decoder, &ctx->is, + NULL, NULL, LZMS_NUM_DELTA_POWER_SYMS, + LZMS_DELTA_POWER_CODE_REBUILD_FREQ); + + + /* Initialize range decoders, all of which wrap around the same + * lzms_range_decoder_raw. */ + lzms_init_range_decoder(&ctx->main_range_decoder, + &ctx->rd, LZMS_NUM_MAIN_STATES); + + lzms_init_range_decoder(&ctx->match_range_decoder, + &ctx->rd, LZMS_NUM_MATCH_STATES); + + lzms_init_range_decoder(&ctx->lz_match_range_decoder, + &ctx->rd, LZMS_NUM_LZ_MATCH_STATES); + + for (size_t i = 0; i < ARRAY_LEN(ctx->lz_repeat_match_range_decoders); i++) + lzms_init_range_decoder(&ctx->lz_repeat_match_range_decoders[i], + &ctx->rd, LZMS_NUM_LZ_REPEAT_MATCH_STATES); + + lzms_init_range_decoder(&ctx->delta_match_range_decoder, + &ctx->rd, LZMS_NUM_DELTA_MATCH_STATES); + + for (size_t i = 0; i < ARRAY_LEN(ctx->delta_repeat_match_range_decoders); i++) + lzms_init_range_decoder(&ctx->delta_repeat_match_range_decoders[i], + &ctx->rd, LZMS_NUM_DELTA_REPEAT_MATCH_STATES); + + /* Initialize LRU match information. */ + lzms_init_lru_queues(&ctx->lru); + + LZMS_DEBUG("Decompressor successfully initialized"); +} + +/* Decode the series of literals and matches from the LZMS-compressed data. + * Returns 0 on success; nonzero if the compressed data is invalid. */ +static int +lzms_decode_items(const u8 *cdata, size_t clen, u8 *ubuf, size_t ulen, + struct lzms_decompressor *ctx) +{ + /* Initialize the LZMS decompressor. */ + lzms_init_decompressor(ctx, cdata, clen, ubuf, ulen); + + /* Decode the sequence of items. */ + while (ctx->out_next != ctx->out_end) { + LZMS_DEBUG("Position %u", ctx->out_next - ctx->out_begin); + if (lzms_decode_item(ctx)) + return -1; + } + return 0; +} + +static int +lzms_decompress(const void *compressed_data, size_t compressed_size, + void *uncompressed_data, size_t uncompressed_size, void *_ctx) +{ + struct lzms_decompressor *ctx = _ctx; + + /* The range decoder requires that a minimum of 4 bytes of compressed + * data be initially available. */ + if (compressed_size < 4) { + LZMS_DEBUG("Compressed size too small (got %zu, expected >= 4)", + compressed_size); + return -1; + } + + /* A LZMS-compressed data block should be evenly divisible into 16-bit + * integers. */ + if (compressed_size % 2 != 0) { + LZMS_DEBUG("Compressed size not divisible by 2 (got %zu)", + compressed_size); + return -1; } - if (actual_ulen != ulen) { - ERROR("Unexpected actual uncompressed length (got %u, expected %u)", - actual_ulen, ulen); - ret = -1; - goto out_close_decompressor; + /* Handle the trivial case where nothing needs to be decompressed. + * (Necessary because a window of size 0 does not have a valid position + * slot.) */ + if (uncompressed_size == 0) + return 0; + + /* The x86 post-processor requires that the uncompressed length fit into + * a signed 32-bit integer. Also, the position slot table cannot be + * searched for a position of INT32_MAX or greater. */ + if (uncompressed_size >= INT32_MAX) { + LZMS_DEBUG("Uncompressed length too large " + "(got %zu, expected < INT32_MAX)", + uncompressed_size); + return -1; } - ERROR("Successfully decompressed data."); - ret = 0; -out_close_decompressor: - CloseDecompressor(h); -out: - return ret; + /* Decode the literals and matches. */ + if (lzms_decode_items(compressed_data, compressed_size, + uncompressed_data, uncompressed_size, ctx)) + return -1; + + /* Postprocess the data. */ + lzms_x86_filter(uncompressed_data, uncompressed_size, + ctx->last_target_usages, true); + + LZMS_DEBUG("Decompression successful."); + return 0; +} + +static void +lzms_free_decompressor(void *_ctx) +{ + struct lzms_decompressor *ctx = _ctx; + + FREE(ctx); } + +static int +lzms_create_decompressor(size_t max_block_size, + const struct wimlib_decompressor_params_header *params, + void **ctx_ret) +{ + struct lzms_decompressor *ctx; + + ctx = MALLOC(sizeof(struct lzms_decompressor)); + if (ctx == NULL) + return WIMLIB_ERR_NOMEM; + + /* Initialize position and length slot data if not done already. */ + lzms_init_slots(); + + *ctx_ret = ctx; + return 0; +} + +const struct decompressor_ops lzms_decompressor_ops = { + .create_decompressor = lzms_create_decompressor, + .decompress = lzms_decompress, + .free_decompressor = lzms_free_decompressor, +};