/**

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

 * All rights reserved.

 *

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

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

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

 */

/* *************************************

*  Dependencies

***************************************/

#include <stdlib.h>

#include "error_private.h"

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

#define ZBUFF_STATIC_LINKING_ONLY

#include "zbuff.h"

/* *************************************

*  Constants

***************************************/

static size_t const ZBUFF_endFrameSize = ZSTD_BLOCKHEADERSIZE;

/*-***********************************************************

*  Streaming compression

*

*  A ZBUFF_CCtx object is required to track streaming operation.

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

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

*  ZBUFF_CCtx objects can be reused multiple times.

*

*  Use ZBUFF_compressContinue() repetitively to consume your input.

*  *srcSizePtr and *dstCapacityPtr can be any size.

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

*  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.

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

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

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

*

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

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

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

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

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

*

*  ZBUFF_compressEnd() instructs to finish a frame.

*  It will perform a flush and write frame epilogue.

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

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

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

*

*  Hint : recommended buffer sizes (not compulsory)

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

*  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.

* ***********************************************************/

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

/* *** Resources *** */

struct ZBUFF_CCtx_s {

    ZSTD_CCtx* zc;

    char*  inBuff;

    size_t inBuffSize;

    size_t inToCompress;

    size_t inBuffPos;

    size_t inBuffTarget;

    size_t blockSize;

    char*  outBuff;

    size_t outBuffSize;

    size_t outBuffContentSize;

    size_t outBuffFlushedSize;

    ZBUFF_cStage stage;

    U32    checksum;

    U32    frameEnded;

    ZSTD_customMem customMem;

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

ZBUFF_CCtx* ZBUFF_createCCtx(void)

{

    return ZBUFF_createCCtx_advanced(defaultCustomMem);

}

ZBUFF_CCtx* ZBUFF_createCCtx_advanced(ZSTD_customMem customMem)

{

    ZBUFF_CCtx* zbc;

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

        customMem = defaultCustomMem;

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

        return NULL;

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

    if (zbc==NULL) return NULL;

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

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

    zbc->zc = ZSTD_createCCtx_advanced(customMem);

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

    return zbc;

}

size_t ZBUFF_freeCCtx(ZBUFF_CCtx* zbc)

{

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

    ZSTD_freeCCtx(zbc->zc);

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

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

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

    return 0;

}

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

size_t ZBUFF_compressInit_advanced(ZBUFF_CCtx* zbc,

                                   const void* dict, size_t dictSize,

                                   ZSTD_parameters params, unsigned long long pledgedSrcSize)

{

    /* allocate buffers */

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

        if (zbc->inBuffSize < neededInBuffSize) {

            zbc->inBuffSize = neededInBuffSize;

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

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

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

        }

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

    }

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

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

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

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

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

    }

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

      if (ZSTD_isError(errorCode)) return errorCode; }

    zbc->inToCompress = 0;

    zbc->inBuffPos = 0;

    zbc->inBuffTarget = zbc->blockSize;

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

    zbc->stage = ZBUFFcs_load;

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

    zbc->frameEnded = 0;

    return 0;   /* ready to go */

}

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

{

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

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

}

size_t ZBUFF_compressInit(ZBUFF_CCtx* zbc, int compressionLevel)

{

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

}

/* internal util function */

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

{

    size_t const length = MIN(dstCapacity, srcSize);

    memcpy(dst, src, length);

    return length;

}

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

typedef enum { zbf_gather, zbf_flush, zbf_end } ZBUFF_flush_e;

static size_t ZBUFF_compressContinue_generic(ZBUFF_CCtx* zbc,

                              void* dst, size_t* dstCapacityPtr,

                        const void* src, size_t* srcSizePtr,

                              ZBUFF_flush_e const flush)

{

    U32 someMoreWork = 1;

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

    const char* const iend = istart + *srcSizePtr;

    const char* ip = istart;

    char* const ostart = (char*)dst;

    char* const oend = ostart + *dstCapacityPtr;

    char* op = ostart;

    while (someMoreWork) {

        switch(zbc->stage)

        {

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

        case ZBUFFcs_load:

            /* complete inBuffer */

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

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

                zbc->inBuffPos += loaded;

                ip += loaded;

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

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

            }   }

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

            {   void* cDst;

                size_t cSize;

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

                size_t oSize = oend-op;

                if (oSize >= ZSTD_compressBound(iSize))

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

                else

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

                cSize = (flush == zbf_end) ?

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

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

                if (ZSTD_isError(cSize)) return cSize;

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

                /* prepare next block */

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

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

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

                zbc->inToCompress = zbc->inBuffPos;

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

                zbc->outBuffContentSize = cSize;

                zbc->outBuffFlushedSize = 0;

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

            }

        case ZBUFFcs_flush:

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

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

                op += flushed;

                zbc->outBuffFlushedSize += flushed;

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

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

                zbc->stage = ZBUFFcs_load;

                break;

            }

        case ZBUFFcs_final:

            someMoreWork = 0;   /* do nothing */

            break;

        default:

            return ERROR(GENERIC);   /* impossible */

        }

    }

    *srcSizePtr = ip - istart;

    *dstCapacityPtr = op - ostart;

    if (zbc->frameEnded) return 0;

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

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

        return hintInSize;

    }

}

size_t ZBUFF_compressContinue(ZBUFF_CCtx* zbc,

                              void* dst, size_t* dstCapacityPtr,

                        const void* src, size_t* srcSizePtr)

{

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

}

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

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

{

    size_t srcSize = 0;

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

    return zbc->outBuffContentSize - zbc->outBuffFlushedSize;

}

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

{

    BYTE* const ostart = (BYTE*)dst;

    BYTE* const oend = ostart + *dstCapacityPtr;

    BYTE* op = ostart;

    if (zbc->stage != ZBUFFcs_final) {

        /* flush whatever remains */

        size_t outSize = *dstCapacityPtr;

        size_t srcSize = 0;

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

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

        op += outSize;

        if (remainingToFlush) {

            *dstCapacityPtr = op-ostart;

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

        }

        /* create epilogue */

        zbc->stage = ZBUFFcs_final;

        zbc->outBuffContentSize = !notEnded ? 0 :

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

    }

    /* flush epilogue */

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

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

        op += flushed;

        zbc->outBuffFlushedSize += flushed;

        *dstCapacityPtr = op-ostart;

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

        return toFlush - flushed;

    }

}

/* *************************************

*  Tool functions

***************************************/

size_t ZBUFF_recommendedCInSize(void)  { return ZSTD_BLOCKSIZE_ABSOLUTEMAX; }

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