New compression/decompression API
[wimlib] / src / compress_common.c
1 /*
2  * compress_common.c
3  *
4  * Code for compression shared among multiple compression formats.
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 #include "wimlib/assert.h"
31 #include "wimlib/endianness.h"
32 #include "wimlib/compiler.h"
33 #include "wimlib/compress_common.h"
34 #include "wimlib/util.h"
35
36 #include <stdlib.h>
37 #include <string.h>
38
39 /* Writes @num_bits bits, given by the @num_bits least significant bits of
40  * @bits, to the output @ostream. */
41 void
42 bitstream_put_bits(struct output_bitstream *ostream, u32 bits,
43                    unsigned num_bits)
44 {
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
52                  * compression.  */
53
54                 /* There must be at least 2 bytes of space remaining.  */
55                 if (unlikely(ostream->bytes_remaining < 2)) {
56                         ostream->overrun = true;
57                         return;
58                 }
59
60                 /* Fill the buffer with as many bits that fit.  */
61                 unsigned fill_bits = ostream->free_bits;
62
63                 ostream->bitbuf <<= fill_bits;
64                 ostream->bitbuf |= bits >> (num_bits - fill_bits);
65
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;
69                 ostream->output += 2;
70                 ostream->bytes_remaining -= 2;
71
72                 ostream->free_bits = 16;
73                 num_bits -= fill_bits;
74                 bits &= (1U << num_bits) - 1;
75         }
76
77         /* Buffer variable has space for the new bits.  */
78         ostream->bitbuf = (ostream->bitbuf << num_bits) | bits;
79         ostream->free_bits -= num_bits;
80 }
81
82 void
83 bitstream_put_byte(struct output_bitstream *ostream, u8 n)
84 {
85         if (unlikely(ostream->bytes_remaining < 1)) {
86                 ostream->overrun = true;
87                 return;
88         }
89         *ostream->output++ = n;
90         ostream->bytes_remaining--;
91 }
92
93 /* Flushes any remaining bits to the output bitstream.
94  *
95  * Returns -1 if the stream has overrun; otherwise returns the total number of
96  * bytes in the output.  */
97 input_idx_t
98 flush_output_bitstream(struct output_bitstream *ostream)
99 {
100         if (unlikely(ostream->overrun))
101                 return ~(input_idx_t)0;
102
103         *(le16*)ostream->bit_output =
104                 cpu_to_le16((u16)((u32)ostream->bitbuf << ostream->free_bits));
105         *(le16*)ostream->next_bit_output =
106                 cpu_to_le16(0);
107
108         return ostream->output - ostream->output_start;
109 }
110
111 /* Initializes an output bit buffer to write its output to the memory location
112  * pointer to by @data. */
113 void
114 init_output_bitstream(struct output_bitstream *ostream,
115                       void *data, unsigned num_bytes)
116 {
117         wimlib_assert(num_bytes >= 4);
118
119         ostream->bitbuf              = 0;
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 - 4;
126         ostream->overrun             = false;
127 }
128
129 typedef struct {
130         input_idx_t freq;
131         u16 sym;
132         union {
133                 u16 path_len;
134                 u16 height;
135         };
136 } HuffmanNode;
137
138 typedef struct HuffmanIntermediateNode {
139         HuffmanNode node_base;
140         HuffmanNode *left_child;
141         HuffmanNode *right_child;
142 } HuffmanIntermediateNode;
143
144
145 /* Comparator function for HuffmanNodes.  Sorts primarily by symbol
146  * frequency and secondarily by symbol value. */
147 static int
148 cmp_nodes_by_freq(const void *_leaf1, const void *_leaf2)
149 {
150         const HuffmanNode *leaf1 = _leaf1;
151         const HuffmanNode *leaf2 = _leaf2;
152
153         if (leaf1->freq > leaf2->freq)
154                 return 1;
155         else if (leaf1->freq < leaf2->freq)
156                 return -1;
157         else
158                 return (int)leaf1->sym - (int)leaf2->sym;
159 }
160
161 /* Comparator function for HuffmanNodes.  Sorts primarily by code length and
162  * secondarily by symbol value. */
163 static int
164 cmp_nodes_by_code_len(const void *_leaf1, const void *_leaf2)
165 {
166         const HuffmanNode *leaf1 = _leaf1;
167         const HuffmanNode *leaf2 = _leaf2;
168
169         int code_len_diff = (int)leaf1->path_len - (int)leaf2->path_len;
170
171         if (code_len_diff == 0)
172                 return (int)leaf1->sym - (int)leaf2->sym;
173         else
174                 return code_len_diff;
175 }
176
177 #define INVALID_SYMBOL 0xffff
178
179 /* Recursive function to calculate the depth of the leaves in a Huffman tree.
180  * */
181 static void
182 huffman_tree_compute_path_lengths(HuffmanNode *base_node, u16 cur_len)
183 {
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);
189         } else {
190                 /* Leaf node. */
191                 base_node->path_len = cur_len;
192         }
193 }
194
195 /* make_canonical_huffman_code: - Creates a canonical Huffman code from an array
196  *                                of symbol frequencies.
197  *
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.
208  *
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)
221  * operation.
222  *
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.
231  *
232  * @num_syms:  The number of symbols in the alphabet.
233  *
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.
239  *
240  * @freq_tab:  An array of length @num_syms that contains the frequencies
241  *                      of each symbol in the uncompressed data.
242  *
243  * @lens:          An array of length @num_syms into which the lengths of the
244  *                      codewords for each symbol will be written.
245  *
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
249  *                      for symbol i.
250  */
251 void
252 make_canonical_huffman_code(unsigned num_syms,
253                             unsigned max_codeword_len,
254                             const input_idx_t freq_tab[restrict],
255                             u8 lens[restrict],
256                             u16 codewords[restrict])
257 {
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);
262
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]));
266
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)
272                         num_used_symbols++;
273
274
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);
278
279         /* Initialize the array of leaf nodes with the symbols and their
280          * frequencies. */
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;
288                         leaf_idx++;
289                 }
290         }
291
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.) */
300                 } else {
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. */
306
307                         unsigned sym = leaves[0].sym;
308
309                         codewords[0] = 0;
310                         lens[0]      = 1;
311                         if (sym == 0) {
312                                 /* dummy symbol is 1, real symbol is 0 */
313                                 codewords[1] = 1;
314                                 lens[1]      = 1;
315                         } else {
316                                 /* dummy symbol is 0, real symbol is sym */
317                                 codewords[sym] = 1;
318                                 lens[sym]      = 1;
319                         }
320                 }
321                 return;
322         }
323
324         /* Otherwise, there are at least 2 symbols in the input, so we need to
325          * find a real Huffman code. */
326
327
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.
335          *
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];
344
345
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;
349
350         /* Pointer past the end of the array of leaves. */
351         HuffmanNode *end_leaf = &leaves[num_used_symbols];
352
353         /* Pointer to the intermediate node of lowest frequency. */
354         HuffmanIntermediateNode     *cur_inode;
355
356         /* Pointer to the next unallocated intermediate node. */
357         HuffmanIntermediateNode     *next_inode;
358
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:
362
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);
367
368         cur_leaf   = &leaves[0];
369         cur_inode  = &inodes[0];
370         next_inode = &inodes[0];
371
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. */
378         while (1) {
379
380                 /* Lowest frequency node. */
381                 HuffmanNode *f1;
382
383                 /* Second lowest frequency node. */
384                 HuffmanNode *f2;
385
386                 /* Get the lowest and second lowest frequency nodes from the
387                  * remaining leaves or from the intermediate nodes. */
388
389                 if (cur_leaf != end_leaf && (cur_inode == next_inode ||
390                                         cur_leaf->freq <= cur_inode->node_base.freq)) {
391                         f1 = cur_leaf++;
392                 } else if (cur_inode != next_inode) {
393                         f1 = (HuffmanNode*)cur_inode++;
394                 }
395
396                 if (cur_leaf != end_leaf && (cur_inode == next_inode ||
397                                         cur_leaf->freq <= cur_inode->node_base.freq)) {
398                         f2 = cur_leaf++;
399                 } else if (cur_inode != next_inode) {
400                         f2 = (HuffmanNode*)cur_inode++;
401                 } else {
402                         /* All nodes used up! */
403                         break;
404                 }
405
406                 /* next_inode becomes the parent of f1 and f2. */
407
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;
412
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;
418
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).
429                          * */
430                         for (unsigned i = 0; i < num_used_symbols; i++)
431                                 leaves[i].freq = (leaves[i].freq + 1) >> 1;
432
433                         goto try_building_tree_again;
434                 }
435                 next_inode++;
436         }
437
438         /* The Huffman tree is now complete, and its height is no more than
439          * max_codeword_len.  */
440
441         HuffmanIntermediateNode *root = next_inode - 1;
442         wimlib_assert(root->node_base.height <= max_codeword_len);
443
444         /* Compute the path lengths for the leaf nodes. */
445         huffman_tree_compute_path_lengths(&root->node_base, 0);
446
447         /* Sort the leaf nodes primarily by code length and secondarily by
448          * symbol.  */
449         qsort(leaves, num_used_symbols, sizeof(leaves[0]), cmp_nodes_by_code_len);
450
451         u16 cur_codeword = 0;
452         unsigned cur_codeword_len = 0;
453         for (unsigned i = 0; i < num_used_symbols; i++) {
454
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
459                  * codeword.  */
460
461                 unsigned len_diff = leaves[i].path_len - cur_codeword_len;
462                 cur_codeword <<= len_diff;
463                 cur_codeword_len += len_diff;
464
465                 u16 sym = leaves[i].sym;
466                 codewords[sym] = cur_codeword;
467                 lens[sym] = cur_codeword_len;
468
469                 cur_codeword++;
470         }
471 }