1 /**

     2  * Copyright (c) 2016-present, Yann Collet, Facebook, Inc.

     3  * All rights reserved.

     4  *

     5  * This source code is licensed under the BSD-style license found in the

     6  * LICENSE file in the root directory of this source tree. An additional grant

     7  * of patent rights can be found in the PATENTS file in the same directory.

     8  */

     9 

    10 

    11 

    12 /* *************************************

    13 *  Dependencies

    14 ***************************************/

    15 #include <stdlib.h>

    16 #include "error_private.h"

    17 #include "zstd_internal.h"  /* MIN, ZSTD_BLOCKHEADERSIZE, defaultCustomMem */

    18 #define ZBUFF_STATIC_LINKING_ONLY

    19 #include "zbuff.h"

    20 

    21 

    22 /* *************************************

    23 *  Constants

    24 ***************************************/

    25 static size_t const ZBUFF_endFrameSize = ZSTD_BLOCKHEADERSIZE;

    26 

    27 

    28 /*-***********************************************************

    29 *  Streaming compression

    30 *

    31 *  A ZBUFF_CCtx object is required to track streaming operation.

    32 *  Use ZBUFF_createCCtx() and ZBUFF_freeCCtx() to create/release resources.

    33 *  Use ZBUFF_compressInit() to start a new compression operation.

    34 *  ZBUFF_CCtx objects can be reused multiple times.

    35 *

    36 *  Use ZBUFF_compressContinue() repetitively to consume your input.

    37 *  *srcSizePtr and *dstCapacityPtr can be any size.

    38 *  The function will report how many bytes were read or written by modifying *srcSizePtr and *dstCapacityPtr.

    39 *  Note that it may not consume the entire input, in which case it's up to the caller to call again the function with remaining input.

    40 *  The content of dst will be overwritten (up to *dstCapacityPtr) at each function call, so save its content if it matters or change dst .

    41 *  @return : a hint to preferred nb of bytes to use as input for next function call (it's only a hint, to improve latency)

    42 *            or an error code, which can be tested using ZBUFF_isError().

    43 *

    44 *  ZBUFF_compressFlush() can be used to instruct ZBUFF to compress and output whatever remains within its buffer.

    45 *  Note that it will not output more than *dstCapacityPtr.

    46 *  Therefore, some content might still be left into its internal buffer if dst buffer is too small.

    47 *  @return : nb of bytes still present into internal buffer (0 if it's empty)

    48 *            or an error code, which can be tested using ZBUFF_isError().

    49 *

    50 *  ZBUFF_compressEnd() instructs to finish a frame.

    51 *  It will perform a flush and write frame epilogue.

    52 *  Similar to ZBUFF_compressFlush(), it may not be able to output the entire internal buffer content if *dstCapacityPtr is too small.

    53 *  @return : nb of bytes still present into internal buffer (0 if it's empty)

    54 *            or an error code, which can be tested using ZBUFF_isError().

    55 *

    56 *  Hint : recommended buffer sizes (not compulsory)

    57 *  input : ZSTD_BLOCKSIZE_MAX (128 KB), internal unit size, it improves latency to use this value.

    58 *  output : ZSTD_compressBound(ZSTD_BLOCKSIZE_MAX) + ZSTD_blockHeaderSize + ZBUFF_endFrameSize : ensures it's always possible to write/flush/end a full block at best speed.

    59 * ***********************************************************/

    60 

    61 typedef enum { ZBUFFcs_init, ZBUFFcs_load, ZBUFFcs_flush, ZBUFFcs_final } ZBUFF_cStage;

    62 

    63 /* *** Resources *** */

    64 struct ZBUFF_CCtx_s {

    65     ZSTD_CCtx* zc;

    66     char*  inBuff;

    67     size_t inBuffSize;

    68     size_t inToCompress;

    69     size_t inBuffPos;

    70     size_t inBuffTarget;

    71     size_t blockSize;

    72     char*  outBuff;

    73     size_t outBuffSize;

    74     size_t outBuffContentSize;

    75     size_t outBuffFlushedSize;

    76     ZBUFF_cStage stage;

    77     U32    checksum;

    78     U32    frameEnded;

    79     ZSTD_customMem customMem;

    80 };   /* typedef'd tp ZBUFF_CCtx within "zbuff.h" */

    81 

    82 ZBUFF_CCtx* ZBUFF_createCCtx(void)

    83 {

    84     return ZBUFF_createCCtx_advanced(defaultCustomMem);

    85 }

    86 

    87 ZBUFF_CCtx* ZBUFF_createCCtx_advanced(ZSTD_customMem customMem)

    88 {

    89     ZBUFF_CCtx* zbc;

    90 

    91     if (!customMem.customAlloc && !customMem.customFree)

    92         customMem = defaultCustomMem;

    93 

    94     if (!customMem.customAlloc || !customMem.customFree)

    95         return NULL;

    96 

    97     zbc = (ZBUFF_CCtx*)customMem.customAlloc(customMem.opaque, sizeof(ZBUFF_CCtx));

    98     if (zbc==NULL) return NULL;

    99     memset(zbc, 0, sizeof(ZBUFF_CCtx));

   100     memcpy(&zbc->customMem, &customMem, sizeof(ZSTD_customMem));

   101     zbc->zc = ZSTD_createCCtx_advanced(customMem);

   102     if (zbc->zc == NULL) { ZBUFF_freeCCtx(zbc); return NULL; }

   103     return zbc;

   104 }

   105 

   106 size_t ZBUFF_freeCCtx(ZBUFF_CCtx* zbc)

   107 {

   108     if (zbc==NULL) return 0;   /* support free on NULL */

   109     ZSTD_freeCCtx(zbc->zc);

   110     if (zbc->inBuff) zbc->customMem.customFree(zbc->customMem.opaque, zbc->inBuff);

   111     if (zbc->outBuff) zbc->customMem.customFree(zbc->customMem.opaque, zbc->outBuff);

   112     zbc->customMem.customFree(zbc->customMem.opaque, zbc);

   113     return 0;

   114 }

   115 

   116 

   117 /* ======   Initialization   ====== */

   118 

   119 size_t ZBUFF_compressInit_advanced(ZBUFF_CCtx* zbc,

   120                                    const void* dict, size_t dictSize,

   121                                    ZSTD_parameters params, unsigned long long pledgedSrcSize)

   122 {

   123     /* allocate buffers */

   124     {   size_t const neededInBuffSize = (size_t)1 << params.cParams.windowLog;

   125         if (zbc->inBuffSize < neededInBuffSize) {

   126             zbc->inBuffSize = neededInBuffSize;

   127             zbc->customMem.customFree(zbc->customMem.opaque, zbc->inBuff);   /* should not be necessary */

   128             zbc->inBuff = (char*)zbc->customMem.customAlloc(zbc->customMem.opaque, neededInBuffSize);

   129             if (zbc->inBuff == NULL) return ERROR(memory_allocation);

   130         }

   131         zbc->blockSize = MIN(ZSTD_BLOCKSIZE_ABSOLUTEMAX, neededInBuffSize);

   132     }

   133     if (zbc->outBuffSize < ZSTD_compressBound(zbc->blockSize)+1) {

   134         zbc->outBuffSize = ZSTD_compressBound(zbc->blockSize)+1;

   135         zbc->customMem.customFree(zbc->customMem.opaque, zbc->outBuff);   /* should not be necessary */

   136         zbc->outBuff = (char*)zbc->customMem.customAlloc(zbc->customMem.opaque, zbc->outBuffSize);

   137         if (zbc->outBuff == NULL) return ERROR(memory_allocation);

   138     }

   139 

   140     { size_t const errorCode = ZSTD_compressBegin_advanced(zbc->zc, dict, dictSize, params, pledgedSrcSize);

   141       if (ZSTD_isError(errorCode)) return errorCode; }

   142 

   143     zbc->inToCompress = 0;

   144     zbc->inBuffPos = 0;

   145     zbc->inBuffTarget = zbc->blockSize;

   146     zbc->outBuffContentSize = zbc->outBuffFlushedSize = 0;

   147     zbc->stage = ZBUFFcs_load;

   148     zbc->checksum = params.fParams.checksumFlag > 0;

   149     zbc->frameEnded = 0;

   150     return 0;   /* ready to go */

   151 }

   152 

   153 

   154 size_t ZBUFF_compressInitDictionary(ZBUFF_CCtx* zbc, const void* dict, size_t dictSize, int compressionLevel)

   155 {

   156     ZSTD_parameters const params = ZSTD_getParams(compressionLevel, 0, dictSize);

   157     return ZBUFF_compressInit_advanced(zbc, dict, dictSize, params, 0);

   158 }

   159 

   160 size_t ZBUFF_compressInit(ZBUFF_CCtx* zbc, int compressionLevel)

   161 {

   162     return ZBUFF_compressInitDictionary(zbc, NULL, 0, compressionLevel);

   163 }

   164 

   165 

   166 /* internal util function */

   167 MEM_STATIC size_t ZBUFF_limitCopy(void* dst, size_t dstCapacity, const void* src, size_t srcSize)

   168 {

   169     size_t const length = MIN(dstCapacity, srcSize);

   170     memcpy(dst, src, length);

   171     return length;

   172 }

   173 

   174 

   175 /* ======   Compression   ====== */

   176 

   177 typedef enum { zbf_gather, zbf_flush, zbf_end } ZBUFF_flush_e;

   178 

   179 static size_t ZBUFF_compressContinue_generic(ZBUFF_CCtx* zbc,

   180                               void* dst, size_t* dstCapacityPtr,

   181                         const void* src, size_t* srcSizePtr,

   182                               ZBUFF_flush_e const flush)

   183 {

   184     U32 someMoreWork = 1;

   185     const char* const istart = (const char*)src;

   186     const char* const iend = istart + *srcSizePtr;

   187     const char* ip = istart;

   188     char* const ostart = (char*)dst;

   189     char* const oend = ostart + *dstCapacityPtr;

   190     char* op = ostart;

   191 

   192     while (someMoreWork) {

   193         switch(zbc->stage)

   194         {

   195         case ZBUFFcs_init: return ERROR(init_missing);   /* call ZBUFF_compressInit() first ! */

   196 

   197         case ZBUFFcs_load:

   198             /* complete inBuffer */

   199             {   size_t const toLoad = zbc->inBuffTarget - zbc->inBuffPos;

   200                 size_t const loaded = ZBUFF_limitCopy(zbc->inBuff + zbc->inBuffPos, toLoad, ip, iend-ip);

   201                 zbc->inBuffPos += loaded;

   202                 ip += loaded;

   203                 if ( (zbc->inBuffPos==zbc->inToCompress) || (!flush && (toLoad != loaded)) ) {

   204                     someMoreWork = 0; break;  /* not enough input to get a full block : stop there, wait for more */

   205             }   }

   206             /* compress current block (note : this stage cannot be stopped in the middle) */

   207             {   void* cDst;

   208                 size_t cSize;

   209                 size_t const iSize = zbc->inBuffPos - zbc->inToCompress;

   210                 size_t oSize = oend-op;

   211                 if (oSize >= ZSTD_compressBound(iSize))

   212                     cDst = op;   /* compress directly into output buffer (avoid flush stage) */

   213                 else

   214                     cDst = zbc->outBuff, oSize = zbc->outBuffSize;

   215                 cSize = (flush == zbf_end) ?

   216                         ZSTD_compressEnd(zbc->zc, cDst, oSize, zbc->inBuff + zbc->inToCompress, iSize) :

   217                         ZSTD_compressContinue(zbc->zc, cDst, oSize, zbc->inBuff + zbc->inToCompress, iSize);

   218                 if (ZSTD_isError(cSize)) return cSize;

   219                 if (flush == zbf_end) zbc->frameEnded = 1;

   220                 /* prepare next block */

   221                 zbc->inBuffTarget = zbc->inBuffPos + zbc->blockSize;

   222                 if (zbc->inBuffTarget > zbc->inBuffSize)

   223                     zbc->inBuffPos = 0, zbc->inBuffTarget = zbc->blockSize;   /* note : inBuffSize >= blockSize */

   224                 zbc->inToCompress = zbc->inBuffPos;

   225                 if (cDst == op) { op += cSize; break; }   /* no need to flush */

   226                 zbc->outBuffContentSize = cSize;

   227                 zbc->outBuffFlushedSize = 0;

   228                 zbc->stage = ZBUFFcs_flush;   /* continue to flush stage */

   229             }

   230 

   231         case ZBUFFcs_flush:

   232             {   size_t const toFlush = zbc->outBuffContentSize - zbc->outBuffFlushedSize;

   233                 size_t const flushed = ZBUFF_limitCopy(op, oend-op, zbc->outBuff + zbc->outBuffFlushedSize, toFlush);

   234                 op += flushed;

   235                 zbc->outBuffFlushedSize += flushed;

   236                 if (toFlush!=flushed) { someMoreWork = 0; break; } /* dst too small to store flushed data : stop there */

   237                 zbc->outBuffContentSize = zbc->outBuffFlushedSize = 0;

   238                 zbc->stage = ZBUFFcs_load;

   239                 break;

   240             }

   241 

   242         case ZBUFFcs_final:

   243             someMoreWork = 0;   /* do nothing */

   244             break;

   245 

   246         default:

   247             return ERROR(GENERIC);   /* impossible */

   248         }

   249     }

   250 

   251     *srcSizePtr = ip - istart;

   252     *dstCapacityPtr = op - ostart;

   253     if (zbc->frameEnded) return 0;

   254     {   size_t hintInSize = zbc->inBuffTarget - zbc->inBuffPos;

   255         if (hintInSize==0) hintInSize = zbc->blockSize;

   256         return hintInSize;

   257     }

   258 }

   259 

   260 size_t ZBUFF_compressContinue(ZBUFF_CCtx* zbc,

   261                               void* dst, size_t* dstCapacityPtr,

   262                         const void* src, size_t* srcSizePtr)

   263 {

   264     return ZBUFF_compressContinue_generic(zbc, dst, dstCapacityPtr, src, srcSizePtr, zbf_gather);

   265 }

   266 

   267 

   268 

   269 /* ======   Finalize   ====== */

   270 

   271 size_t ZBUFF_compressFlush(ZBUFF_CCtx* zbc, void* dst, size_t* dstCapacityPtr)

   272 {

   273     size_t srcSize = 0;

   274     ZBUFF_compressContinue_generic(zbc, dst, dstCapacityPtr, &srcSize, &srcSize, zbf_flush);  /* use a valid src address instead of NULL */

   275     return zbc->outBuffContentSize - zbc->outBuffFlushedSize;

   276 }

   277 

   278 

   279 size_t ZBUFF_compressEnd(ZBUFF_CCtx* zbc, void* dst, size_t* dstCapacityPtr)

   280 {

   281     BYTE* const ostart = (BYTE*)dst;

   282     BYTE* const oend = ostart + *dstCapacityPtr;

   283     BYTE* op = ostart;

   284 

   285     if (zbc->stage != ZBUFFcs_final) {

   286         /* flush whatever remains */

   287         size_t outSize = *dstCapacityPtr;

   288         size_t srcSize = 0;

   289         size_t const notEnded = ZBUFF_compressContinue_generic(zbc, dst, &outSize, &srcSize, &srcSize, zbf_end);  /* use a valid address instead of NULL */

   290         size_t const remainingToFlush = zbc->outBuffContentSize - zbc->outBuffFlushedSize;

   291         op += outSize;

   292         if (remainingToFlush) {

   293             *dstCapacityPtr = op-ostart;

   294             return remainingToFlush + ZBUFF_endFrameSize + (zbc->checksum * 4);

   295         }

   296         /* create epilogue */

   297         zbc->stage = ZBUFFcs_final;

   298         zbc->outBuffContentSize = !notEnded ? 0 :

   299             ZSTD_compressEnd(zbc->zc, zbc->outBuff, zbc->outBuffSize, NULL, 0);  /* write epilogue into outBuff */

   300     }

   301 

   302     /* flush epilogue */

   303     {   size_t const toFlush = zbc->outBuffContentSize - zbc->outBuffFlushedSize;

   304         size_t const flushed = ZBUFF_limitCopy(op, oend-op, zbc->outBuff + zbc->outBuffFlushedSize, toFlush);

   305         op += flushed;

   306         zbc->outBuffFlushedSize += flushed;

   307         *dstCapacityPtr = op-ostart;

   308         if (toFlush==flushed) zbc->stage = ZBUFFcs_init;  /* end reached */

   309         return toFlush - flushed;

   310     }

   311 }

   312 

   313 

   314 

   315 /* *************************************

   316 *  Tool functions

   317 ***************************************/

   318 size_t ZBUFF_recommendedCInSize(void)  { return ZSTD_BLOCKSIZE_ABSOLUTEMAX; }

   319 size_t ZBUFF_recommendedCOutSize(void) { return ZSTD_compressBound(ZSTD_BLOCKSIZE_ABSOLUTEMAX) + ZSTD_blockHeaderSize + ZBUFF_endFrameSize; }