lib/zstd/huf_compress.c

   1 /*
   2  * Huffman encoder, part of New Generation Entropy library
   3  * Copyright (C) 2013-2016, Yann Collet.
   4  *
   5  * BSD 2-Clause License (http://www.opensource.org/licenses/bsd-license.php)
   6  *
   7  * Redistribution and use in source and binary forms, with or without
   8  * modification, are permitted provided that the following conditions are
   9  * met:
  10  *
  11  *   * Redistributions of source code must retain the above copyright
  12  * notice, this list of conditions and the following disclaimer.
  13  *   * Redistributions in binary form must reproduce the above
  14  * copyright notice, this list of conditions and the following disclaimer
  15  * in the documentation and/or other materials provided with the
  16  * distribution.
  17  *
  18  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  19  * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  20  * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  21  * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  22  * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  23  * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  24  * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  25  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  26  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  27  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  28  * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  29  *
  30  * This program is free software; you can redistribute it and/or modify it under
  31  * the terms of the GNU General Public License version 2 as published by the
  32  * Free Software Foundation. This program is dual-licensed; you may select
  33  * either version 2 of the GNU General Public License ("GPL") or BSD license
  34  * ("BSD").
  35  *
  36  * You can contact the author at :
  37  * - Source repository : https://github.com/Cyan4973/FiniteStateEntropy
  38  */
  39
  40 /* **************************************************************
  41 *  Includes
  42 ****************************************************************/
  43 #include "bitstream.h"
  44 #include "fse.h" /* header compression */
  45 #include "huf.h"
  46 #include <linux/kernel.h>
  47 #include <linux/string.h> /* memcpy, memset */
  48
  49 /* **************************************************************
  50 *  Error Management
  51 ****************************************************************/
  52 #define HUF_STATIC_ASSERT(c)                                   \
  53         {                                                      \
  54                 enum { HUF_static_assert = 1 / (int)(!!(c)) }; \
  55         } /* use only *after* variable declarations */
  56 #define CHECK_V_F(e, f)     \
  57         size_t const e = f; \
  58         if (ERR_isError(e)) \
  59         return f
  60 #define CHECK_F(f)                        \
  61         {                                 \
  62                 CHECK_V_F(_var_err__, f); \
  63         }
  64
  65 /* **************************************************************
  66 *  Utils
  67 ****************************************************************/
  68 unsigned HUF_optimalTableLog(unsigned maxTableLog, size_t srcSize, unsigned maxSymbolValue)
  69 {
  70         return FSE_optimalTableLog_internal(maxTableLog, srcSize, maxSymbolValue, 1);
  71 }
  72
  73 /* *******************************************************
  74 *  HUF : Huffman block compression
  75 *********************************************************/
  76 /* HUF_compressWeights() :
  77  * Same as FSE_compress(), but dedicated to huff0's weights compression.
  78  * The use case needs much less stack memory.
  79  * Note : all elements within weightTable are supposed to be <= HUF_TABLELOG_MAX.
  80  */
  81 #define MAX_FSE_TABLELOG_FOR_HUFF_HEADER 6
  82 size_t HUF_compressWeights_wksp(void *dst, size_t dstSize, const void *weightTable, size_t wtSize, void *workspace, size_t workspaceSize)
  83 {
  84         BYTE *const ostart = (BYTE *)dst;
  85         BYTE *op = ostart;
  86         BYTE *const oend = ostart + dstSize;
  87
  88         U32 maxSymbolValue = HUF_TABLELOG_MAX;
  89         U32 tableLog = MAX_FSE_TABLELOG_FOR_HUFF_HEADER;
  90
  91         FSE_CTable *CTable;
  92         U32 *count;
  93         S16 *norm;
  94         size_t spaceUsed32 = 0;
  95
  96         HUF_STATIC_ASSERT(sizeof(FSE_CTable) == sizeof(U32));
  97
  98         CTable = (FSE_CTable *)((U32 *)workspace + spaceUsed32);
  99         spaceUsed32 += FSE_CTABLE_SIZE_U32(MAX_FSE_TABLELOG_FOR_HUFF_HEADER, HUF_TABLELOG_MAX);
 100         count = (U32 *)workspace + spaceUsed32;
 101         spaceUsed32 += HUF_TABLELOG_MAX + 1;
 102         norm = (S16 *)((U32 *)workspace + spaceUsed32);
 103         spaceUsed32 += ALIGN(sizeof(S16) * (HUF_TABLELOG_MAX + 1), sizeof(U32)) >> 2;
 104
 105         if ((spaceUsed32 << 2) > workspaceSize)
 106                 return ERROR(tableLog_tooLarge);
 107         workspace = (U32 *)workspace + spaceUsed32;
 108         workspaceSize -= (spaceUsed32 << 2);
 109
 110         /* init conditions */
 111         if (wtSize <= 1)
 112                 return 0; /* Not compressible */
 113
 114         /* Scan input and build symbol stats */
 115         {
 116                 CHECK_V_F(maxCount, FSE_count_simple(count, &maxSymbolValue, weightTable, wtSize));
 117                 if (maxCount == wtSize)
 118                         return 1; /* only a single symbol in src : rle */
 119                 if (maxCount == 1)
 120                         return 0; /* each symbol present maximum once => not compressible */
 121         }
 122
 123         tableLog = FSE_optimalTableLog(tableLog, wtSize, maxSymbolValue);
 124         CHECK_F(FSE_normalizeCount(norm, tableLog, count, wtSize, maxSymbolValue));
 125
 126         /* Write table description header */
 127         {
 128                 CHECK_V_F(hSize, FSE_writeNCount(op, oend - op, norm, maxSymbolValue, tableLog));
 129                 op += hSize;
 130         }
 131
 132         /* Compress */
 133         CHECK_F(FSE_buildCTable_wksp(CTable, norm, maxSymbolValue, tableLog, workspace, workspaceSize));
 134         {
 135                 CHECK_V_F(cSize, FSE_compress_usingCTable(op, oend - op, weightTable, wtSize, CTable));
 136                 if (cSize == 0)
 137                         return 0; /* not enough space for compressed data */
 138                 op += cSize;
 139         }
 140
 141         return op - ostart;
 142 }
 143
 144 struct HUF_CElt_s {
 145         U16 val;
 146         BYTE nbBits;
 147 }; /* typedef'd to HUF_CElt within "huf.h" */
 148
 149 /*! HUF_writeCTable_wksp() :
 150         `CTable` : Huffman tree to save, using huf representation.
 151         @return : size of saved CTable */
 152 size_t HUF_writeCTable_wksp(void *dst, size_t maxDstSize, const HUF_CElt *CTable, U32 maxSymbolValue, U32 huffLog, void *workspace, size_t workspaceSize)
 153 {
 154         BYTE *op = (BYTE *)dst;
 155         U32 n;
 156
 157         BYTE *bitsToWeight;
 158         BYTE *huffWeight;
 159         size_t spaceUsed32 = 0;
 160
 161         bitsToWeight = (BYTE *)((U32 *)workspace + spaceUsed32);
 162         spaceUsed32 += ALIGN(HUF_TABLELOG_MAX + 1, sizeof(U32)) >> 2;
 163         huffWeight = (BYTE *)((U32 *)workspace + spaceUsed32);
 164         spaceUsed32 += ALIGN(HUF_SYMBOLVALUE_MAX, sizeof(U32)) >> 2;
 165
 166         if ((spaceUsed32 << 2) > workspaceSize)
 167                 return ERROR(tableLog_tooLarge);
 168         workspace = (U32 *)workspace + spaceUsed32;
 169         workspaceSize -= (spaceUsed32 << 2);
 170
 171         /* check conditions */
 172         if (maxSymbolValue > HUF_SYMBOLVALUE_MAX)
 173                 return ERROR(maxSymbolValue_tooLarge);
 174
 175         /* convert to weight */
 176         bitsToWeight[0] = 0;
 177         for (n = 1; n < huffLog + 1; n++)
 178                 bitsToWeight[n] = (BYTE)(huffLog + 1 - n);
 179         for (n = 0; n < maxSymbolValue; n++)
 180                 huffWeight[n] = bitsToWeight[CTable[n].nbBits];
 181
 182         /* attempt weights compression by FSE */
 183         {
 184                 CHECK_V_F(hSize, HUF_compressWeights_wksp(op + 1, maxDstSize - 1, huffWeight, maxSymbolValue, workspace, workspaceSize));
 185                 if ((hSize > 1) & (hSize < maxSymbolValue / 2)) { /* FSE compressed */
 186                         op[0] = (BYTE)hSize;
 187                         return hSize + 1;
 188                 }
 189         }
 190
 191         /* write raw values as 4-bits (max : 15) */
 192         if (maxSymbolValue > (256 - 128))
 193                 return ERROR(GENERIC); /* should not happen : likely means source cannot be compressed */
 194         if (((maxSymbolValue + 1) / 2) + 1 > maxDstSize)
 195                 return ERROR(dstSize_tooSmall); /* not enough space within dst buffer */
 196         op[0] = (BYTE)(128 /*special case*/ + (maxSymbolValue - 1));
 197         huffWeight[maxSymbolValue] = 0; /* to be sure it doesn't cause msan issue in final combination */
 198         for (n = 0; n < maxSymbolValue; n += 2)
 199                 op[(n / 2) + 1] = (BYTE)((huffWeight[n] << 4) + huffWeight[n + 1]);
 200         return ((maxSymbolValue + 1) / 2) + 1;
 201 }
 202
 203 size_t HUF_readCTable_wksp(HUF_CElt *CTable, U32 maxSymbolValue, const void *src, size_t srcSize, void *workspace, size_t workspaceSize)
 204 {
 205         U32 *rankVal;
 206         BYTE *huffWeight;
 207         U32 tableLog = 0;
 208         U32 nbSymbols = 0;
 209         size_t readSize;
 210         size_t spaceUsed32 = 0;
 211
 212         rankVal = (U32 *)workspace + spaceUsed32;
 213         spaceUsed32 += HUF_TABLELOG_ABSOLUTEMAX + 1;
 214         huffWeight = (BYTE *)((U32 *)workspace + spaceUsed32);
 215         spaceUsed32 += ALIGN(HUF_SYMBOLVALUE_MAX + 1, sizeof(U32)) >> 2;
 216
 217         if ((spaceUsed32 << 2) > workspaceSize)
 218                 return ERROR(tableLog_tooLarge);
 219         workspace = (U32 *)workspace + spaceUsed32;
 220         workspaceSize -= (spaceUsed32 << 2);
 221
 222         /* get symbol weights */
 223         readSize = HUF_readStats_wksp(huffWeight, HUF_SYMBOLVALUE_MAX + 1, rankVal, &nbSymbols, &tableLog, src, srcSize, workspace, workspaceSize);
 224         if (ERR_isError(readSize))
 225                 return readSize;
 226
 227         /* check result */
 228         if (tableLog > HUF_TABLELOG_MAX)
 229                 return ERROR(tableLog_tooLarge);
 230         if (nbSymbols > maxSymbolValue + 1)
 231                 return ERROR(maxSymbolValue_tooSmall);
 232
 233         /* Prepare base value per rank */
 234         {
 235                 U32 n, nextRankStart = 0;
 236                 for (n = 1; n <= tableLog; n++) {
 237                         U32 curr = nextRankStart;
 238                         nextRankStart += (rankVal[n] << (n - 1));
 239                         rankVal[n] = curr;
 240                 }
 241         }
 242
 243         /* fill nbBits */
 244         {
 245                 U32 n;
 246                 for (n = 0; n < nbSymbols; n++) {
 247                         const U32 w = huffWeight[n];
 248                         CTable[n].nbBits = (BYTE)(tableLog + 1 - w);
 249                 }
 250         }
 251
 252         /* fill val */
 253         {
 254                 U16 nbPerRank[HUF_TABLELOG_MAX + 2] = {0}; /* support w=0=>n=tableLog+1 */
 255                 U16 valPerRank[HUF_TABLELOG_MAX + 2] = {0};
 256                 {
 257                         U32 n;
 258                         for (n = 0; n < nbSymbols; n++)
 259                                 nbPerRank[CTable[n].nbBits]++;
 260                 }
 261                 /* determine stating value per rank */
 262                 valPerRank[tableLog + 1] = 0; /* for w==0 */
 263                 {
 264                         U16 min = 0;
 265                         U32 n;
 266                         for (n = tableLog; n > 0; n--) { /* start at n=tablelog <-> w=1 */
 267                                 valPerRank[n] = min;     /* get starting value within each rank */
 268                                 min += nbPerRank[n];
 269                                 min >>= 1;
 270                         }
 271                 }
 272                 /* assign value within rank, symbol order */
 273                 {
 274                         U32 n;
 275                         for (n = 0; n <= maxSymbolValue; n++)
 276                                 CTable[n].val = valPerRank[CTable[n].nbBits]++;
 277                 }
 278         }
 279
 280         return readSize;
 281 }
 282
 283 typedef struct nodeElt_s {
 284         U32 count;
 285         U16 parent;
 286         BYTE byte;
 287         BYTE nbBits;
 288 } nodeElt;
 289
 290 static U32 HUF_setMaxHeight(nodeElt *huffNode, U32 lastNonNull, U32 maxNbBits)
 291 {
 292         const U32 largestBits = huffNode[lastNonNull].nbBits;
 293         if (largestBits <= maxNbBits)
 294                 return largestBits; /* early exit : no elt > maxNbBits */
 295
 296         /* there are several too large elements (at least >= 2) */
 297         {
 298                 int totalCost = 0;
 299                 const U32 baseCost = 1 << (largestBits - maxNbBits);
 300                 U32 n = lastNonNull;
 301
 302                 while (huffNode[n].nbBits > maxNbBits) {
 303                         totalCost += baseCost - (1 << (largestBits - huffNode[n].nbBits));
 304                         huffNode[n].nbBits = (BYTE)maxNbBits;
 305                         n--;
 306                 } /* n stops at huffNode[n].nbBits <= maxNbBits */
 307                 while (huffNode[n].nbBits == maxNbBits)
 308                         n--; /* n end at index of smallest symbol using < maxNbBits */
 309
 310                 /* renorm totalCost */
 311                 totalCost >>= (largestBits - maxNbBits); /* note : totalCost is necessarily a multiple of baseCost */
 312
 313                 /* repay normalized cost */
 314                 {
 315                         U32 const noSymbol = 0xF0F0F0F0;
 316                         U32 rankLast[HUF_TABLELOG_MAX + 2];
 317                         int pos;
 318
 319                         /* Get pos of last (smallest) symbol per rank */
 320                         memset(rankLast, 0xF0, sizeof(rankLast));
 321                         {
 322                                 U32 currNbBits = maxNbBits;
 323                                 for (pos = n; pos >= 0; pos--) {
 324                                         if (huffNode[pos].nbBits >= currNbBits)
 325                                                 continue;
 326                                         currNbBits = huffNode[pos].nbBits; /* < maxNbBits */
 327                                         rankLast[maxNbBits - currNbBits] = pos;
 328                                 }
 329                         }
 330
 331                         while (totalCost > 0) {
 332                                 U32 nBitsToDecrease = BIT_highbit32(totalCost) + 1;
 333                                 for (; nBitsToDecrease > 1; nBitsToDecrease--) {
 334                                         U32 highPos = rankLast[nBitsToDecrease];
 335                                         U32 lowPos = rankLast[nBitsToDecrease - 1];
 336                                         if (highPos == noSymbol)
 337                                                 continue;
 338                                         if (lowPos == noSymbol)
 339                                                 break;
 340                                         {
 341                                                 U32 const highTotal = huffNode[highPos].count;
 342                                                 U32 const lowTotal = 2 * huffNode[lowPos].count;
 343                                                 if (highTotal <= lowTotal)
 344                                                         break;
 345                                         }
 346                                 }
 347                                 /* only triggered when no more rank 1 symbol left => find closest one (note : there is necessarily at least one !) */
 348                                 /* HUF_MAX_TABLELOG test just to please gcc 5+; but it should not be necessary */
 349                                 while ((nBitsToDecrease <= HUF_TABLELOG_MAX) && (rankLast[nBitsToDecrease] == noSymbol))
 350                                         nBitsToDecrease++;
 351                                 totalCost -= 1 << (nBitsToDecrease - 1);
 352                                 if (rankLast[nBitsToDecrease - 1] == noSymbol)
 353                                         rankLast[nBitsToDecrease - 1] = rankLast[nBitsToDecrease]; /* this rank is no longer empty */
 354                                 huffNode[rankLast[nBitsToDecrease]].nbBits++;
 355                                 if (rankLast[nBitsToDecrease] == 0) /* special case, reached largest symbol */
 356                                         rankLast[nBitsToDecrease] = noSymbol;
 357                                 else {
 358                                         rankLast[nBitsToDecrease]--;
 359                                         if (huffNode[rankLast[nBitsToDecrease]].nbBits != maxNbBits - nBitsToDecrease)
 360                                                 rankLast[nBitsToDecrease] = noSymbol; /* this rank is now empty */
 361                                 }
 362                         } /* while (totalCost > 0) */
 363
 364                         while (totalCost < 0) {                /* Sometimes, cost correction overshoot */
 365                                 if (rankLast[1] == noSymbol) { /* special case : no rank 1 symbol (using maxNbBits-1); let's create one from largest rank 0
 366                                                                   (using maxNbBits) */
 367                                         while (huffNode[n].nbBits == maxNbBits)
 368                                                 n--;
 369                                         huffNode[n + 1].nbBits--;
 370                                         rankLast[1] = n + 1;
 371                                         totalCost++;
 372                                         continue;
 373                                 }
 374                                 huffNode[rankLast[1] + 1].nbBits--;
 375                                 rankLast[1]++;
 376                                 totalCost++;
 377                         }
 378                 }
 379         } /* there are several too large elements (at least >= 2) */
 380
 381         return maxNbBits;
 382 }
 383
 384 typedef struct {
 385         U32 base;
 386         U32 curr;
 387 } rankPos;
 388
 389 static void HUF_sort(nodeElt *huffNode, const U32 *count, U32 maxSymbolValue)
 390 {
 391         rankPos rank[32];
 392         U32 n;
 393
 394         memset(rank, 0, sizeof(rank));
 395         for (n = 0; n <= maxSymbolValue; n++) {
 396                 U32 r = BIT_highbit32(count[n] + 1);
 397                 rank[r].base++;
 398         }
 399         for (n = 30; n > 0; n--)
 400                 rank[n - 1].base += rank[n].base;
 401         for (n = 0; n < 32; n++)
 402                 rank[n].curr = rank[n].base;
 403         for (n = 0; n <= maxSymbolValue; n++) {
 404                 U32 const c = count[n];
 405                 U32 const r = BIT_highbit32(c + 1) + 1;
 406                 U32 pos = rank[r].curr++;
 407                 while ((pos > rank[r].base) && (c > huffNode[pos - 1].count))
 408                         huffNode[pos] = huffNode[pos - 1], pos--;
 409                 huffNode[pos].count = c;
 410                 huffNode[pos].byte = (BYTE)n;
 411         }
 412 }
 413
 414 /** HUF_buildCTable_wksp() :
 415  *  Same as HUF_buildCTable(), but using externally allocated scratch buffer.
 416  *  `workSpace` must be aligned on 4-bytes boundaries, and be at least as large as a table of 1024 unsigned.
 417  */
 418 #define STARTNODE (HUF_SYMBOLVALUE_MAX + 1)
 419 typedef nodeElt huffNodeTable[2 * HUF_SYMBOLVALUE_MAX + 1 + 1];
 420 size_t HUF_buildCTable_wksp(HUF_CElt *tree, const U32 *count, U32 maxSymbolValue, U32 maxNbBits, void *workSpace, size_t wkspSize)
 421 {
 422         nodeElt *const huffNode0 = (nodeElt *)workSpace;
 423         nodeElt *const huffNode = huffNode0 + 1;
 424         U32 n, nonNullRank;
 425         int lowS, lowN;
 426         U16 nodeNb = STARTNODE;
 427         U32 nodeRoot;
 428
 429         /* safety checks */
 430         if (wkspSize < sizeof(huffNodeTable))
 431                 return ERROR(GENERIC); /* workSpace is not large enough */
 432         if (maxNbBits == 0)
 433                 maxNbBits = HUF_TABLELOG_DEFAULT;
 434         if (maxSymbolValue > HUF_SYMBOLVALUE_MAX)
 435                 return ERROR(GENERIC);
 436         memset(huffNode0, 0, sizeof(huffNodeTable));
 437
 438         /* sort, decreasing order */
 439         HUF_sort(huffNode, count, maxSymbolValue);
 440
 441         /* init for parents */
 442         nonNullRank = maxSymbolValue;
 443         while (huffNode[nonNullRank].count == 0)
 444                 nonNullRank--;
 445         lowS = nonNullRank;
 446         nodeRoot = nodeNb + lowS - 1;
 447         lowN = nodeNb;
 448         huffNode[nodeNb].count = huffNode[lowS].count + huffNode[lowS - 1].count;
 449         huffNode[lowS].parent = huffNode[lowS - 1].parent = nodeNb;
 450         nodeNb++;
 451         lowS -= 2;
 452         for (n = nodeNb; n <= nodeRoot; n++)
 453                 huffNode[n].count = (U32)(1U << 30);
 454         huffNode0[0].count = (U32)(1U << 31); /* fake entry, strong barrier */
 455
 456         /* create parents */
 457         while (nodeNb <= nodeRoot) {
 458                 U32 n1 = (huffNode[lowS].count < huffNode[lowN].count) ? lowS-- : lowN++;
 459                 U32 n2 = (huffNode[lowS].count < huffNode[lowN].count) ? lowS-- : lowN++;
 460                 huffNode[nodeNb].count = huffNode[n1].count + huffNode[n2].count;
 461                 huffNode[n1].parent = huffNode[n2].parent = nodeNb;
 462                 nodeNb++;
 463         }
 464
 465         /* distribute weights (unlimited tree height) */
 466         huffNode[nodeRoot].nbBits = 0;
 467         for (n = nodeRoot - 1; n >= STARTNODE; n--)
 468                 huffNode[n].nbBits = huffNode[huffNode[n].parent].nbBits + 1;
 469         for (n = 0; n <= nonNullRank; n++)
 470                 huffNode[n].nbBits = huffNode[huffNode[n].parent].nbBits + 1;
 471
 472         /* enforce maxTableLog */
 473         maxNbBits = HUF_setMaxHeight(huffNode, nonNullRank, maxNbBits);
 474
 475         /* fill result into tree (val, nbBits) */
 476         {
 477                 U16 nbPerRank[HUF_TABLELOG_MAX + 1] = {0};
 478                 U16 valPerRank[HUF_TABLELOG_MAX + 1] = {0};
 479                 if (maxNbBits > HUF_TABLELOG_MAX)
 480                         return ERROR(GENERIC); /* check fit into table */
 481                 for (n = 0; n <= nonNullRank; n++)
 482                         nbPerRank[huffNode[n].nbBits]++;
 483                 /* determine stating value per rank */
 484                 {
 485                         U16 min = 0;
 486                         for (n = maxNbBits; n > 0; n--) {
 487                                 valPerRank[n] = min; /* get starting value within each rank */
 488                                 min += nbPerRank[n];
 489                                 min >>= 1;
 490                         }
 491                 }
 492                 for (n = 0; n <= maxSymbolValue; n++)
 493                         tree[huffNode[n].byte].nbBits = huffNode[n].nbBits; /* push nbBits per symbol, symbol order */
 494                 for (n = 0; n <= maxSymbolValue; n++)
 495                         tree[n].val = valPerRank[tree[n].nbBits]++; /* assign value within rank, symbol order */
 496         }
 497
 498         return maxNbBits;
 499 }
 500
 501 static size_t HUF_estimateCompressedSize(HUF_CElt *CTable, const unsigned *count, unsigned maxSymbolValue)
 502 {
 503         size_t nbBits = 0;
 504         int s;
 505         for (s = 0; s <= (int)maxSymbolValue; ++s) {
 506                 nbBits += CTable[s].nbBits * count[s];
 507         }
 508         return nbBits >> 3;
 509 }
 510
 511 static int HUF_validateCTable(const HUF_CElt *CTable, const unsigned *count, unsigned maxSymbolValue)
 512 {
 513         int bad = 0;
 514         int s;
 515         for (s = 0; s <= (int)maxSymbolValue; ++s) {
 516                 bad |= (count[s] != 0) & (CTable[s].nbBits == 0);
 517         }
 518         return !bad;
 519 }
 520
 521 static void HUF_encodeSymbol(BIT_CStream_t *bitCPtr, U32 symbol, const HUF_CElt *CTable)
 522 {
 523         BIT_addBitsFast(bitCPtr, CTable[symbol].val, CTable[symbol].nbBits);
 524 }
 525
 526 size_t HUF_compressBound(size_t size) { return HUF_COMPRESSBOUND(size); }
 527
 528 #define HUF_FLUSHBITS(s)  BIT_flushBits(s)
 529
 530 #define HUF_FLUSHBITS_1(stream)                                            \
 531         if (sizeof((stream)->bitContainer) * 8 < HUF_TABLELOG_MAX * 2 + 7) \
 532         HUF_FLUSHBITS(stream)
 533
 534 #define HUF_FLUSHBITS_2(stream)                                            \
 535         if (sizeof((stream)->bitContainer) * 8 < HUF_TABLELOG_MAX * 4 + 7) \
 536         HUF_FLUSHBITS(stream)
 537
 538 size_t HUF_compress1X_usingCTable(void *dst, size_t dstSize, const void *src, size_t srcSize, const HUF_CElt *CTable)
 539 {
 540         const BYTE *ip = (const BYTE *)src;
 541         BYTE *const ostart = (BYTE *)dst;
 542         BYTE *const oend = ostart + dstSize;
 543         BYTE *op = ostart;
 544         size_t n;
 545         BIT_CStream_t bitC;
 546
 547         /* init */
 548         if (dstSize < 8)
 549                 return 0; /* not enough space to compress */
 550         {
 551                 size_t const initErr = BIT_initCStream(&bitC, op, oend - op);
 552                 if (HUF_isError(initErr))
 553                         return 0;
 554         }
 555
 556         n = srcSize & ~3; /* join to mod 4 */
 557         switch (srcSize & 3) {
 558         case 3: HUF_encodeSymbol(&bitC, ip[n + 2], CTable); HUF_FLUSHBITS_2(&bitC);
 559                 fallthrough;
 560         case 2: HUF_encodeSymbol(&bitC, ip[n + 1], CTable); HUF_FLUSHBITS_1(&bitC);
 561                 fallthrough;
 562         case 1: HUF_encodeSymbol(&bitC, ip[n + 0], CTable); HUF_FLUSHBITS(&bitC);
 563                 fallthrough;
 564         case 0:
 565         default:;
 566         }
 567
 568         for (; n > 0; n -= 4) { /* note : n&3==0 at this stage */
 569                 HUF_encodeSymbol(&bitC, ip[n - 1], CTable);
 570                 HUF_FLUSHBITS_1(&bitC);
 571                 HUF_encodeSymbol(&bitC, ip[n - 2], CTable);
 572                 HUF_FLUSHBITS_2(&bitC);
 573                 HUF_encodeSymbol(&bitC, ip[n - 3], CTable);
 574                 HUF_FLUSHBITS_1(&bitC);
 575                 HUF_encodeSymbol(&bitC, ip[n - 4], CTable);
 576                 HUF_FLUSHBITS(&bitC);
 577         }
 578
 579         return BIT_closeCStream(&bitC);
 580 }
 581
 582 size_t HUF_compress4X_usingCTable(void *dst, size_t dstSize, const void *src, size_t srcSize, const HUF_CElt *CTable)
 583 {
 584         size_t const segmentSize = (srcSize + 3) / 4; /* first 3 segments */
 585         const BYTE *ip = (const BYTE *)src;
 586         const BYTE *const iend = ip + srcSize;
 587         BYTE *const ostart = (BYTE *)dst;
 588         BYTE *const oend = ostart + dstSize;
 589         BYTE *op = ostart;
 590
 591         if (dstSize < 6 + 1 + 1 + 1 + 8)
 592                 return 0; /* minimum space to compress successfully */
 593         if (srcSize < 12)
 594                 return 0; /* no saving possible : too small input */
 595         op += 6;          /* jumpTable */
 596
 597         {
 598                 CHECK_V_F(cSize, HUF_compress1X_usingCTable(op, oend - op, ip, segmentSize, CTable));
 599                 if (cSize == 0)
 600                         return 0;
 601                 ZSTD_writeLE16(ostart, (U16)cSize);
 602                 op += cSize;
 603         }
 604
 605         ip += segmentSize;
 606         {
 607                 CHECK_V_F(cSize, HUF_compress1X_usingCTable(op, oend - op, ip, segmentSize, CTable));
 608                 if (cSize == 0)
 609                         return 0;
 610                 ZSTD_writeLE16(ostart + 2, (U16)cSize);
 611                 op += cSize;
 612         }
 613
 614         ip += segmentSize;
 615         {
 616                 CHECK_V_F(cSize, HUF_compress1X_usingCTable(op, oend - op, ip, segmentSize, CTable));
 617                 if (cSize == 0)
 618                         return 0;
 619                 ZSTD_writeLE16(ostart + 4, (U16)cSize);
 620                 op += cSize;
 621         }
 622
 623         ip += segmentSize;
 624         {
 625                 CHECK_V_F(cSize, HUF_compress1X_usingCTable(op, oend - op, ip, iend - ip, CTable));
 626                 if (cSize == 0)
 627                         return 0;
 628                 op += cSize;
 629         }
 630
 631         return op - ostart;
 632 }
 633
 634 static size_t HUF_compressCTable_internal(BYTE *const ostart, BYTE *op, BYTE *const oend, const void *src, size_t srcSize, unsigned singleStream,
 635                                           const HUF_CElt *CTable)
 636 {
 637         size_t const cSize =
 638             singleStream ? HUF_compress1X_usingCTable(op, oend - op, src, srcSize, CTable) : HUF_compress4X_usingCTable(op, oend - op, src, srcSize, CTable);
 639         if (HUF_isError(cSize)) {
 640                 return cSize;
 641         }
 642         if (cSize == 0) {
 643                 return 0;
 644         } /* uncompressible */
 645         op += cSize;
 646         /* check compressibility */
 647         if ((size_t)(op - ostart) >= srcSize - 1) {
 648                 return 0;
 649         }
 650         return op - ostart;
 651 }
 652
 653 /* `workSpace` must a table of at least 1024 unsigned */
 654 static size_t HUF_compress_internal(void *dst, size_t dstSize, const void *src, size_t srcSize, unsigned maxSymbolValue, unsigned huffLog,
 655                                     unsigned singleStream, void *workSpace, size_t wkspSize, HUF_CElt *oldHufTable, HUF_repeat *repeat, int preferRepeat)
 656 {
 657         BYTE *const ostart = (BYTE *)dst;
 658         BYTE *const oend = ostart + dstSize;
 659         BYTE *op = ostart;
 660
 661         U32 *count;
 662         size_t const countSize = sizeof(U32) * (HUF_SYMBOLVALUE_MAX + 1);
 663         HUF_CElt *CTable;
 664         size_t const CTableSize = sizeof(HUF_CElt) * (HUF_SYMBOLVALUE_MAX + 1);
 665
 666         /* checks & inits */
 667         if (wkspSize < sizeof(huffNodeTable) + countSize + CTableSize)
 668                 return ERROR(GENERIC);
 669         if (!srcSize)
 670                 return 0; /* Uncompressed (note : 1 means rle, so first byte must be correct) */
 671         if (!dstSize)
 672                 return 0; /* cannot fit within dst budget */
 673         if (srcSize > HUF_BLOCKSIZE_MAX)
 674                 return ERROR(srcSize_wrong); /* curr block size limit */
 675         if (huffLog > HUF_TABLELOG_MAX)
 676                 return ERROR(tableLog_tooLarge);
 677         if (!maxSymbolValue)
 678                 maxSymbolValue = HUF_SYMBOLVALUE_MAX;
 679         if (!huffLog)
 680                 huffLog = HUF_TABLELOG_DEFAULT;
 681
 682         count = (U32 *)workSpace;
 683         workSpace = (BYTE *)workSpace + countSize;
 684         wkspSize -= countSize;
 685         CTable = (HUF_CElt *)workSpace;
 686         workSpace = (BYTE *)workSpace + CTableSize;
 687         wkspSize -= CTableSize;
 688
 689         /* Heuristic : If we don't need to check the validity of the old table use the old table for small inputs */
 690         if (preferRepeat && repeat && *repeat == HUF_repeat_valid) {
 691                 return HUF_compressCTable_internal(ostart, op, oend, src, srcSize, singleStream, oldHufTable);
 692         }
 693
 694         /* Scan input and build symbol stats */
 695         {
 696                 CHECK_V_F(largest, FSE_count_wksp(count, &maxSymbolValue, (const BYTE *)src, srcSize, (U32 *)workSpace));
 697                 if (largest == srcSize) {
 698                         *ostart = ((const BYTE *)src)[0];
 699                         return 1;
 700                 } /* single symbol, rle */
 701                 if (largest <= (srcSize >> 7) + 1)
 702                         return 0; /* Fast heuristic : not compressible enough */
 703         }
 704
 705         /* Check validity of previous table */
 706         if (repeat && *repeat == HUF_repeat_check && !HUF_validateCTable(oldHufTable, count, maxSymbolValue)) {
 707                 *repeat = HUF_repeat_none;
 708         }
 709         /* Heuristic : use existing table for small inputs */
 710         if (preferRepeat && repeat && *repeat != HUF_repeat_none) {
 711                 return HUF_compressCTable_internal(ostart, op, oend, src, srcSize, singleStream, oldHufTable);
 712         }
 713
 714         /* Build Huffman Tree */
 715         huffLog = HUF_optimalTableLog(huffLog, srcSize, maxSymbolValue);
 716         {
 717                 CHECK_V_F(maxBits, HUF_buildCTable_wksp(CTable, count, maxSymbolValue, huffLog, workSpace, wkspSize));
 718                 huffLog = (U32)maxBits;
 719                 /* Zero the unused symbols so we can check it for validity */
 720                 memset(CTable + maxSymbolValue + 1, 0, CTableSize - (maxSymbolValue + 1) * sizeof(HUF_CElt));
 721         }
 722
 723         /* Write table description header */
 724         {
 725                 CHECK_V_F(hSize, HUF_writeCTable_wksp(op, dstSize, CTable, maxSymbolValue, huffLog, workSpace, wkspSize));
 726                 /* Check if using the previous table will be beneficial */
 727                 if (repeat && *repeat != HUF_repeat_none) {
 728                         size_t const oldSize = HUF_estimateCompressedSize(oldHufTable, count, maxSymbolValue);
 729                         size_t const newSize = HUF_estimateCompressedSize(CTable, count, maxSymbolValue);
 730                         if (oldSize <= hSize + newSize || hSize + 12 >= srcSize) {
 731                                 return HUF_compressCTable_internal(ostart, op, oend, src, srcSize, singleStream, oldHufTable);
 732                         }
 733                 }
 734                 /* Use the new table */
 735                 if (hSize + 12ul >= srcSize) {
 736                         return 0;
 737                 }
 738                 op += hSize;
 739                 if (repeat) {
 740                         *repeat = HUF_repeat_none;
 741                 }
 742                 if (oldHufTable) {
 743                         memcpy(oldHufTable, CTable, CTableSize);
 744                 } /* Save the new table */
 745         }
 746         return HUF_compressCTable_internal(ostart, op, oend, src, srcSize, singleStream, CTable);
 747 }
 748
 749 size_t HUF_compress1X_wksp(void *dst, size_t dstSize, const void *src, size_t srcSize, unsigned maxSymbolValue, unsigned huffLog, void *workSpace,
 750                            size_t wkspSize)
 751 {
 752         return HUF_compress_internal(dst, dstSize, src, srcSize, maxSymbolValue, huffLog, 1 /* single stream */, workSpace, wkspSize, NULL, NULL, 0);
 753 }
 754
 755 size_t HUF_compress1X_repeat(void *dst, size_t dstSize, const void *src, size_t srcSize, unsigned maxSymbolValue, unsigned huffLog, void *workSpace,
 756                              size_t wkspSize, HUF_CElt *hufTable, HUF_repeat *repeat, int preferRepeat)
 757 {
 758         return HUF_compress_internal(dst, dstSize, src, srcSize, maxSymbolValue, huffLog, 1 /* single stream */, workSpace, wkspSize, hufTable, repeat,
 759                                      preferRepeat);
 760 }
 761
 762 size_t HUF_compress4X_wksp(void *dst, size_t dstSize, const void *src, size_t srcSize, unsigned maxSymbolValue, unsigned huffLog, void *workSpace,
 763                            size_t wkspSize)
 764 {
 765         return HUF_compress_internal(dst, dstSize, src, srcSize, maxSymbolValue, huffLog, 0 /* 4 streams */, workSpace, wkspSize, NULL, NULL, 0);
 766 }
 767
 768 size_t HUF_compress4X_repeat(void *dst, size_t dstSize, const void *src, size_t srcSize, unsigned maxSymbolValue, unsigned huffLog, void *workSpace,
 769                              size_t wkspSize, HUF_CElt *hufTable, HUF_repeat *repeat, int preferRepeat)
 770 {
 771         return HUF_compress_internal(dst, dstSize, src, srcSize, maxSymbolValue, huffLog, 0 /* 4 streams */, workSpace, wkspSize, hufTable, repeat,
 772                                      preferRepeat);
 773 }