1 // SPDX-License-Identifier: MIT 2 /* 3 * decompress_common.c - Code shared by the XPRESS and LZX decompressors 4 * 5 * This is a port of the upstream wimlib "decompress_common.c" which builds 6 * subtable-based Huffman decode tables, as opposed to the older 7 * binary-tree-based format previously used in this library. The vectorized 8 * (SSE2/AVX2) fill paths are omitted for portability in the kernel. 9 * 10 * Copyright (C) 2022 Eric Biggers 11 */ 12 13 #include "decompress_common.h" 14 15 /* Compute the number of bits with which a subtable must be indexed for a 16 * codeword of length @codeword_len, given that the root table is indexed with 17 * @table_bits bits. 18 */ 19 static u32 compute_subtable_bits(u32 table_bits, 20 u32 codeword_len, u16 len_counts[]) 21 { 22 u32 subtable_bits = codeword_len - table_bits; 23 s32 remainder = (s32)1 << subtable_bits; 24 25 for (;;) { 26 remainder -= len_counts[table_bits + subtable_bits]; 27 if (remainder <= 0) 28 break; 29 subtable_bits++; 30 remainder <<= 1; 31 } 32 return subtable_bits; 33 } 34 35 /* Build the subtables for codewords longer than table_bits. */ 36 static int build_subtables(u16 decode_table[], u32 num_syms, u32 table_bits, 37 u16 len_counts[], const u16 sorted_syms[], u32 sym_idx, 38 u32 decode_table_pos, u32 decode_table_size) 39 { 40 u32 subtable_pos = 1U << table_bits; 41 u32 subtable_bits = table_bits; 42 u32 subtable_prefix = (u32)-1; 43 u32 codeword_len = table_bits + 1; 44 u32 codeword = decode_table_pos << 1; 45 u32 prefix; 46 u16 entry; 47 u32 n; 48 49 for (; sym_idx < num_syms; sym_idx++) { 50 while (len_counts[codeword_len] == 0) { 51 codeword_len++; 52 codeword <<= 1; 53 } 54 55 prefix = codeword >> (codeword_len - table_bits); 56 57 if (prefix != subtable_prefix) { 58 subtable_prefix = prefix; 59 subtable_bits = compute_subtable_bits(table_bits, codeword_len, 60 len_counts); 61 decode_table[subtable_prefix] = 62 MAKE_DECODE_TABLE_ENTRY(subtable_pos, subtable_bits); 63 } 64 65 entry = MAKE_DECODE_TABLE_ENTRY(sorted_syms[sym_idx], 66 codeword_len - table_bits); 67 n = 1U << (subtable_bits - (codeword_len - table_bits)); 68 69 /* Defensive bound check: 'lens' is derived from untrusted 70 * on-disk compressed data, and subtable growth depends on 71 * its content. This should never trigger for a correctly 72 * sized DECODE_TABLE_ENOUGH() value, but turns a wrong value 73 * into a clean decode failure instead of writing past the 74 * caller's decode_table[]. 75 */ 76 if (unlikely(subtable_pos + n > decode_table_size)) 77 return -1; 78 79 do { 80 decode_table[subtable_pos++] = entry; 81 } while (--n); 82 83 len_counts[codeword_len]--; 84 codeword++; 85 } 86 87 return 0; 88 } 89 90 /* 91 * Given an alphabet of symbols and the length of each symbol's codeword in a 92 * canonical prefix code, build a table for quickly decoding symbols that were 93 * encoded with that code. 94 * 95 * The root table is indexed with 'table_bits' bits. Codewords not longer than 96 * 'table_bits' are decoded directly from the root table. Longer codewords are 97 * decoded via subtables: the corresponding root entry is a pointer (the index 98 * of the subtable plus the number of bits with which the subtable is indexed), 99 * and the subtable is indexed with the remaining bits of the codeword. 100 * 101 * Each entry stores both the symbol (high 12 bits) and the codeword length (low 102 * 4 bits), so a single lookup yields the symbol and lets the bitstream be 103 * advanced by the correct number of bits. 104 * 105 * @decode_table: array in which to build the table (declared with 106 * DECODE_TABLE()). May alias @lens. 107 * @num_syms: number of symbols in the alphabet. 108 * @table_bits: log2 of the number of root table entries. 109 * @lens: array of @num_syms codeword lengths, indexed by symbol. 110 * @max_codeword_len: longest codeword length allowed for this code. 111 * @working_space: temporary array declared with DECODE_TABLE_WORKING_SPACE(). 112 * @decode_table_size: number of u16 entries in @decode_table (i.e. 113 * ARRAY_SIZE(decode_table) at the call site). Used only as a 114 * defensive bound check against @lens-dependent subtable growth. 115 * 116 * Returns 0 on success, or -1 if the lengths do not form a valid prefix code, 117 * or if building the subtables would overflow @decode_table_size entries. 118 */ 119 int make_huffman_decode_table(u16 decode_table[], u32 num_syms, u32 table_bits, 120 const u8 lens[], u32 max_codeword_len, 121 u16 working_space[], u32 decode_table_size) 122 { 123 u16 *const len_counts = &working_space[0]; 124 u16 *const offsets = &working_space[1 * (max_codeword_len + 1)]; 125 u16 *const sorted_syms = &working_space[2 * (max_codeword_len + 1)]; 126 u32 decode_table_pos = 0; 127 u32 sym_idx; 128 u32 codeword_len; 129 s32 remainder = 1; 130 void *entry_ptr = decode_table; 131 u32 len; 132 u32 sym; 133 134 /* Count how many codewords have each length, including 0. */ 135 for (len = 0; len <= max_codeword_len; len++) 136 len_counts[len] = 0; 137 for (sym = 0; sym < num_syms; sym++) 138 len_counts[lens[sym]]++; 139 140 /* A codeword of length n should require a proportion of the codespace 141 * equaling (1/2)^n. The code is complete iff the codespace is exactly 142 * filled by the lengths. 143 */ 144 for (len = 1; len <= max_codeword_len; len++) { 145 remainder = (remainder << 1) - len_counts[len]; 146 if (unlikely(remainder < 0)) 147 return -1; /* over-subscribed */ 148 } 149 150 if (remainder != 0) { 151 /* Incomplete code. Permitted only if the code is empty. */ 152 if (unlikely(remainder != (s32)(1U << max_codeword_len))) 153 return -1; 154 155 /* Empty code: zero the root table so lookups yield symbol 0 156 * without consuming any bits. 157 */ 158 memset(decode_table, 0, sizeof(decode_table[0]) << table_bits); 159 return 0; 160 } 161 162 /* Sort the symbols primarily by increasing codeword length and 163 * secondarily by increasing symbol value. 164 */ 165 offsets[0] = 0; 166 for (len = 0; len < max_codeword_len; len++) 167 offsets[len + 1] = offsets[len] + len_counts[len]; 168 for (sym = 0; sym < num_syms; sym++) 169 sorted_syms[offsets[lens[sym]]++] = sym; 170 171 /* Fill the root table entries for codewords no longer than table_bits. */ 172 sym_idx = offsets[0]; 173 codeword_len = 1; 174 for (; codeword_len <= table_bits; codeword_len++) { 175 u32 stores_per_loop = 1U << (table_bits - codeword_len); 176 u32 end_sym_idx = sym_idx + len_counts[codeword_len]; 177 178 for (; sym_idx < end_sym_idx; sym_idx++) { 179 u16 v = MAKE_DECODE_TABLE_ENTRY(sorted_syms[sym_idx], 180 codeword_len); 181 u32 n = stores_per_loop; 182 u16 *p = entry_ptr; 183 184 do { 185 *p++ = v; 186 } while (--n); 187 entry_ptr = p; 188 } 189 } 190 decode_table_pos = (u16 *)entry_ptr - decode_table; 191 192 /* If all symbols were processed, no subtables are required. */ 193 if (sym_idx == num_syms) 194 return 0; 195 196 /* At least one subtable is required. Process the remaining symbols. */ 197 return build_subtables(decode_table, num_syms, table_bits, len_counts, 198 sorted_syms, sym_idx, decode_table_pos, 199 decode_table_size); 200 } 201