1 // SPDX-License-Identifier: GPL-2.0-or-later
2 /*
3  * decompress_common.c - Code shared by the XPRESS and LZX decompressors
4  *
5  * Copyright (C) 2015 Eric Biggers
6  */
7 
8 #include "decompress_common.h"
9 
10 /*
11  * make_huffman_decode_table() -
12  *
13  * Build a decoding table for a canonical prefix code, or "Huffman code".
14  *
15  * This is an internal function, not part of the library API!
16  *
17  * This takes as input the length of the codeword for each symbol in the
18  * alphabet and produces as output a table that can be used for fast
19  * decoding of prefix-encoded symbols using read_huffsym().
20  *
21  * Strictly speaking, a canonical prefix code might not be a Huffman
22  * code.  But this algorithm will work either way; and in fact, since
23  * Huffman codes are defined in terms of symbol frequencies, there is no
24  * way for the decompressor to know whether the code is a true Huffman
25  * code or not until all symbols have been decoded.
26  *
27  * Because the prefix code is assumed to be "canonical", it can be
28  * reconstructed directly from the codeword lengths.  A prefix code is
29  * canonical if and only if a longer codeword never lexicographically
30  * precedes a shorter codeword, and the lexicographic ordering of
31  * codewords of the same length is the same as the lexicographic ordering
32  * of the corresponding symbols.  Consequently, we can sort the symbols
33  * primarily by codeword length and secondarily by symbol value, then
34  * reconstruct the prefix code by generating codewords lexicographically
35  * in that order.
36  *
37  * This function does not, however, generate the prefix code explicitly.
38  * Instead, it directly builds a table for decoding symbols using the
39  * code.  The basic idea is this: given the next 'max_codeword_len' bits
40  * in the input, we can look up the decoded symbol by indexing a table
41  * containing 2**max_codeword_len entries.  A codeword with length
42  * 'max_codeword_len' will have exactly one entry in this table, whereas
43  * a codeword shorter than 'max_codeword_len' will have multiple entries
44  * in this table.  Precisely, a codeword of length n will be represented
45  * by 2**(max_codeword_len - n) entries in this table.  The 0-based index
46  * of each such entry will contain the corresponding codeword as a prefix
47  * when zero-padded on the left to 'max_codeword_len' binary digits.
48  *
49  * That's the basic idea, but we implement two optimizations regarding
50  * the format of the decode table itself:
51  *
52  * - For many compression formats, the maximum codeword length is too
53  *   long for it to be efficient to build the full decoding table
54  *   whenever a new prefix code is used.  Instead, we can build the table
55  *   using only 2**table_bits entries, where 'table_bits' is some number
56  *   less than or equal to 'max_codeword_len'.  Then, only codewords of
57  *   length 'table_bits' and shorter can be directly looked up.  For
58  *   longer codewords, the direct lookup instead produces the root of a
59  *   binary tree.  Using this tree, the decoder can do traditional
60  *   bit-by-bit decoding of the remainder of the codeword.  Child nodes
61  *   are allocated in extra entries at the end of the table; leaf nodes
62  *   contain symbols.  Note that the long-codeword case is, in general,
63  *   not performance critical, since in Huffman codes the most frequently
64  *   used symbols are assigned the shortest codeword lengths.
65  *
66  * - When we decode a symbol using a direct lookup of the table, we still
67  *   need to know its length so that the bitstream can be advanced by the
68  *   appropriate number of bits.  The simple solution is to simply retain
69  *   the 'lens' array and use the decoded symbol as an index into it.
70  *   However, this requires two separate array accesses in the fast path.
71  *   The optimization is to store the length directly in the decode
72  *   table.  We use the bottom 11 bits for the symbol and the top 5 bits
73  *   for the length.  In addition, to combine this optimization with the
74  *   previous one, we introduce a special case where the top 2 bits of
75  *   the length are both set if the entry is actually the root of a
76  *   binary tree.
77  *
78  * @decode_table:
79  *	The array in which to create the decoding table.  This must have
80  *	a length of at least ((2**table_bits) + 2 * num_syms) entries.
81  *
82  * @num_syms:
83  *	The number of symbols in the alphabet; also, the length of the
84  *	'lens' array.  Must be less than or equal to 2048.
85  *
86  * @table_bits:
87  *	The order of the decode table size, as explained above.  Must be
88  *	less than or equal to 13.
89  *
90  * @lens:
91  *	An array of length @num_syms, indexable by symbol, that gives the
92  *	length of the codeword, in bits, for that symbol.  The length can
93  *	be 0, which means that the symbol does not have a codeword
94  *	assigned.
95  *
96  * @max_codeword_len:
97  *	The longest codeword length allowed in the compression format.
98  *	All entries in 'lens' must be less than or equal to this value.
99  *	This must be less than or equal to 23.
100  *
101  * @working_space
102  *	A temporary array of length '2 * (max_codeword_len + 1) +
103  *	num_syms'.
104  *
105  * Returns 0 on success, or -1 if the lengths do not form a valid prefix
106  * code.
107  */
make_huffman_decode_table(u16 decode_table[],const u32 num_syms,const u32 table_bits,const u8 lens[],const u32 max_codeword_len,u16 working_space[])108 int make_huffman_decode_table(u16 decode_table[], const u32 num_syms,
109 			      const u32 table_bits, const u8 lens[],
110 			      const u32 max_codeword_len,
111 			      u16 working_space[])
112 {
113 	const u32 table_num_entries = 1 << table_bits;
114 	u16 * const len_counts = &working_space[0];
115 	u16 * const offsets = &working_space[1 * (max_codeword_len + 1)];
116 	u16 * const sorted_syms = &working_space[2 * (max_codeword_len + 1)];
117 	int left;
118 	void *decode_table_ptr;
119 	u32 sym_idx;
120 	u32 codeword_len;
121 	u32 stores_per_loop;
122 	u32 decode_table_pos;
123 	u32 len;
124 	u32 sym;
125 
126 	/* Count how many symbols have each possible codeword length.
127 	 * Note that a length of 0 indicates the corresponding symbol is not
128 	 * used in the code and therefore does not have a codeword.
129 	 */
130 	for (len = 0; len <= max_codeword_len; len++)
131 		len_counts[len] = 0;
132 	for (sym = 0; sym < num_syms; sym++)
133 		len_counts[lens[sym]]++;
134 
135 	/* We can assume all lengths are <= max_codeword_len, but we
136 	 * cannot assume they form a valid prefix code.  A codeword of
137 	 * length n should require a proportion of the codespace equaling
138 	 * (1/2)^n.  The code is valid if and only if the codespace is
139 	 * exactly filled by the lengths, by this measure.
140 	 */
141 	left = 1;
142 	for (len = 1; len <= max_codeword_len; len++) {
143 		left <<= 1;
144 		left -= len_counts[len];
145 		if (left < 0) {
146 			/* The lengths overflow the codespace; that is, the code
147 			 * is over-subscribed.
148 			 */
149 			return -1;
150 		}
151 	}
152 
153 	if (left) {
154 		/* The lengths do not fill the codespace; that is, they form an
155 		 * incomplete set.
156 		 */
157 		if (left == (1 << max_codeword_len)) {
158 			/* The code is completely empty.  This is arguably
159 			 * invalid, but in fact it is valid in LZX and XPRESS,
160 			 * so we must allow it.  By definition, no symbols can
161 			 * be decoded with an empty code.  Consequently, we
162 			 * technically don't even need to fill in the decode
163 			 * table.  However, to avoid accessing uninitialized
164 			 * memory if the algorithm nevertheless attempts to
165 			 * decode symbols using such a code, we zero out the
166 			 * decode table.
167 			 */
168 			memset(decode_table, 0,
169 			       table_num_entries * sizeof(decode_table[0]));
170 			return 0;
171 		}
172 		return -1;
173 	}
174 
175 	/* Sort the symbols primarily by length and secondarily by symbol order.
176 	 */
177 
178 	/* Initialize 'offsets' so that offsets[len] for 1 <= len <=
179 	 * max_codeword_len is the number of codewords shorter than 'len' bits.
180 	 */
181 	offsets[1] = 0;
182 	for (len = 1; len < max_codeword_len; len++)
183 		offsets[len + 1] = offsets[len] + len_counts[len];
184 
185 	/* Use the 'offsets' array to sort the symbols.  Note that we do not
186 	 * include symbols that are not used in the code.  Consequently, fewer
187 	 * than 'num_syms' entries in 'sorted_syms' may be filled.
188 	 */
189 	for (sym = 0; sym < num_syms; sym++)
190 		if (lens[sym])
191 			sorted_syms[offsets[lens[sym]]++] = sym;
192 
193 	/* Fill entries for codewords with length <= table_bits
194 	 * --- that is, those short enough for a direct mapping.
195 	 *
196 	 * The table will start with entries for the shortest codeword(s), which
197 	 * have the most entries.  From there, the number of entries per
198 	 * codeword will decrease.
199 	 */
200 	decode_table_ptr = decode_table;
201 	sym_idx = 0;
202 	codeword_len = 1;
203 	stores_per_loop = (1 << (table_bits - codeword_len));
204 	for (; stores_per_loop != 0; codeword_len++, stores_per_loop >>= 1) {
205 		u32 end_sym_idx = sym_idx + len_counts[codeword_len];
206 
207 		for (; sym_idx < end_sym_idx; sym_idx++) {
208 			u16 entry;
209 			u16 *p;
210 			u32 n;
211 
212 			entry = ((u32)codeword_len << 11) | sorted_syms[sym_idx];
213 			p = (u16 *)decode_table_ptr;
214 			n = stores_per_loop;
215 
216 			do {
217 				*p++ = entry;
218 			} while (--n);
219 
220 			decode_table_ptr = p;
221 		}
222 	}
223 
224 	/* If we've filled in the entire table, we are done.  Otherwise,
225 	 * there are codewords longer than table_bits for which we must
226 	 * generate binary trees.
227 	 */
228 	decode_table_pos = (u16 *)decode_table_ptr - decode_table;
229 	if (decode_table_pos != table_num_entries) {
230 		u32 j;
231 		u32 next_free_tree_slot;
232 		u32 cur_codeword;
233 
234 		/* First, zero out the remaining entries.  This is
235 		 * necessary so that these entries appear as
236 		 * "unallocated" in the next part.  Each of these entries
237 		 * will eventually be filled with the representation of
238 		 * the root node of a binary tree.
239 		 */
240 		j = decode_table_pos;
241 		do {
242 			decode_table[j] = 0;
243 		} while (++j != table_num_entries);
244 
245 		/* We allocate child nodes starting at the end of the
246 		 * direct lookup table.  Note that there should be
247 		 * 2*num_syms extra entries for this purpose, although
248 		 * fewer than this may actually be needed.
249 		 */
250 		next_free_tree_slot = table_num_entries;
251 
252 		/* Iterate through each codeword with length greater than
253 		 * 'table_bits', primarily in order of codeword length
254 		 * and secondarily in order of symbol.
255 		 */
256 		for (cur_codeword = decode_table_pos << 1;
257 		     codeword_len <= max_codeword_len;
258 		     codeword_len++, cur_codeword <<= 1) {
259 			u32 end_sym_idx = sym_idx + len_counts[codeword_len];
260 
261 			for (; sym_idx < end_sym_idx; sym_idx++, cur_codeword++) {
262 				/* 'sorted_sym' is the symbol represented by the
263 				 * codeword.
264 				 */
265 				u32 sorted_sym = sorted_syms[sym_idx];
266 				u32 extra_bits = codeword_len - table_bits;
267 				u32 node_idx = cur_codeword >> extra_bits;
268 
269 				/* Go through each bit of the current codeword
270 				 * beyond the prefix of length @table_bits and
271 				 * walk the appropriate binary tree, allocating
272 				 * any slots that have not yet been allocated.
273 				 *
274 				 * Note that the 'pointer' entry to the binary
275 				 * tree, which is stored in the direct lookup
276 				 * portion of the table, is represented
277 				 * identically to other internal (non-leaf)
278 				 * nodes of the binary tree; it can be thought
279 				 * of as simply the root of the tree.  The
280 				 * representation of these internal nodes is
281 				 * simply the index of the left child combined
282 				 * with the special bits 0xC000 to distinguish
283 				 * the entry from direct mapping and leaf node
284 				 * entries.
285 				 */
286 				do {
287 					/* At least one bit remains in the
288 					 * codeword, but the current node is an
289 					 * unallocated leaf.  Change it to an
290 					 * internal node.
291 					 */
292 					if (decode_table[node_idx] == 0) {
293 						decode_table[node_idx] =
294 							next_free_tree_slot | 0xC000;
295 						decode_table[next_free_tree_slot++] = 0;
296 						decode_table[next_free_tree_slot++] = 0;
297 					}
298 
299 					/* Go to the left child if the next bit
300 					 * in the codeword is 0; otherwise go to
301 					 * the right child.
302 					 */
303 					node_idx = decode_table[node_idx] & 0x3FFF;
304 					--extra_bits;
305 					node_idx += (cur_codeword >> extra_bits) & 1;
306 				} while (extra_bits != 0);
307 
308 				/* We've traversed the tree using the entire
309 				 * codeword, and we're now at the entry where
310 				 * the actual symbol will be stored.  This is
311 				 * distinguished from internal nodes by not
312 				 * having its high two bits set.
313 				 */
314 				decode_table[node_idx] = sorted_sym;
315 			}
316 		}
317 	}
318 	return 0;
319 }
320