4 * Functions used for compression.
8 * Copyright (C) 2012, 2013 Eric Biggers
10 * This file is part of wimlib, a library for working with WIM files.
12 * wimlib is free software; you can redistribute it and/or modify it under the
13 * terms of the GNU General Public License as published by the Free
14 * Software Foundation; either version 3 of the License, or (at your option)
17 * wimlib is distributed in the hope that it will be useful, but WITHOUT ANY
18 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
19 * A PARTICULAR PURPOSE. See the GNU General Public License for more
22 * You should have received a copy of the GNU General Public License
23 * along with wimlib; if not, see http://www.gnu.org/licenses/.
30 #include "wimlib/assert.h"
31 #include "wimlib/endianness.h"
32 #include "wimlib/compiler.h"
33 #include "wimlib/compress.h"
34 #include "wimlib/util.h"
39 /* Writes @num_bits bits, given by the @num_bits least significant bits of
40 * @bits, to the output @ostream. */
42 bitstream_put_bits(struct output_bitstream *ostream, u32 bits,
45 bits &= (1U << num_bits) - 1;
46 while (num_bits > ostream->free_bits) {
47 /* Buffer variable does not have space for the new bits. It
48 * needs to be flushed as a 16-bit integer. Bits in the second
49 * byte logically precede those in the first byte
50 * (little-endian), but within each byte the bits are ordered
51 * from high to low. This is true for both XPRESS and LZX
54 /* There must be at least 2 bytes of space remaining. */
55 if (unlikely(ostream->bytes_remaining < 2)) {
56 ostream->overrun = true;
60 /* Fill the buffer with as many bits that fit. */
61 unsigned fill_bits = ostream->free_bits;
63 ostream->bitbuf <<= fill_bits;
64 ostream->bitbuf |= bits >> (num_bits - fill_bits);
66 *(le16*)ostream->bit_output = cpu_to_le16(ostream->bitbuf);
67 ostream->bit_output = ostream->next_bit_output;
68 ostream->next_bit_output = ostream->output;
70 ostream->bytes_remaining -= 2;
72 ostream->free_bits = 16;
73 num_bits -= fill_bits;
74 bits &= (1U << num_bits) - 1;
77 /* Buffer variable has space for the new bits. */
78 ostream->bitbuf = (ostream->bitbuf << num_bits) | bits;
79 ostream->free_bits -= num_bits;
83 bitstream_put_byte(struct output_bitstream *ostream, u8 n)
85 if (unlikely(ostream->bytes_remaining < 1)) {
86 ostream->overrun = true;
89 *ostream->output++ = n;
90 ostream->bytes_remaining--;
93 /* Flushes any remaining bits to the output bitstream.
95 * Returns -1 if the stream has overrun; otherwise returns the total number of
96 * bytes in the output. */
98 flush_output_bitstream(struct output_bitstream *ostream)
100 if (unlikely(ostream->overrun))
101 return ~(input_idx_t)0;
103 *(le16*)ostream->bit_output =
104 cpu_to_le16((u16)((u32)ostream->bitbuf << ostream->free_bits));
105 *(le16*)ostream->next_bit_output =
108 return ostream->output - ostream->output_start;
111 /* Initializes an output bit buffer to write its output to the memory location
112 * pointer to by @data. */
114 init_output_bitstream(struct output_bitstream *ostream,
115 void *data, unsigned num_bytes)
117 wimlib_assert(num_bytes >= 4);
120 ostream->free_bits = 16;
121 ostream->output_start = data;
122 ostream->bit_output = data;
123 ostream->next_bit_output = data + 2;
124 ostream->output = data + 4;
125 ostream->bytes_remaining = num_bytes;
126 ostream->overrun = false;
138 typedef struct HuffmanIntermediateNode {
139 HuffmanNode node_base;
140 HuffmanNode *left_child;
141 HuffmanNode *right_child;
142 } HuffmanIntermediateNode;
145 /* Comparator function for HuffmanNodes. Sorts primarily by symbol
146 * frequency and secondarily by symbol value. */
148 cmp_nodes_by_freq(const void *_leaf1, const void *_leaf2)
150 const HuffmanNode *leaf1 = _leaf1;
151 const HuffmanNode *leaf2 = _leaf2;
153 if (leaf1->freq > leaf2->freq)
155 else if (leaf1->freq < leaf2->freq)
158 return (int)leaf1->sym - (int)leaf2->sym;
161 /* Comparator function for HuffmanNodes. Sorts primarily by code length and
162 * secondarily by symbol value. */
164 cmp_nodes_by_code_len(const void *_leaf1, const void *_leaf2)
166 const HuffmanNode *leaf1 = _leaf1;
167 const HuffmanNode *leaf2 = _leaf2;
169 int code_len_diff = (int)leaf1->path_len - (int)leaf2->path_len;
171 if (code_len_diff == 0)
172 return (int)leaf1->sym - (int)leaf2->sym;
174 return code_len_diff;
177 #define INVALID_SYMBOL 0xffff
179 /* Recursive function to calculate the depth of the leaves in a Huffman tree.
182 huffman_tree_compute_path_lengths(HuffmanNode *base_node, u16 cur_len)
184 if (base_node->sym == INVALID_SYMBOL) {
185 /* Intermediate node. */
186 HuffmanIntermediateNode *node = (HuffmanIntermediateNode*)base_node;
187 huffman_tree_compute_path_lengths(node->left_child, cur_len + 1);
188 huffman_tree_compute_path_lengths(node->right_child, cur_len + 1);
191 base_node->path_len = cur_len;
195 /* make_canonical_huffman_code: - Creates a canonical Huffman code from an array
196 * of symbol frequencies.
198 * The algorithm used is similar to the well-known algorithm that builds a
199 * Huffman tree using a minheap. In that algorithm, the leaf nodes are
200 * initialized and inserted into the minheap with the frequency as the key.
201 * Repeatedly, the top two nodes (nodes with the lowest frequency) are taken out
202 * of the heap and made the children of a new node that has a frequency equal to
203 * the sum of the two frequencies of its children. This new node is inserted
204 * into the heap. When all the nodes have been removed from the heap, what
205 * remains is the Huffman tree. The Huffman code for a symbol is given by the
206 * path to it in the tree, where each left pointer is mapped to a 0 bit and each
207 * right pointer is mapped to a 1 bit.
209 * The algorithm used here uses an optimization that removes the need to
210 * actually use a heap. The leaf nodes are first sorted by frequency, as
211 * opposed to being made into a heap. Note that this sorting step takes O(n log
212 * n) time vs. O(n) time for heapifying the array, where n is the number of
213 * symbols. However, the heapless method is probably faster overall, due to the
214 * time saved later. In the heapless method, whenever an intermediate node is
215 * created, it is not inserted into the sorted array. Instead, the intermediate
216 * nodes are kept in a separate array, which is easily kept sorted because every
217 * time an intermediate node is initialized, it will have a frequency at least
218 * as high as that of the previous intermediate node that was initialized. So
219 * whenever we want the 2 nodes, leaf or intermediate, that have the lowest
220 * frequency, we check the low-frequency ends of both arrays, which is an O(1)
223 * The function builds a canonical Huffman code, not just any Huffman code. A
224 * Huffman code is canonical if the codeword for each symbol numerically
225 * precedes the codeword for all other symbols of the same length that are
226 * numbered higher than the symbol, and additionally, all shorter codewords,
227 * 0-extended, numerically precede longer codewords. A canonical Huffman code
228 * is useful because it can be reconstructed by only knowing the path lengths in
229 * the tree. See the make_huffman_decode_table() function to see how to
230 * reconstruct a canonical Huffman code from only the lengths of the codes.
232 * @num_syms: The number of symbols in the alphabet.
234 * @max_codeword_len: The maximum allowed length of a codeword in the code.
235 * Note that if the code being created runs up against
236 * this restriction, the code ultimately created will be
237 * suboptimal, although there are some advantages for
238 * limiting the length of the codewords.
240 * @freq_tab: An array of length @num_syms that contains the frequencies
241 * of each symbol in the uncompressed data.
243 * @lens: An array of length @num_syms into which the lengths of the
244 * codewords for each symbol will be written.
246 * @codewords: An array of @num_syms short integers into which the
247 * codewords for each symbol will be written. The first
248 * lens[i] bits of codewords[i] will contain the codeword
252 make_canonical_huffman_code(unsigned num_syms,
253 unsigned max_codeword_len,
254 const input_idx_t freq_tab[restrict],
256 u16 codewords[restrict])
258 /* We require at least 2 possible symbols in the alphabet to produce a
259 * valid Huffman decoding table. It is allowed that fewer than 2 symbols
260 * are actually used, though. */
261 wimlib_assert(num_syms >= 2 && num_syms < INVALID_SYMBOL);
263 /* Initialize the lengths and codewords to 0 */
264 memset(lens, 0, num_syms * sizeof(lens[0]));
265 memset(codewords, 0, num_syms * sizeof(codewords[0]));
267 /* Calculate how many symbols have non-zero frequency. These are the
268 * symbols that actually appeared in the input. */
269 unsigned num_used_symbols = 0;
270 for (unsigned i = 0; i < num_syms; i++)
271 if (freq_tab[i] != 0)
275 /* It is impossible to make a code for num_used_symbols symbols if there
276 * aren't enough code bits to uniquely represent all of them. */
277 wimlib_assert((1 << max_codeword_len) > num_used_symbols);
279 /* Initialize the array of leaf nodes with the symbols and their
281 HuffmanNode leaves[num_used_symbols];
282 unsigned leaf_idx = 0;
283 for (unsigned i = 0; i < num_syms; i++) {
284 if (freq_tab[i] != 0) {
285 leaves[leaf_idx].freq = freq_tab[i];
286 leaves[leaf_idx].sym = i;
287 leaves[leaf_idx].height = 0;
292 /* Deal with the special cases where num_used_symbols < 2. */
293 if (num_used_symbols < 2) {
294 if (num_used_symbols == 0) {
295 /* If num_used_symbols is 0, there are no symbols in the
296 * input, so it must be empty. This should be an error,
297 * but the LZX format expects this case to succeed. All
298 * the codeword lengths are simply marked as 0 (which
299 * was already done.) */
301 /* If only one symbol is present, the LZX format
302 * requires that the Huffman code include two codewords.
303 * One is not used. Note that this doesn't make the
304 * encoded data take up more room anyway, since binary
305 * data itself has 2 symbols. */
307 unsigned sym = leaves[0].sym;
312 /* dummy symbol is 1, real symbol is 0 */
316 /* dummy symbol is 0, real symbol is sym */
324 /* Otherwise, there are at least 2 symbols in the input, so we need to
325 * find a real Huffman code. */
328 /* Declare the array of intermediate nodes. An intermediate node is not
329 * associated with a symbol. Instead, it represents some binary code
330 * prefix that is shared between at least 2 codewords. There can be at
331 * most num_used_symbols - 1 intermediate nodes when creating a Huffman
332 * code. This is because if there were at least num_used_symbols nodes,
333 * the code would be suboptimal because there would be at least one
334 * unnecessary intermediate node.
336 * The worst case (greatest number of intermediate nodes) would be if
337 * all the intermediate nodes were chained together. This results in
338 * num_used_symbols - 1 intermediate nodes. If num_used_symbols is at
339 * least 17, this configuration would not be allowed because the LZX
340 * format constrains codes to 16 bits or less each. However, it is
341 * still possible for there to be more than 16 intermediate nodes, as
342 * long as no leaf has a depth of more than 16. */
343 HuffmanIntermediateNode inodes[num_used_symbols - 1];
346 /* Pointer to the leaf node of lowest frequency that hasn't already been
347 * added as the child of some intermediate note. */
348 HuffmanNode *cur_leaf;
350 /* Pointer past the end of the array of leaves. */
351 HuffmanNode *end_leaf = &leaves[num_used_symbols];
353 /* Pointer to the intermediate node of lowest frequency. */
354 HuffmanIntermediateNode *cur_inode;
356 /* Pointer to the next unallocated intermediate node. */
357 HuffmanIntermediateNode *next_inode;
359 /* Only jump back to here if the maximum length of the codewords allowed
360 * by the LZX format (16 bits) is exceeded. */
361 try_building_tree_again:
363 /* Sort the leaves from those that correspond to the least frequent
364 * symbol, to those that correspond to the most frequent symbol. If two
365 * leaves have the same frequency, they are sorted by symbol. */
366 qsort(leaves, num_used_symbols, sizeof(leaves[0]), cmp_nodes_by_freq);
368 cur_leaf = &leaves[0];
369 cur_inode = &inodes[0];
370 next_inode = &inodes[0];
372 /* The following loop takes the two lowest frequency nodes of those
373 * remaining and makes them the children of the next available
374 * intermediate node. It continues until all the leaf nodes and
375 * intermediate nodes have been used up, or the maximum allowed length
376 * for the codewords is exceeded. For the latter case, we must adjust
377 * the frequencies to be more equal and then execute this loop again. */
380 /* Lowest frequency node. */
383 /* Second lowest frequency node. */
386 /* Get the lowest and second lowest frequency nodes from the
387 * remaining leaves or from the intermediate nodes. */
389 if (cur_leaf != end_leaf && (cur_inode == next_inode ||
390 cur_leaf->freq <= cur_inode->node_base.freq)) {
392 } else if (cur_inode != next_inode) {
393 f1 = (HuffmanNode*)cur_inode++;
396 if (cur_leaf != end_leaf && (cur_inode == next_inode ||
397 cur_leaf->freq <= cur_inode->node_base.freq)) {
399 } else if (cur_inode != next_inode) {
400 f2 = (HuffmanNode*)cur_inode++;
402 /* All nodes used up! */
406 /* next_inode becomes the parent of f1 and f2. */
408 next_inode->node_base.freq = f1->freq + f2->freq;
409 next_inode->node_base.sym = INVALID_SYMBOL;
410 next_inode->left_child = f1;
411 next_inode->right_child = f2;
413 /* We need to keep track of the height so that we can detect if
414 * the length of a codeword has execeed max_codeword_len. The
415 * parent node has a height one higher than the maximum height
416 * of its children. */
417 next_inode->node_base.height = max(f1->height, f2->height) + 1;
419 /* Check to see if the code length of the leaf farthest away
420 * from next_inode has exceeded the maximum code length. */
421 if (next_inode->node_base.height > max_codeword_len) {
422 /* The code lengths can be made more uniform by making
423 * the frequencies more uniform. Divide all the
424 * frequencies by 2, leaving 1 as the minimum frequency.
425 * If this keeps happening, the symbol frequencies will
426 * approach equality, which makes their Huffman
427 * codewords approach the length
428 * log_2(num_used_symbols).
430 for (unsigned i = 0; i < num_used_symbols; i++)
431 leaves[i].freq = (leaves[i].freq + 1) >> 1;
433 goto try_building_tree_again;
438 /* The Huffman tree is now complete, and its height is no more than
439 * max_codeword_len. */
441 HuffmanIntermediateNode *root = next_inode - 1;
442 wimlib_assert(root->node_base.height <= max_codeword_len);
444 /* Compute the path lengths for the leaf nodes. */
445 huffman_tree_compute_path_lengths(&root->node_base, 0);
447 /* Sort the leaf nodes primarily by code length and secondarily by
449 qsort(leaves, num_used_symbols, sizeof(leaves[0]), cmp_nodes_by_code_len);
451 u16 cur_codeword = 0;
452 unsigned cur_codeword_len = 0;
453 for (unsigned i = 0; i < num_used_symbols; i++) {
455 /* Each time a codeword becomes one longer, the current codeword
456 * is left shifted by one place. This is part of the procedure
457 * for enumerating the canonical Huffman code. Additionally,
458 * whenever a codeword is used, 1 is added to the current
461 unsigned len_diff = leaves[i].path_len - cur_codeword_len;
462 cur_codeword <<= len_diff;
463 cur_codeword_len += len_diff;
465 u16 sym = leaves[i].sym;
466 codewords[sym] = cur_codeword;
467 lens[sym] = cur_codeword_len;