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