xref: /linux/lib/zstd/common/entropy_common.c (revision e61f33273ca755b3e2ebee4520a76097199dc7a8)
1 // SPDX-License-Identifier: GPL-2.0+ OR BSD-3-Clause
2 /* ******************************************************************
3  * Common functions of New Generation Entropy library
4  * Copyright (c) Meta Platforms, Inc. and affiliates.
5  *
6  *  You can contact the author at :
7  *  - FSE+HUF source repository : https://github.com/Cyan4973/FiniteStateEntropy
8  *  - Public forum : https://groups.google.com/forum/#!forum/lz4c
9  *
10  * This source code is licensed under both the BSD-style license (found in the
11  * LICENSE file in the root directory of this source tree) and the GPLv2 (found
12  * in the COPYING file in the root directory of this source tree).
13  * You may select, at your option, one of the above-listed licenses.
14 ****************************************************************** */
15 
16 /* *************************************
17 *  Dependencies
18 ***************************************/
19 #include "mem.h"
20 #include "error_private.h"       /* ERR_*, ERROR */
21 #define FSE_STATIC_LINKING_ONLY  /* FSE_MIN_TABLELOG */
22 #include "fse.h"
23 #include "huf.h"
24 #include "bits.h"                /* ZSDT_highbit32, ZSTD_countTrailingZeros32 */
25 
26 
27 /*===   Version   ===*/
FSE_versionNumber(void)28 unsigned FSE_versionNumber(void) { return FSE_VERSION_NUMBER; }
29 
30 
31 /*===   Error Management   ===*/
FSE_isError(size_t code)32 unsigned FSE_isError(size_t code) { return ERR_isError(code); }
FSE_getErrorName(size_t code)33 const char* FSE_getErrorName(size_t code) { return ERR_getErrorName(code); }
34 
HUF_isError(size_t code)35 unsigned HUF_isError(size_t code) { return ERR_isError(code); }
HUF_getErrorName(size_t code)36 const char* HUF_getErrorName(size_t code) { return ERR_getErrorName(code); }
37 
38 
39 /*-**************************************************************
40 *  FSE NCount encoding-decoding
41 ****************************************************************/
42 FORCE_INLINE_TEMPLATE
FSE_readNCount_body(short * normalizedCounter,unsigned * maxSVPtr,unsigned * tableLogPtr,const void * headerBuffer,size_t hbSize)43 size_t FSE_readNCount_body(short* normalizedCounter, unsigned* maxSVPtr, unsigned* tableLogPtr,
44                            const void* headerBuffer, size_t hbSize)
45 {
46     const BYTE* const istart = (const BYTE*) headerBuffer;
47     const BYTE* const iend = istart + hbSize;
48     const BYTE* ip = istart;
49     int nbBits;
50     int remaining;
51     int threshold;
52     U32 bitStream;
53     int bitCount;
54     unsigned charnum = 0;
55     unsigned const maxSV1 = *maxSVPtr + 1;
56     int previous0 = 0;
57 
58     if (hbSize < 8) {
59         /* This function only works when hbSize >= 8 */
60         char buffer[8] = {0};
61         ZSTD_memcpy(buffer, headerBuffer, hbSize);
62         {   size_t const countSize = FSE_readNCount(normalizedCounter, maxSVPtr, tableLogPtr,
63                                                     buffer, sizeof(buffer));
64             if (FSE_isError(countSize)) return countSize;
65             if (countSize > hbSize) return ERROR(corruption_detected);
66             return countSize;
67     }   }
68     assert(hbSize >= 8);
69 
70     /* init */
71     ZSTD_memset(normalizedCounter, 0, (*maxSVPtr+1) * sizeof(normalizedCounter[0]));   /* all symbols not present in NCount have a frequency of 0 */
72     bitStream = MEM_readLE32(ip);
73     nbBits = (bitStream & 0xF) + FSE_MIN_TABLELOG;   /* extract tableLog */
74     if (nbBits > FSE_TABLELOG_ABSOLUTE_MAX) return ERROR(tableLog_tooLarge);
75     bitStream >>= 4;
76     bitCount = 4;
77     *tableLogPtr = nbBits;
78     remaining = (1<<nbBits)+1;
79     threshold = 1<<nbBits;
80     nbBits++;
81 
82     for (;;) {
83         if (previous0) {
84             /* Count the number of repeats. Each time the
85              * 2-bit repeat code is 0b11 there is another
86              * repeat.
87              * Avoid UB by setting the high bit to 1.
88              */
89             int repeats = ZSTD_countTrailingZeros32(~bitStream | 0x80000000) >> 1;
90             while (repeats >= 12) {
91                 charnum += 3 * 12;
92                 if (LIKELY(ip <= iend-7)) {
93                     ip += 3;
94                 } else {
95                     bitCount -= (int)(8 * (iend - 7 - ip));
96                     bitCount &= 31;
97                     ip = iend - 4;
98                 }
99                 bitStream = MEM_readLE32(ip) >> bitCount;
100                 repeats = ZSTD_countTrailingZeros32(~bitStream | 0x80000000) >> 1;
101             }
102             charnum += 3 * repeats;
103             bitStream >>= 2 * repeats;
104             bitCount += 2 * repeats;
105 
106             /* Add the final repeat which isn't 0b11. */
107             assert((bitStream & 3) < 3);
108             charnum += bitStream & 3;
109             bitCount += 2;
110 
111             /* This is an error, but break and return an error
112              * at the end, because returning out of a loop makes
113              * it harder for the compiler to optimize.
114              */
115             if (charnum >= maxSV1) break;
116 
117             /* We don't need to set the normalized count to 0
118              * because we already memset the whole buffer to 0.
119              */
120 
121             if (LIKELY(ip <= iend-7) || (ip + (bitCount>>3) <= iend-4)) {
122                 assert((bitCount >> 3) <= 3); /* For first condition to work */
123                 ip += bitCount>>3;
124                 bitCount &= 7;
125             } else {
126                 bitCount -= (int)(8 * (iend - 4 - ip));
127                 bitCount &= 31;
128                 ip = iend - 4;
129             }
130             bitStream = MEM_readLE32(ip) >> bitCount;
131         }
132         {
133             int const max = (2*threshold-1) - remaining;
134             int count;
135 
136             if ((bitStream & (threshold-1)) < (U32)max) {
137                 count = bitStream & (threshold-1);
138                 bitCount += nbBits-1;
139             } else {
140                 count = bitStream & (2*threshold-1);
141                 if (count >= threshold) count -= max;
142                 bitCount += nbBits;
143             }
144 
145             count--;   /* extra accuracy */
146             /* When it matters (small blocks), this is a
147              * predictable branch, because we don't use -1.
148              */
149             if (count >= 0) {
150                 remaining -= count;
151             } else {
152                 assert(count == -1);
153                 remaining += count;
154             }
155             normalizedCounter[charnum++] = (short)count;
156             previous0 = !count;
157 
158             assert(threshold > 1);
159             if (remaining < threshold) {
160                 /* This branch can be folded into the
161                  * threshold update condition because we
162                  * know that threshold > 1.
163                  */
164                 if (remaining <= 1) break;
165                 nbBits = ZSTD_highbit32(remaining) + 1;
166                 threshold = 1 << (nbBits - 1);
167             }
168             if (charnum >= maxSV1) break;
169 
170             if (LIKELY(ip <= iend-7) || (ip + (bitCount>>3) <= iend-4)) {
171                 ip += bitCount>>3;
172                 bitCount &= 7;
173             } else {
174                 bitCount -= (int)(8 * (iend - 4 - ip));
175                 bitCount &= 31;
176                 ip = iend - 4;
177             }
178             bitStream = MEM_readLE32(ip) >> bitCount;
179     }   }
180     if (remaining != 1) return ERROR(corruption_detected);
181     /* Only possible when there are too many zeros. */
182     if (charnum > maxSV1) return ERROR(maxSymbolValue_tooSmall);
183     if (bitCount > 32) return ERROR(corruption_detected);
184     *maxSVPtr = charnum-1;
185 
186     ip += (bitCount+7)>>3;
187     return ip-istart;
188 }
189 
190 /* Avoids the FORCE_INLINE of the _body() function. */
FSE_readNCount_body_default(short * normalizedCounter,unsigned * maxSVPtr,unsigned * tableLogPtr,const void * headerBuffer,size_t hbSize)191 static size_t FSE_readNCount_body_default(
192         short* normalizedCounter, unsigned* maxSVPtr, unsigned* tableLogPtr,
193         const void* headerBuffer, size_t hbSize)
194 {
195     return FSE_readNCount_body(normalizedCounter, maxSVPtr, tableLogPtr, headerBuffer, hbSize);
196 }
197 
198 #if DYNAMIC_BMI2
FSE_readNCount_body_bmi2(short * normalizedCounter,unsigned * maxSVPtr,unsigned * tableLogPtr,const void * headerBuffer,size_t hbSize)199 BMI2_TARGET_ATTRIBUTE static size_t FSE_readNCount_body_bmi2(
200         short* normalizedCounter, unsigned* maxSVPtr, unsigned* tableLogPtr,
201         const void* headerBuffer, size_t hbSize)
202 {
203     return FSE_readNCount_body(normalizedCounter, maxSVPtr, tableLogPtr, headerBuffer, hbSize);
204 }
205 #endif
206 
FSE_readNCount_bmi2(short * normalizedCounter,unsigned * maxSVPtr,unsigned * tableLogPtr,const void * headerBuffer,size_t hbSize,int bmi2)207 size_t FSE_readNCount_bmi2(
208         short* normalizedCounter, unsigned* maxSVPtr, unsigned* tableLogPtr,
209         const void* headerBuffer, size_t hbSize, int bmi2)
210 {
211 #if DYNAMIC_BMI2
212     if (bmi2) {
213         return FSE_readNCount_body_bmi2(normalizedCounter, maxSVPtr, tableLogPtr, headerBuffer, hbSize);
214     }
215 #endif
216     (void)bmi2;
217     return FSE_readNCount_body_default(normalizedCounter, maxSVPtr, tableLogPtr, headerBuffer, hbSize);
218 }
219 
FSE_readNCount(short * normalizedCounter,unsigned * maxSVPtr,unsigned * tableLogPtr,const void * headerBuffer,size_t hbSize)220 size_t FSE_readNCount(
221         short* normalizedCounter, unsigned* maxSVPtr, unsigned* tableLogPtr,
222         const void* headerBuffer, size_t hbSize)
223 {
224     return FSE_readNCount_bmi2(normalizedCounter, maxSVPtr, tableLogPtr, headerBuffer, hbSize, /* bmi2 */ 0);
225 }
226 
227 
228 /*! HUF_readStats() :
229     Read compact Huffman tree, saved by HUF_writeCTable().
230     `huffWeight` is destination buffer.
231     `rankStats` is assumed to be a table of at least HUF_TABLELOG_MAX U32.
232     @return : size read from `src` , or an error Code .
233     Note : Needed by HUF_readCTable() and HUF_readDTableX?() .
234 */
HUF_readStats(BYTE * huffWeight,size_t hwSize,U32 * rankStats,U32 * nbSymbolsPtr,U32 * tableLogPtr,const void * src,size_t srcSize)235 size_t HUF_readStats(BYTE* huffWeight, size_t hwSize, U32* rankStats,
236                      U32* nbSymbolsPtr, U32* tableLogPtr,
237                      const void* src, size_t srcSize)
238 {
239     U32 wksp[HUF_READ_STATS_WORKSPACE_SIZE_U32];
240     return HUF_readStats_wksp(huffWeight, hwSize, rankStats, nbSymbolsPtr, tableLogPtr, src, srcSize, wksp, sizeof(wksp), /* flags */ 0);
241 }
242 
243 FORCE_INLINE_TEMPLATE size_t
HUF_readStats_body(BYTE * huffWeight,size_t hwSize,U32 * rankStats,U32 * nbSymbolsPtr,U32 * tableLogPtr,const void * src,size_t srcSize,void * workSpace,size_t wkspSize,int bmi2)244 HUF_readStats_body(BYTE* huffWeight, size_t hwSize, U32* rankStats,
245                    U32* nbSymbolsPtr, U32* tableLogPtr,
246                    const void* src, size_t srcSize,
247                    void* workSpace, size_t wkspSize,
248                    int bmi2)
249 {
250     U32 weightTotal;
251     const BYTE* ip = (const BYTE*) src;
252     size_t iSize;
253     size_t oSize;
254 
255     if (!srcSize) return ERROR(srcSize_wrong);
256     iSize = ip[0];
257     /* ZSTD_memset(huffWeight, 0, hwSize);   *//* is not necessary, even though some analyzer complain ... */
258 
259     if (iSize >= 128) {  /* special header */
260         oSize = iSize - 127;
261         iSize = ((oSize+1)/2);
262         if (iSize+1 > srcSize) return ERROR(srcSize_wrong);
263         if (oSize >= hwSize) return ERROR(corruption_detected);
264         ip += 1;
265         {   U32 n;
266             for (n=0; n<oSize; n+=2) {
267                 huffWeight[n]   = ip[n/2] >> 4;
268                 huffWeight[n+1] = ip[n/2] & 15;
269     }   }   }
270     else  {   /* header compressed with FSE (normal case) */
271         if (iSize+1 > srcSize) return ERROR(srcSize_wrong);
272         /* max (hwSize-1) values decoded, as last one is implied */
273         oSize = FSE_decompress_wksp_bmi2(huffWeight, hwSize-1, ip+1, iSize, 6, workSpace, wkspSize, bmi2);
274         if (FSE_isError(oSize)) return oSize;
275     }
276 
277     /* collect weight stats */
278     ZSTD_memset(rankStats, 0, (HUF_TABLELOG_MAX + 1) * sizeof(U32));
279     weightTotal = 0;
280     {   U32 n; for (n=0; n<oSize; n++) {
281             if (huffWeight[n] > HUF_TABLELOG_MAX) return ERROR(corruption_detected);
282             rankStats[huffWeight[n]]++;
283             weightTotal += (1 << huffWeight[n]) >> 1;
284     }   }
285     if (weightTotal == 0) return ERROR(corruption_detected);
286 
287     /* get last non-null symbol weight (implied, total must be 2^n) */
288     {   U32 const tableLog = ZSTD_highbit32(weightTotal) + 1;
289         if (tableLog > HUF_TABLELOG_MAX) return ERROR(corruption_detected);
290         *tableLogPtr = tableLog;
291         /* determine last weight */
292         {   U32 const total = 1 << tableLog;
293             U32 const rest = total - weightTotal;
294             U32 const verif = 1 << ZSTD_highbit32(rest);
295             U32 const lastWeight = ZSTD_highbit32(rest) + 1;
296             if (verif != rest) return ERROR(corruption_detected);    /* last value must be a clean power of 2 */
297             huffWeight[oSize] = (BYTE)lastWeight;
298             rankStats[lastWeight]++;
299     }   }
300 
301     /* check tree construction validity */
302     if ((rankStats[1] < 2) || (rankStats[1] & 1)) return ERROR(corruption_detected);   /* by construction : at least 2 elts of rank 1, must be even */
303 
304     /* results */
305     *nbSymbolsPtr = (U32)(oSize+1);
306     return iSize+1;
307 }
308 
309 /* Avoids the FORCE_INLINE of the _body() function. */
HUF_readStats_body_default(BYTE * huffWeight,size_t hwSize,U32 * rankStats,U32 * nbSymbolsPtr,U32 * tableLogPtr,const void * src,size_t srcSize,void * workSpace,size_t wkspSize)310 static size_t HUF_readStats_body_default(BYTE* huffWeight, size_t hwSize, U32* rankStats,
311                      U32* nbSymbolsPtr, U32* tableLogPtr,
312                      const void* src, size_t srcSize,
313                      void* workSpace, size_t wkspSize)
314 {
315     return HUF_readStats_body(huffWeight, hwSize, rankStats, nbSymbolsPtr, tableLogPtr, src, srcSize, workSpace, wkspSize, 0);
316 }
317 
318 #if DYNAMIC_BMI2
HUF_readStats_body_bmi2(BYTE * huffWeight,size_t hwSize,U32 * rankStats,U32 * nbSymbolsPtr,U32 * tableLogPtr,const void * src,size_t srcSize,void * workSpace,size_t wkspSize)319 static BMI2_TARGET_ATTRIBUTE size_t HUF_readStats_body_bmi2(BYTE* huffWeight, size_t hwSize, U32* rankStats,
320                      U32* nbSymbolsPtr, U32* tableLogPtr,
321                      const void* src, size_t srcSize,
322                      void* workSpace, size_t wkspSize)
323 {
324     return HUF_readStats_body(huffWeight, hwSize, rankStats, nbSymbolsPtr, tableLogPtr, src, srcSize, workSpace, wkspSize, 1);
325 }
326 #endif
327 
HUF_readStats_wksp(BYTE * huffWeight,size_t hwSize,U32 * rankStats,U32 * nbSymbolsPtr,U32 * tableLogPtr,const void * src,size_t srcSize,void * workSpace,size_t wkspSize,int flags)328 size_t HUF_readStats_wksp(BYTE* huffWeight, size_t hwSize, U32* rankStats,
329                      U32* nbSymbolsPtr, U32* tableLogPtr,
330                      const void* src, size_t srcSize,
331                      void* workSpace, size_t wkspSize,
332                      int flags)
333 {
334 #if DYNAMIC_BMI2
335     if (flags & HUF_flags_bmi2) {
336         return HUF_readStats_body_bmi2(huffWeight, hwSize, rankStats, nbSymbolsPtr, tableLogPtr, src, srcSize, workSpace, wkspSize);
337     }
338 #endif
339     (void)flags;
340     return HUF_readStats_body_default(huffWeight, hwSize, rankStats, nbSymbolsPtr, tableLogPtr, src, srcSize, workSpace, wkspSize);
341 }
342