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