Style
[wimlib] / src / compress.c
1 /*
2  * compress.c
3  *
4  * Functions used for compression.
5  */
6
7 /*
8  * Copyright (C) 2012, 2013 Eric Biggers
9  *
10  * This file is part of wimlib, a library for working with WIM files.
11  *
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)
15  * any later version.
16  *
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
20  * details.
21  *
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/.
24  */
25
26 #include "compress.h"
27 #include <stdlib.h>
28 #include <string.h>
29
30 static inline void
31 flush_bits(struct output_bitstream *ostream)
32 {
33         *(u16*)ostream->bit_output = cpu_to_le16(ostream->bitbuf);
34         ostream->bit_output = ostream->next_bit_output;
35         ostream->next_bit_output = ostream->output;
36         ostream->output += 2;
37         ostream->num_bytes_remaining -= 2;
38 }
39
40 /* Writes @num_bits bits, given by the @num_bits least significant bits of
41  * @bits, to the output @ostream. */
42 int
43 bitstream_put_bits(struct output_bitstream *ostream, output_bitbuf_t bits,
44                    unsigned num_bits)
45 {
46         unsigned rem_bits;
47
48         wimlib_assert(num_bits <= 16);
49         if (num_bits <= ostream->free_bits) {
50                 ostream->bitbuf = (ostream->bitbuf << num_bits) | bits;
51                 ostream->free_bits -= num_bits;
52         } else {
53
54                 if (ostream->num_bytes_remaining + (ostream->output -
55                                                 ostream->bit_output) < 2)
56                         return 1;
57
58                 /* It is tricky to output the bits correctly.  The correct way
59                  * is to output little-endian 2-byte words, such that the bits
60                  * in the SECOND byte logically precede those in the FIRST byte.
61                  * While the byte order is little-endian, the bit order is
62                  * big-endian; the first bit in a byte is the high-order one.
63                  * Any multi-bit numbers are in bit-big-endian form, so the
64                  * low-order bit of a multi-bit number is the LAST bit to be
65                  * output. */
66                 rem_bits = num_bits - ostream->free_bits;
67                 ostream->bitbuf <<= ostream->free_bits;
68                 ostream->bitbuf |= bits >> rem_bits;
69                 flush_bits(ostream);
70                 ostream->free_bits = 16 - rem_bits;
71                 ostream->bitbuf = bits;
72
73         }
74         return 0;
75 }
76
77 /* Flushes any remaining bits in the output buffer to the output byte stream. */
78 int
79 flush_output_bitstream(struct output_bitstream *ostream)
80 {
81         if (ostream->num_bytes_remaining + (ostream->output -
82                                         ostream->bit_output) < 2)
83                 return 1;
84         if (ostream->free_bits != 16) {
85                 ostream->bitbuf <<= ostream->free_bits;
86                 flush_bits(ostream);
87         }
88         return 0;
89 }
90
91 /* Initializes an output bit buffer to write its output to the memory location
92  * pointer to by @data. */
93 void
94 init_output_bitstream(struct output_bitstream *ostream, void *data,
95                       unsigned num_bytes)
96 {
97         wimlib_assert(num_bytes >= 4);
98
99         ostream->bitbuf              = 0;
100         ostream->free_bits           = 16;
101         ostream->bit_output          = (u8*)data;
102         ostream->next_bit_output     = (u8*)data + 2;
103         ostream->output              = (u8*)data + 4;
104         ostream->num_bytes_remaining = num_bytes - 4;
105 }
106
107 /* Intermediate (non-leaf) node in a Huffman tree. */
108 typedef struct HuffmanNode {
109         u32 freq;
110         u16 sym;
111         union {
112                 u16 path_len;
113                 u16 height;
114         };
115         struct HuffmanNode *left_child;
116         struct HuffmanNode *right_child;
117 } HuffmanNode;
118
119 /* Leaf node in a Huffman tree.  The fields are in the same order as the
120  * HuffmanNode, so it can be cast to a HuffmanNode.  There are no pointers to
121  * the children in the leaf node. */
122 typedef struct {
123         u32 freq;
124         u16 sym;
125         union {
126                 u16 path_len;
127                 u16 height;
128         };
129 } HuffmanLeafNode;
130
131 /* Comparator function for HuffmanLeafNodes.  Sorts primarily by symbol
132  * frequency and secondarily by symbol value. */
133 static int
134 cmp_leaves_by_freq(const void *__leaf1, const void *__leaf2)
135 {
136         const HuffmanLeafNode *leaf1 = __leaf1;
137         const HuffmanLeafNode *leaf2 = __leaf2;
138
139         int freq_diff = (int)leaf1->freq - (int)leaf2->freq;
140
141         if (freq_diff == 0)
142                 return (int)leaf1->sym - (int)leaf2->sym;
143         else
144                 return freq_diff;
145 }
146
147 /* Comparator function for HuffmanLeafNodes.  Sorts primarily by code length and
148  * secondarily by symbol value. */
149 static int
150 cmp_leaves_by_code_len(const void *__leaf1, const void *__leaf2)
151 {
152         const HuffmanLeafNode *leaf1 = __leaf1;
153         const HuffmanLeafNode *leaf2 = __leaf2;
154
155         int code_len_diff = (int)leaf1->path_len - (int)leaf2->path_len;
156
157         if (code_len_diff == 0)
158                 return (int)leaf1->sym - (int)leaf2->sym;
159         else
160                 return code_len_diff;
161 }
162
163 /* Recursive function to calculate the depth of the leaves in a Huffman tree.
164  * */
165 static void
166 huffman_tree_compute_path_lengths(HuffmanNode *node, u16 cur_len)
167 {
168         if (node->sym == (u16)(-1)) {
169                 /* Intermediate node. */
170                 huffman_tree_compute_path_lengths(node->left_child, cur_len + 1);
171                 huffman_tree_compute_path_lengths(node->right_child, cur_len + 1);
172         } else {
173                 /* Leaf node. */
174                 node->path_len = cur_len;
175         }
176 }
177
178 /* make_canonical_huffman_code: - Creates a canonical Huffman code from an array
179  *                                of symbol frequencies.
180  *
181  * The algorithm used is similar to the well-known algorithm that builds a
182  * Huffman tree using a minheap.  In that algorithm, the leaf nodes are
183  * initialized and inserted into the minheap with the frequency as the key.
184  * Repeatedly, the top two nodes (nodes with the lowest frequency) are taken out
185  * of the heap and made the children of a new node that has a frequency equal to
186  * the sum of the two frequencies of its children.  This new node is inserted
187  * into the heap.  When all the nodes have been removed from the heap, what
188  * remains is the Huffman tree. The Huffman code for a symbol is given by the
189  * path to it in the tree, where each left pointer is mapped to a 0 bit and each
190  * right pointer is mapped to a 1 bit.
191  *
192  * The algorithm used here uses an optimization that removes the need to
193  * actually use a heap.  The leaf nodes are first sorted by frequency, as
194  * opposed to being made into a heap.  Note that this sorting step takes O(n log
195  * n) time vs.  O(n) time for heapifying the array, where n is the number of
196  * symbols.  However, the heapless method is probably faster overall, due to the
197  * time saved later.  In the heapless method, whenever an intermediate node is
198  * created, it is not inserted into the sorted array.  Instead, the intermediate
199  * nodes are kept in a separate array, which is easily kept sorted because every
200  * time an intermediate node is initialized, it will have a frequency at least
201  * as high as that of the previous intermediate node that was initialized.  So
202  * whenever we want the 2 nodes, leaf or intermediate, that have the lowest
203  * frequency, we check the low-frequency ends of both arrays, which is an O(1)
204  * operation.
205  *
206  * The function builds a canonical Huffman code, not just any Huffman code.  A
207  * Huffman code is canonical if the codeword for each symbol numerically
208  * precedes the codeword for all other symbols of the same length that are
209  * numbered higher than the symbol, and additionally, all shorter codewords,
210  * 0-extended, numerically precede longer codewords.  A canonical Huffman code
211  * is useful because it can be reconstructed by only knowing the path lengths in
212  * the tree.  See the make_huffman_decode_table() function to see how to
213  * reconstruct a canonical Huffman code from only the lengths of the codes.
214  *
215  * @num_syms:  The number of symbols in the alphabet.
216  *
217  * @max_codeword_len:  The maximum allowed length of a codeword in the code.
218  *                      Note that if the code being created runs up against
219  *                      this restriction, the code ultimately created will be
220  *                      suboptimal, although there are some advantages for
221  *                      limiting the length of the codewords.
222  *
223  * @freq_tab:  An array of length @num_syms that contains the frequencies
224  *                      of each symbol in the uncompressed data.
225  *
226  * @lens:          An array of length @num_syms into which the lengths of the
227  *                      codewords for each symbol will be written.
228  *
229  * @codewords:     An array of @num_syms short integers into which the
230  *                      codewords for each symbol will be written.  The first
231  *                      lens[i] bits of codewords[i] will contain the codeword
232  *                      for symbol i.
233  */
234 void
235 make_canonical_huffman_code(unsigned num_syms, unsigned max_codeword_len,
236                             const freq_t freq_tab[], u8 lens[],
237                             u16 codewords[])
238 {
239         /* We require at least 2 possible symbols in the alphabet to produce a
240          * valid Huffman decoding table. It is allowed that fewer than 2 symbols
241          * are actually used, though. */
242         wimlib_assert(num_syms >= 2);
243
244         /* Initialize the lengths and codewords to 0 */
245         memset(lens, 0, num_syms * sizeof(lens[0]));
246         memset(codewords, 0, num_syms * sizeof(codewords[0]));
247
248         /* Calculate how many symbols have non-zero frequency.  These are the
249          * symbols that actually appeared in the input. */
250         unsigned num_used_symbols = 0;
251         for (unsigned i = 0; i < num_syms; i++)
252                 if (freq_tab[i] != 0)
253                         num_used_symbols++;
254
255
256         /* It is impossible to make a code for num_used_symbols symbols if there
257          * aren't enough code bits to uniquely represent all of them. */
258         wimlib_assert((1 << max_codeword_len) > num_used_symbols);
259
260         /* Initialize the array of leaf nodes with the symbols and their
261          * frequencies. */
262         HuffmanLeafNode leaves[num_used_symbols];
263         unsigned leaf_idx = 0;
264         for (unsigned i = 0; i < num_syms; i++) {
265                 if (freq_tab[i] != 0) {
266                         leaves[leaf_idx].freq = freq_tab[i];
267                         leaves[leaf_idx].sym  = i;
268                         leaves[leaf_idx].height = 0;
269                         leaf_idx++;
270                 }
271         }
272
273         /* Deal with the special cases where num_used_symbols < 2. */
274         if (num_used_symbols < 2) {
275                 if (num_used_symbols == 0) {
276                         /* If num_used_symbols is 0, there are no symbols in the
277                          * input, so it must be empty.  This should be an error,
278                          * but the LZX format expects this case to succeed.  All
279                          * the codeword lengths are simply marked as 0 (which
280                          * was already done.) */
281                 } else {
282                         /* If only one symbol is present, the LZX format
283                          * requires that the Huffman code include two codewords.
284                          * One is not used.  Note that this doesn't make the
285                          * encoded data take up more room anyway, since binary
286                          * data itself has 2 symbols. */
287
288                         unsigned sym = leaves[0].sym;
289
290                         codewords[0] = 0;
291                         lens[0]      = 1;
292                         if (sym == 0) {
293                                 /* dummy symbol is 1, real symbol is 0 */
294                                 codewords[1] = 1;
295                                 lens[1]      = 1;
296                         } else {
297                                 /* dummy symbol is 0, real symbol is sym */
298                                 codewords[sym] = 1;
299                                 lens[sym]      = 1;
300                         }
301                 }
302                 return;
303         }
304
305         /* Otherwise, there are at least 2 symbols in the input, so we need to
306          * find a real Huffman code. */
307
308
309         /* Declare the array of intermediate nodes.  An intermediate node is not
310          * associated with a symbol. Instead, it represents some binary code
311          * prefix that is shared between at least 2 codewords.  There can be at
312          * most num_used_symbols - 1 intermediate nodes when creating a Huffman
313          * code.  This is because if there were at least num_used_symbols nodes,
314          * the code would be suboptimal because there would be at least one
315          * unnecessary intermediate node.
316          *
317          * The worst case (greatest number of intermediate nodes) would be if
318          * all the intermediate nodes were chained together.  This results in
319          * num_used_symbols - 1 intermediate nodes.  If num_used_symbols is at
320          * least 17, this configuration would not be allowed because the LZX
321          * format constrains codes to 16 bits or less each.  However, it is
322          * still possible for there to be more than 16 intermediate nodes, as
323          * long as no leaf has a depth of more than 16.  */
324         HuffmanNode inodes[num_used_symbols - 1];
325
326
327         /* Pointer to the leaf node of lowest frequency that hasn't already been
328          * added as the child of some intermediate note. */
329         HuffmanLeafNode *cur_leaf;
330
331         /* Pointer past the end of the array of leaves. */
332         HuffmanLeafNode *end_leaf = &leaves[num_used_symbols];
333
334         /* Pointer to the intermediate node of lowest frequency. */
335         HuffmanNode     *cur_inode;
336
337         /* Pointer to the next unallocated intermediate node. */
338         HuffmanNode     *next_inode;
339
340         /* Only jump back to here if the maximum length of the codewords allowed
341          * by the LZX format (16 bits) is exceeded. */
342 try_building_tree_again:
343
344         /* Sort the leaves from those that correspond to the least frequent
345          * symbol, to those that correspond to the most frequent symbol.  If two
346          * leaves have the same frequency, they are sorted by symbol. */
347         qsort(leaves, num_used_symbols, sizeof(leaves[0]), cmp_leaves_by_freq);
348
349         cur_leaf   = &leaves[0];
350         cur_inode  = &inodes[0];
351         next_inode = &inodes[0];
352
353         /* The following loop takes the two lowest frequency nodes of those
354          * remaining and makes them the children of the next available
355          * intermediate node.  It continues until all the leaf nodes and
356          * intermediate nodes have been used up, or the maximum allowed length
357          * for the codewords is exceeded.  For the latter case, we must adjust
358          * the frequencies to be more equal and then execute this loop again. */
359         while (1) {
360
361                 /* Lowest frequency node. */
362                 HuffmanNode *f1;
363
364                 /* Second lowest frequency node. */
365                 HuffmanNode *f2;
366
367                 /* Get the lowest and second lowest frequency nodes from the
368                  * remaining leaves or from the intermediate nodes. */
369
370                 if (cur_leaf != end_leaf && (cur_inode == next_inode ||
371                                         cur_leaf->freq <= cur_inode->freq)) {
372                         f1 = (HuffmanNode*)cur_leaf++;
373                 } else if (cur_inode != next_inode) {
374                         f1 = cur_inode++;
375                 }
376
377                 if (cur_leaf != end_leaf && (cur_inode == next_inode ||
378                                         cur_leaf->freq <= cur_inode->freq)) {
379                         f2 = (HuffmanNode*)cur_leaf++;
380                 } else if (cur_inode != next_inode) {
381                         f2 = cur_inode++;
382                 } else {
383                         /* All nodes used up! */
384                         break;
385                 }
386
387                 /* next_inode becomes the parent of f1 and f2. */
388
389                 next_inode->freq   = f1->freq + f2->freq;
390                 next_inode->sym    = (u16)(-1); /* Invalid symbol. */
391                 next_inode->left_child   = f1;
392                 next_inode->right_child  = f2;
393
394                 /* We need to keep track of the height so that we can detect if
395                  * the length of a codeword has execeed max_codeword_len.   The
396                  * parent node has a height one higher than the maximum height
397                  * of its children. */
398                 next_inode->height = max(f1->height, f2->height) + 1;
399
400                 /* Check to see if the code length of the leaf farthest away
401                  * from next_inode has exceeded the maximum code length. */
402                 if (next_inode->height > max_codeword_len) {
403                         /* The code lengths can be made more uniform by making
404                          * the frequencies more uniform.  Divide all the
405                          * frequencies by 2, leaving 1 as the minimum frequency.
406                          * If this keeps happening, the symbol frequencies will
407                          * approach equality, which makes their Huffman
408                          * codewords approach the length
409                          * log_2(num_used_symbols).
410                          * */
411                         for (unsigned i = 0; i < num_used_symbols; i++)
412                                 if (leaves[i].freq > 1)
413                                         leaves[i].freq >>= 1;
414                         goto try_building_tree_again;
415                 }
416                 next_inode++;
417         }
418
419         /* The Huffman tree is now complete, and its height is no more than
420          * max_codeword_len.  */
421
422         HuffmanNode *root = next_inode - 1;
423         wimlib_assert(root->height <= max_codeword_len);
424
425         /* Compute the path lengths for the leaf nodes. */
426         huffman_tree_compute_path_lengths(root, 0);
427
428         /* Sort the leaf nodes primarily by code length and secondarily by
429          * symbol.  */
430         qsort(leaves, num_used_symbols, sizeof(leaves[0]), cmp_leaves_by_code_len);
431
432         u16 cur_codeword = 0;
433         unsigned cur_codeword_len = 0;
434         for (unsigned i = 0; i < num_used_symbols; i++) {
435
436                 /* Each time a codeword becomes one longer, the current codeword
437                  * is left shifted by one place.  This is part of the procedure
438                  * for enumerating the canonical Huffman code.  Additionally,
439                  * whenever a codeword is used, 1 is added to the current
440                  * codeword.  */
441
442                 unsigned len_diff = leaves[i].path_len - cur_codeword_len;
443                 cur_codeword <<= len_diff;
444                 cur_codeword_len += len_diff;
445
446                 u16 sym = leaves[i].sym;
447                 codewords[sym] = cur_codeword;
448                 lens[sym] = cur_codeword_len;
449
450                 cur_codeword++;
451         }
452 }