959e1deb2d704b359acba3cdb4a8d5b1b04d3522
[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, 2014 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 <string.h>
37
38 /* Writes @num_bits bits, given by the @num_bits least significant bits of
39  * @bits, to the output @ostream. */
40 void
41 bitstream_put_bits(struct output_bitstream *ostream, u32 bits,
42                    unsigned num_bits)
43 {
44         bits &= (1U << num_bits) - 1;
45         while (num_bits > ostream->free_bits) {
46                 /* Buffer variable does not have space for the new bits.  It
47                  * needs to be flushed as a 16-bit integer.  Bits in the second
48                  * byte logically precede those in the first byte
49                  * (little-endian), but within each byte the bits are ordered
50                  * from high to low.  This is true for both XPRESS and LZX
51                  * compression.  */
52
53                 /* There must be at least 2 bytes of space remaining.  */
54                 if (unlikely(ostream->bytes_remaining < 2)) {
55                         ostream->overrun = true;
56                         return;
57                 }
58
59                 /* Fill the buffer with as many bits that fit.  */
60                 unsigned fill_bits = ostream->free_bits;
61
62                 ostream->bitbuf <<= fill_bits;
63                 ostream->bitbuf |= bits >> (num_bits - fill_bits);
64
65                 *(le16*)ostream->bit_output = cpu_to_le16(ostream->bitbuf);
66                 ostream->bit_output = ostream->next_bit_output;
67                 ostream->next_bit_output = ostream->output;
68                 ostream->output += 2;
69                 ostream->bytes_remaining -= 2;
70
71                 ostream->free_bits = 16;
72                 num_bits -= fill_bits;
73                 bits &= (1U << num_bits) - 1;
74         }
75
76         /* Buffer variable has space for the new bits.  */
77         ostream->bitbuf = (ostream->bitbuf << num_bits) | bits;
78         ostream->free_bits -= num_bits;
79 }
80
81 void
82 bitstream_put_byte(struct output_bitstream *ostream, u8 n)
83 {
84         if (unlikely(ostream->bytes_remaining < 1)) {
85                 ostream->overrun = true;
86                 return;
87         }
88         *ostream->output++ = n;
89         ostream->bytes_remaining--;
90 }
91
92 /* Flushes any remaining bits to the output bitstream.
93  *
94  * Returns -1 if the stream has overrun; otherwise returns the total number of
95  * bytes in the output.  */
96 input_idx_t
97 flush_output_bitstream(struct output_bitstream *ostream)
98 {
99         if (unlikely(ostream->overrun))
100                 return ~(input_idx_t)0;
101
102         *(le16*)ostream->bit_output =
103                 cpu_to_le16((u16)((u32)ostream->bitbuf << ostream->free_bits));
104         *(le16*)ostream->next_bit_output =
105                 cpu_to_le16(0);
106
107         return ostream->output - ostream->output_start;
108 }
109
110 /* Initializes an output bit buffer to write its output to the memory location
111  * pointer to by @data. */
112 void
113 init_output_bitstream(struct output_bitstream *ostream,
114                       void *data, unsigned num_bytes)
115 {
116         wimlib_assert(num_bytes >= 4);
117
118         ostream->bitbuf              = 0;
119         ostream->free_bits           = 16;
120         ostream->output_start        = data;
121         ostream->bit_output          = data;
122         ostream->next_bit_output     = data + 2;
123         ostream->output              = data + 4;
124         ostream->bytes_remaining     = num_bytes - 4;
125         ostream->overrun             = false;
126 }
127
128 /* Given the binary tree node A[subtree_idx] whose children already
129  * satisfy the maxheap property, swap the node with its greater child
130  * until it is greater than both its children, so that the maxheap
131  * property is satisfied in the subtree rooted at A[subtree_idx].  */
132 static void
133 heapify_subtree(u32 A[], unsigned length, unsigned subtree_idx)
134 {
135         unsigned parent_idx;
136         unsigned child_idx;
137         u32 v;
138
139         v = A[subtree_idx];
140         parent_idx = subtree_idx;
141         while ((child_idx = parent_idx * 2) <= length) {
142                 if (child_idx < length && A[child_idx + 1] > A[child_idx])
143                         child_idx++;
144                 if (v >= A[child_idx])
145                         break;
146                 A[parent_idx] = A[child_idx];
147                 parent_idx = child_idx;
148         }
149         A[parent_idx] = v;
150 }
151
152 /* Rearrange the array 'A' so that it satisfies the maxheap property.
153  * 'A' uses 1-based indices, so the children of A[i] are A[i*2] and A[i*2 + 1].
154  */
155 static void
156 heapify_array(u32 A[], unsigned length)
157 {
158         for (unsigned subtree_idx = length / 2; subtree_idx >= 1; subtree_idx--)
159                 heapify_subtree(A, length, subtree_idx);
160 }
161
162 /* Sort the array 'A', which contains 'length' unsigned 32-bit integers.  */
163 static void
164 heapsort(u32 A[], unsigned length)
165 {
166         A--; /* Use 1-based indices  */
167
168         heapify_array(A, length);
169
170         while (length >= 2) {
171                 swap(A[1], A[length]);
172                 length--;
173                 heapify_subtree(A, length, 1);
174         }
175 }
176
177 #define NUM_SYMBOL_BITS 10
178 #define SYMBOL_MASK ((1 << NUM_SYMBOL_BITS) - 1)
179
180 /*
181  * Sort the symbols primarily by frequency and secondarily by symbol
182  * value.  Discard symbols with zero frequency and fill in an array with
183  * the remaining symbols, along with their frequencies.  The low
184  * NUM_SYMBOL_BITS bits of each array entry will contain the symbol
185  * value, and the remaining bits will contain the frequency.
186  *
187  * @num_syms
188  *      Number of symbols in the alphabet.
189  *      Can't be greater than (1 << NUM_SYMBOL_BITS).
190  *
191  * @freqs[num_syms]
192  *      The frequency of each symbol.
193  *
194  * @lens[num_syms]
195  *      An array that eventually will hold the length of each codeword.
196  *      This function only fills in the codeword lengths for symbols that
197  *      have zero frequency, which are not well defined per se but will
198  *      be set to 0.
199  *
200  * @symout[num_syms]
201  *      The output array, described above.
202  *
203  * Returns the number of entries in 'symout' that were filled.  This is
204  * the number of symbols that have nonzero frequency.
205  */
206 static unsigned
207 sort_symbols(unsigned num_syms, const u32 freqs[restrict],
208              u8 lens[restrict], u32 symout[restrict])
209 {
210         unsigned num_used_syms;
211         unsigned num_counters;
212
213         /* We rely on heapsort, but with an added optimization.  Since
214          * it's common for most symbol frequencies to be low, we first do
215          * a count sort using a limited number of counters.  High
216          * frequencies will be counted in the last counter, and only they
217          * will be sorted with heapsort.
218          *
219          * Note: with more symbols, it is generally beneficial to have more
220          * counters.  About 1 counter per 4 symbols seems fast.
221          *
222          * Note: I also tested radix sort, but even for large symbol
223          * counts (> 255) and frequencies bounded at 16 bits (enabling
224          * radix sort by just two base-256 digits), it didn't seem any
225          * faster than the method implemented here.
226          *
227          * Note: I tested the optimized quicksort implementation from
228          * glibc (with indirection overhead removed), but it was only
229          * marginally faster than the simple heapsort implemented here.
230          *
231          * Tests were done with building the codes for LZX.  Results may
232          * vary for different compression algorithms...!  */
233
234         num_counters = (DIV_ROUND_UP(num_syms, 4) + 3) & ~3;
235
236         unsigned counters[num_counters];
237
238         memset(counters, 0, sizeof(counters));
239
240         /* Count the frequencies.  */
241         for (unsigned sym = 0; sym < num_syms; sym++)
242                 counters[min(freqs[sym], num_counters - 1)]++;
243
244         /* Make the counters cumulative, ignoring the zero-th, which
245          * counted symbols with zero frequency.  As a side effect, this
246          * calculates the number of symbols with nonzero frequency.  */
247         num_used_syms = 0;
248         for (unsigned i = 1; i < num_counters; i++) {
249                 unsigned count = counters[i];
250                 counters[i] = num_used_syms;
251                 num_used_syms += count;
252         }
253
254         /* Sort nonzero-frequency symbols using the counters.  At the
255          * same time, set the codeword lengths of zero-frequency symbols
256          * to 0.  */
257         for (unsigned sym = 0; sym < num_syms; sym++) {
258                 u32 freq = freqs[sym];
259                 if (freq != 0) {
260                         symout[counters[min(freq, num_counters - 1)]++] =
261                                 sym | (freq << NUM_SYMBOL_BITS);
262                 } else {
263                         lens[sym] = 0;
264                 }
265         }
266
267         /* Sort the symbols counted in the last counter.  */
268         heapsort(symout + counters[num_counters - 2],
269                  counters[num_counters - 1] - counters[num_counters - 2]);
270
271         return num_used_syms;
272 }
273
274 /*
275  * Build the Huffman tree.
276  *
277  * This is an optimized implementation that
278  *      (a) takes advantage of the frequencies being already sorted;
279  *      (b) only generates non-leaf nodes, since the non-leaf nodes of a
280  *          Huffman tree are sufficient to generate a canonical code;
281  *      (c) Only stores parent pointers, not child pointers;
282  *      (d) Produces the nodes in the same memory used for input
283  *          frequency information.
284  *
285  * Array 'A', which contains 'sym_count' entries, is used for both input
286  * and output.  For this function, 'sym_count' must be at least 2.
287  *
288  * For input, the array must contain the frequencies of the symbols,
289  * sorted in increasing order.  Specifically, each entry must contain a
290  * frequency left shifted by NUM_SYMBOL_BITS bits.  Any data in the low
291  * NUM_SYMBOL_BITS bits of the entries will be ignored by this function.
292  * Although these bits will, in fact, contain the symbols that correspond
293  * to the frequencies, this function is concerned with frequencies only
294  * and keeps the symbols as-is.
295  *
296  * For output, this function will produce the non-leaf nodes of the
297  * Huffman tree.  These nodes will be stored in the first (sym_count - 1)
298  * entries of the array.  Entry A[sym_count - 2] will represent the root
299  * node.  Each other node will contain the zero-based index of its parent
300  * node in 'A', left shifted by NUM_SYMBOL_BITS bits.  The low
301  * NUM_SYMBOL_BITS bits of each entry in A will be kept as-is.  Again,
302  * note that although these low bits will, in fact, contain a symbol
303  * value, this symbol will have *no relationship* with the Huffman tree
304  * node that happens to occupy the same slot.  This is because this
305  * implementation only generates the non-leaf nodes of the tree.
306  */
307 static void
308 build_tree(u32 A[], unsigned sym_count)
309 {
310         /* Index, in 'A', of next lowest frequency symbol that has not
311          * yet been processed.  */
312         unsigned i = 0;
313
314         /* Index, in 'A', of next lowest frequency parentless non-leaf
315          * node; or, if equal to 'e', then no such node exists yet.  */
316         unsigned b = 0;
317
318         /* Index, in 'A', of next node to allocate as a non-leaf.  */
319         unsigned e = 0;
320
321         do {
322                 unsigned m, n;
323                 u32 freq_shifted;
324
325                 /* Choose the two next lowest frequency entries.  */
326
327                 if (i != sym_count &&
328                     (b == e || (A[i] >> NUM_SYMBOL_BITS) <= (A[b] >> NUM_SYMBOL_BITS)))
329                         m = i++;
330                 else
331                         m = b++;
332
333                 if (i != sym_count &&
334                     (b == e || (A[i] >> NUM_SYMBOL_BITS) <= (A[b] >> NUM_SYMBOL_BITS)))
335                         n = i++;
336                 else
337                         n = b++;
338
339                 /* Allocate a non-leaf node and link the entries to it.
340                  *
341                  * If we link an entry that we're visiting for the first
342                  * time (via index 'i'), then we're actually linking a
343                  * leaf node and it will have no effect, since the leaf
344                  * will be overwritten with a non-leaf when index 'e'
345                  * catches up to it.  But it's not any slower to
346                  * unconditionally set the parent index.
347                  *
348                  * We also compute the frequency of the non-leaf node as
349                  * the sum of its two children's frequencies.  */
350
351                 freq_shifted = (A[m] & ~SYMBOL_MASK) + (A[n] & ~SYMBOL_MASK);
352
353                 A[m] = (A[m] & SYMBOL_MASK) | (e << NUM_SYMBOL_BITS);
354                 A[n] = (A[n] & SYMBOL_MASK) | (e << NUM_SYMBOL_BITS);
355                 A[e] = (A[e] & SYMBOL_MASK) | freq_shifted;
356                 e++;
357         } while (sym_count - e > 1);
358                 /* When just one entry remains, it is a "leaf" that was
359                  * linked to some other node.  We ignore it, since the
360                  * rest of the array contains the non-leaves which we
361                  * need.  (Note that we're assuming the cases with 0 or 1
362                  * symbols were handled separately.) */
363 }
364
365 /*
366  * Given the stripped-down Huffman tree constructed by build_tree(),
367  * determine the number of codewords that should be assigned each
368  * possible length, taking into account the length-limited constraint.
369  *
370  * @A
371  *      The array produced by build_tree(), containing parent index
372  *      information for the non-leaf nodes of the Huffman tree.  Each
373  *      entry in this array is a node; a node's parent always has a
374  *      greater index than that node itself.  This function will
375  *      overwrite the parent index information in this array, so
376  *      essentially it will destroy the tree.  However, the data in the
377  *      low NUM_SYMBOL_BITS of each entry will be preserved.
378  *
379  * @root_idx
380  *      The 0-based index of the root node in 'A', and consequently one
381  *      less than the number of tree node entries in 'A'.  (Or, really 2
382  *      less than the actual length of 'A'.)
383  *
384  * @len_counts
385  *      An array of length ('max_codeword_len' + 1) in which the number of
386  *      codewords having each length <= max_codeword_len will be
387  *      returned.
388  *
389  * @max_codeword_len
390  *      The maximum permissible codeword length.
391  */
392 static void
393 compute_length_counts(u32 A[restrict], unsigned root_idx,
394                       unsigned len_counts[restrict], unsigned max_codeword_len)
395 {
396         /* The key observations are:
397          *
398          * (1) We can traverse the non-leaf nodes of the tree, always
399          * visiting a parent before its children, by simply iterating
400          * through the array in reverse order.  Consequently, we can
401          * compute the depth of each node in one pass, overwriting the
402          * parent indices with depths.
403          *
404          * (2) We can initially assume that in the real Huffman tree,
405          * both children of the root are leaves.  This corresponds to two
406          * codewords of length 1.  Then, whenever we visit a (non-leaf)
407          * node during the traversal, we modify this assumption to
408          * account for the current node *not* being a leaf, but rather
409          * its two children being leaves.  This causes the loss of one
410          * codeword for the current depth and the addition of two
411          * codewords for the current depth plus one.
412          *
413          * (3) We can handle the length-limited constraint fairly easily
414          * by simply using the largest length available when a depth
415          * exceeds max_codeword_len.
416          */
417
418         for (unsigned len = 0; len <= max_codeword_len; len++)
419                 len_counts[len] = 0;
420         len_counts[1] = 2;
421
422         /* Set the root node's depth to 0.  */
423         A[root_idx] &= SYMBOL_MASK;
424
425         for (int node = root_idx - 1; node >= 0; node--) {
426
427                 /* Calculate the depth of this node.  */
428
429                 unsigned parent = A[node] >> NUM_SYMBOL_BITS;
430                 unsigned parent_depth = A[parent] >> NUM_SYMBOL_BITS;
431                 unsigned depth = parent_depth + 1;
432                 unsigned len = depth;
433
434                 /* Set the depth of this node so that it is available
435                  * when its children (if any) are processed.  */
436
437                 A[node] = (A[node] & SYMBOL_MASK) | (depth << NUM_SYMBOL_BITS);
438
439                 /* If needed, decrease the length to meet the
440                  * length-limited constraint.  This is not the optimal
441                  * method for generating length-limited Huffman codes!
442                  * But it should be good enough.  */
443                 if (len >= max_codeword_len) {
444                         len = max_codeword_len;
445                         do {
446                                 len--;
447                         } while (len_counts[len] == 0);
448                 }
449
450                 /* Account for the fact that we have a non-leaf node at
451                  * the current depth.  */
452                 len_counts[len]--;
453                 len_counts[len + 1] += 2;
454         }
455 }
456
457 /*
458  * Generate the codewords for a canonical Huffman code.
459  *
460  * @A
461  *      The output array for codewords.  In addition, initially this
462  *      array must contain the symbols, sorted primarily by frequency and
463  *      secondarily by symbol value, in the low NUM_SYMBOL_BITS bits of
464  *      each entry.
465  *
466  * @len
467  *      Output array for codeword lengths.
468  *
469  * @len_counts
470  *      An array that provides the number of codewords that will have
471  *      each possible length <= max_codeword_len.
472  *
473  * @max_codeword_len
474  *      Maximum length, in bits, of each codeword.
475  *
476  * @num_syms
477  *      Number of symbols in the alphabet, including symbols with zero
478  *      frequency.  This is the length of the 'A' and 'len' arrays.
479  */
480 static void
481 gen_codewords(u32 A[restrict], u8 lens[restrict],
482               const unsigned len_counts[restrict],
483               unsigned max_codeword_len, unsigned num_syms)
484 {
485         u32 next_codewords[max_codeword_len + 1];
486
487         /* Given the number of codewords that will have each length,
488          * assign codeword lengths to symbols.  We do this by assigning
489          * the lengths in decreasing order to the symbols sorted
490          * primarily by increasing frequency and secondarily by
491          * increasing symbol value.  */
492         for (unsigned i = 0, len = max_codeword_len; len >= 1; len--) {
493                 unsigned count = len_counts[len];
494                 while (count--)
495                         lens[A[i++] & SYMBOL_MASK] = len;
496         }
497
498         /* Generate the codewords themselves.  We initialize the
499          * 'next_codewords' array to provide the lexicographically first
500          * codeword of each length, then assign codewords in symbol
501          * order.  This produces a canonical code.  */
502         next_codewords[0] = 0;
503         next_codewords[1] = 0;
504         for (unsigned len = 2; len <= max_codeword_len; len++)
505                 next_codewords[len] =
506                         (next_codewords[len - 1] + len_counts[len - 1]) << 1;
507
508         for (unsigned sym = 0; sym < num_syms; sym++)
509                 A[sym] = next_codewords[lens[sym]]++;
510 }
511
512 /*
513  * ---------------------------------------------------------------------
514  *                      make_canonical_huffman_code()
515  * ---------------------------------------------------------------------
516  *
517  * Given an alphabet and the frequency of each symbol in it, construct a
518  * length-limited canonical Huffman code.
519  *
520  * @num_syms
521  *      The number of symbols in the alphabet.  The symbols are the
522  *      integers in the range [0, num_syms - 1].  This parameter must be
523  *      at least 2 and can't be greater than (1 << NUM_SYMBOL_BITS).
524  *
525  * @max_codeword_len
526  *      The maximum permissible codeword length.
527  *
528  * @freqs
529  *      An array of @num_syms entries, each of which specifies the
530  *      frequency of the corresponding symbol.  It is valid for some,
531  *      none, or all of the frequencies to be 0.
532  *
533  * @lens
534  *      An array of @num_syms entries in which this function will return
535  *      the length, in bits, of the codeword assigned to each symbol.
536  *      Symbols with 0 frequency will not have codewords per se, but
537  *      their entries in this array will be set to 0.  No lengths greater
538  *      than @max_codeword_len will be assigned.
539  *
540  * @codewords
541  *      An array of @num_syms entries in which this function will return
542  *      the codeword for each symbol, right-justified and padded on the
543  *      left with zeroes.  Codewords for symbols with 0 frequency will be
544  *      undefined.
545  *
546  * ---------------------------------------------------------------------
547  *
548  * This function builds a length-limited canonical Huffman code.
549  *
550  * A length-limited Huffman code contains no codewords longer than some
551  * specified length, and has exactly (with some algorithms) or
552  * approximately (with the algorithm used here) the minimum weighted path
553  * length from the root, given this constraint.
554  *
555  * A canonical Huffman code satisfies the properties that a codeword
556  * never lexicographically precedes a shorter codeword, and the
557  * lexicographic ordering of codewords of the same length is the same as
558  * the lexicographic ordering of the corresponding symbols.  A canonical
559  * Huffman code, or more generally a canonical prefix code, can be
560  * reconstructed from only a list containing the codeword length of each
561  * symbol.
562  *
563  * The classic algorithm to generate a Huffman code creates a node for
564  * each symbol, then inserts these nodes into a min-heap keyed by symbol
565  * frequency.  Then, repeatedly, the two lowest-frequency nodes are
566  * removed from the min-heap and added as the children of a new node
567  * having frequency equal to the sum of its two children, which is then
568  * inserted into the min-heap.  When only a single node remains in the
569  * min-heap, it is the root of the Huffman tree.  The codeword for each
570  * symbol is determined by the path needed to reach the corresponding
571  * node from the root.  Descending to the left child appends a 0 bit,
572  * whereas descending to the right child appends a 1 bit.
573  *
574  * The classic algorithm is relatively easy to understand, but it is
575  * subject to a number of inefficiencies.  In practice, it is fastest to
576  * first sort the symbols by frequency.  (This itself can be subject to
577  * an optimization based on the fact that most frequencies tend to be
578  * low.)  At the same time, we sort secondarily by symbol value, which
579  * aids the process of generating a canonical code.  Then, during tree
580  * construction, no heap is necessary because both the leaf nodes and the
581  * unparented non-leaf nodes can be easily maintained in sorted order.
582  * Consequently, there can never be more than two possibilities for the
583  * next-lowest-frequency node.
584  *
585  * In addition, because we're generating a canonical code, we actually
586  * don't need the leaf nodes of the tree at all, only the non-leaf nodes.
587  * This is because for canonical code generation we don't need to know
588  * where the symbols are in the tree.  Rather, we only need to know how
589  * many leaf nodes have each depth (codeword length).  And this
590  * information can, in fact, be quickly generated from the tree of
591  * non-leaves only.
592  *
593  * Furthermore, we can build this stripped-down Huffman tree directly in
594  * the array in which the codewords are to be generated, provided that
595  * these array slots are large enough to hold a symbol and frequency
596  * value.
597  *
598  * Still furthermore, we don't even need to maintain explicit child
599  * pointers.  We only need the parent pointers, and even those can be
600  * overwritten in-place with depth information as part of the process of
601  * extracting codeword lengths from the tree.  So in summary, we do NOT
602  * need a big structure like:
603  *
604  *      struct huffman_tree_node {
605  *              unsigned int symbol;
606  *              unsigned int frequency;
607  *              unsigned int depth;
608  *              struct huffman_tree_node *left_child;
609  *              struct huffman_tree_node *right_child;
610  *      };
611  *
612  *
613  *   ... which often gets used in "naive" implementations of Huffman code
614  *   generation.
615  *
616  * Most of these optimizations are based on the implementation in 7-Zip
617  * (source file: C/HuffEnc.c), which has been placed in the public domain
618  * by Igor Pavlov.  But I've rewritten the code with extensive comments,
619  * as it took me a while to figure out what it was doing...!
620  *
621  * ---------------------------------------------------------------------
622  *
623  * NOTE: in general, the same frequencies can be used to generate
624  * different length-limited canonical Huffman codes.  One choice we have
625  * is during tree construction, when we must decide whether to prefer a
626  * leaf or non-leaf when there is a tie in frequency.  Another choice we
627  * have is how to deal with codewords that would exceed @max_codeword_len
628  * bits in length.  Both of these choices affect the resulting codeword
629  * lengths, which otherwise can be mapped uniquely onto the resulting
630  * canonical Huffman code.
631  *
632  * Normally, there is no problem with choosing one valid code over
633  * another, provided that they produce similar compression ratios.
634  * However, the LZMS compression format uses adaptive Huffman coding.  It
635  * requires that both the decompressor and compressor build a canonical
636  * code equivalent to that which can be generated by using the classic
637  * Huffman tree construction algorithm and always processing leaves
638  * before non-leaves when there is a frequency tie.  Therefore, we make
639  * sure to do this.  This method also has the advantage of sometimes
640  * shortening the longest codeword that is generated.
641  *
642  * There also is the issue of how codewords longer than @max_codeword_len
643  * are dealt with.  Fortunately, for LZMS this is irrelevant because
644  * because for the LZMS alphabets no codeword can ever exceed
645  * LZMS_MAX_CODEWORD_LEN (= 15).  Since the LZMS algorithm regularly
646  * halves all frequencies, the frequencies cannot become high enough for
647  * a length 16 codeword to be generated.  Specifically, I think that if
648  * ties are broken in favor of non-leaves (as we do), the lowest total
649  * frequency that would give a length-16 codeword would be the sum of the
650  * frequencies 1 1 1 3 4 7 11 18 29 47 76 123 199 322 521 843 1364, which
651  * is 3570.  And in LZMS we can't get a frequency that high based on the
652  * alphabet sizes, rebuild frequencies, and scaling factors.  This
653  * worst-case scenario is based on the following degenerate case (only
654  * the bottom of the tree shown):
655  *
656  *                          ...
657  *                        17
658  *                       /  \
659  *                      10   7
660  *                     / \
661  *                    6   4
662  *                   / \
663  *                  3   3
664  *                 / \
665  *                2   1
666  *               / \
667  *              1   1
668  *
669  * Excluding the first leaves (those with value 1), each leaf value must
670  * be greater than the non-leaf up 1 and down 2 from it; otherwise that
671  * leaf would have taken precedence over that non-leaf and been combined
672  * with the leaf below, thereby decreasing the height compared to that
673  * shown.
674  *
675  * Interesting fact: if we were to instead prioritize non-leaves over
676  * leaves, then the worst case frequencies would be the Fibonacci
677  * sequence, plus an extra frequency of 1.  In this hypothetical
678  * scenario, it would be slightly easier for longer codewords to be
679  * generated.
680  */
681 void
682 make_canonical_huffman_code(unsigned num_syms, unsigned max_codeword_len,
683                             const u32 freqs[restrict],
684                             u8 lens[restrict], u32 codewords[restrict])
685 {
686         u32 *A = codewords;
687         unsigned num_used_syms;
688
689         /* Assumptions  */
690         wimlib_assert2(num_syms >= 2);
691         wimlib_assert2(num_syms <= (1 << NUM_SYMBOL_BITS));
692         wimlib_assert2(max_codeword_len > 0);
693         wimlib_assert2(max_codeword_len <= 32);
694
695         /* We begin by sorting the symbols primarily by frequency and
696          * secondarily by symbol value.  As an optimization, the array
697          * used for this purpose ('A') shares storage with the space in
698          * which we will eventually return the codewords.  */
699
700         num_used_syms = sort_symbols(num_syms, freqs, lens, A);
701
702         /* 'num_used_syms' is the number of symbols with nonzero
703          * frequency.  This may be less than @num_syms.  'num_used_syms'
704          * is also the number of entries in 'A' that are valid.  Each
705          * entry consists of a distinct symbol and a nonzero frequency
706          * packed into a 32-bit integer.  */
707
708         /* Handle special cases where only 0 or 1 symbols were used (had
709          * nonzero frequency).  */
710
711         if (unlikely(num_used_syms == 0)) {
712                 /* Code is empty.  sort_symbols() already set all lengths
713                  * to 0, so there is nothing more to do.  */
714                 return;
715         }
716
717         if (unlikely(num_used_syms == 1)) {
718                 /* Only one symbol was used, so we only need one
719                  * codeword.  But two codewords are needed to form the
720                  * smallest complete Huffman code, which uses codewords 0
721                  * and 1.  Therefore, we choose another symbol to which
722                  * to assign a codeword.  We use 0 (if the used symbol is
723                  * not 0) or 1 (if the used symbol is 0).  In either
724                  * case, the lesser-valued symbol must be assigned
725                  * codeword 0 so that the resulting code is canonical.  */
726
727                 unsigned sym = A[0] & SYMBOL_MASK;
728                 unsigned nonzero_idx = sym ? sym : 1;
729
730                 codewords[0] = 0;
731                 lens[0] = 1;
732                 codewords[nonzero_idx] = 1;
733                 lens[nonzero_idx] = 1;
734                 return;
735         }
736
737         /* Build a stripped-down version of the Huffman tree, sharing the
738          * array 'A' with the symbol values.  Then extract length counts
739          * from the tree and use them to generate the final codewords.  */
740
741         build_tree(A, num_used_syms);
742
743         {
744                 unsigned len_counts[max_codeword_len + 1];
745
746                 compute_length_counts(A, num_used_syms - 2,
747                                       len_counts, max_codeword_len);
748
749                 gen_codewords(A, lens, len_counts, max_codeword_len, num_syms);
750         }
751 }