added ZBUFF_createCCtx_advanced and ZBUFF_createDCtx_advanced

lib/common/zbuff.h

@@ -38,7 +38,8 @@ extern "C" {
 /* *************************************
 *  Dependencies
 ***************************************/
-#include <stddef.h>   /* size_t */
+#include <stddef.h>      /* size_t */
+#include "zstd_static.h" /* ZSTD_customMem */
 
 
 /* ***************************************************************
@@ -60,6 +61,7 @@ extern "C" {
 ***************************************/
 typedef struct ZBUFF_CCtx_s ZBUFF_CCtx;
 ZSTDLIB_API ZBUFF_CCtx* ZBUFF_createCCtx(void);
+ZSTDLIB_API ZBUFF_CCtx* ZBUFF_createCCtx_advanced(ZSTD_customMem customMem);
 ZSTDLIB_API size_t      ZBUFF_freeCCtx(ZBUFF_CCtx* cctx);
 
 ZSTDLIB_API size_t ZBUFF_compressInit(ZBUFF_CCtx* cctx, int compressionLevel);
@@ -112,6 +114,7 @@ ZSTDLIB_API size_t ZBUFF_compressEnd(ZBUFF_CCtx* cctx, void* dst, size_t* dstCap
 
 typedef struct ZBUFF_DCtx_s ZBUFF_DCtx;
 ZSTDLIB_API ZBUFF_DCtx* ZBUFF_createDCtx(void);
+ZSTDLIB_API ZBUFF_DCtx* ZBUFF_createDCtx_advanced(ZSTD_customMem customMem);
 ZSTDLIB_API size_t      ZBUFF_freeDCtx(ZBUFF_DCtx* dctx);
 
 ZSTDLIB_API size_t ZBUFF_decompressInit(ZBUFF_DCtx* dctx);

lib/compress/zbuff_compress.c

@@ -95,14 +95,39 @@ struct ZBUFF_CCtx_s {
     size_t outBuffContentSize;
     size_t outBuffFlushedSize;
     ZBUFF_cStage stage;
+    ZSTD_allocFunction customAlloc;
+    ZSTD_freeFunction customFree;
 };   /* typedef'd tp ZBUFF_CCtx within "zstd_buffered.h" */
 
 ZBUFF_CCtx* ZBUFF_createCCtx(void)
 {
-    ZBUFF_CCtx* zbc = (ZBUFF_CCtx*)malloc(sizeof(ZBUFF_CCtx));
+    ZSTD_customMem customMem = { NULL, NULL };
+    return ZBUFF_createCCtx_advanced(customMem);
+}
+
+ZBUFF_CCtx* ZBUFF_createCCtx_advanced(ZSTD_customMem customMem)
+{
+    ZBUFF_CCtx* zbc;
+
+    if (!customMem.customAlloc && !customMem.customFree)
+    {
+        zbc = (ZBUFF_CCtx*)calloc(1, sizeof(ZBUFF_CCtx));
+        if (zbc==NULL) return NULL;
+        zbc->customAlloc = malloc;
+        zbc->customFree = free;
+        zbc->zc = ZSTD_createCCtx();
+        return zbc;
+    }
+
+    if (!customMem.customAlloc || !customMem.customFree)
+        return NULL;
+
+    zbc = (ZBUFF_CCtx*)customMem.customAlloc(sizeof(ZBUFF_CCtx));
     if (zbc==NULL) return NULL;
-    memset(zbc, 0, sizeof(*zbc));
+    memset(zbc, 0, sizeof(ZBUFF_CCtx));
-    zbc->zc = ZSTD_createCCtx();
+    zbc->customAlloc = customMem.customAlloc; 
+    zbc->customFree = customMem.customFree;
+    zbc->zc = ZSTD_createCCtx_advanced(customMem);
     return zbc;
 }
 
@@ -110,9 +135,9 @@ size_t ZBUFF_freeCCtx(ZBUFF_CCtx* zbc)
 {
     if (zbc==NULL) return 0;   /* support free on NULL */
     ZSTD_freeCCtx(zbc->zc);
-    free(zbc->inBuff);
+    zbc->customFree(zbc->inBuff);
-    free(zbc->outBuff);
+    zbc->customFree(zbc->outBuff);
-    free(zbc);
+    zbc->customFree(zbc);
     return 0;
 }
 
@@ -127,16 +152,16 @@ size_t ZBUFF_compressInit_advanced(ZBUFF_CCtx* zbc,
     {   size_t const neededInBuffSize = (size_t)1 << params.cParams.windowLog;
         if (zbc->inBuffSize < neededInBuffSize) {
             zbc->inBuffSize = neededInBuffSize;
-            free(zbc->inBuff);   /* should not be necessary */
+            zbc->customFree(zbc->inBuff);   /* should not be necessary */
-            zbc->inBuff = (char*)malloc(neededInBuffSize);
+            zbc->inBuff = (char*)zbc->customAlloc(neededInBuffSize);
             if (zbc->inBuff == NULL) return ERROR(memory_allocation);
         }
         zbc->blockSize = MIN(ZSTD_BLOCKSIZE_MAX, neededInBuffSize/2);
     }
     if (zbc->outBuffSize < ZSTD_compressBound(zbc->blockSize)+1) {
         zbc->outBuffSize = ZSTD_compressBound(zbc->blockSize)+1;
-        free(zbc->outBuff);   /* should not be necessary */
+        zbc->customFree(zbc->outBuff);   /* should not be necessary */
-        zbc->outBuff = (char*)malloc(zbc->outBuffSize);
+        zbc->outBuff = (char*)zbc->customAlloc(zbc->outBuffSize);
         if (zbc->outBuff == NULL) return ERROR(memory_allocation);
     }
 

lib/compress/zstd_compress.c

@@ -128,7 +128,7 @@ ZSTD_CCtx* ZSTD_createCCtx_advanced(ZSTD_customMem customMem)
 {
     ZSTD_CCtx* ctx;
 
-    if (!customMem.customAlloc || !customMem.customFree)
+    if (!customMem.customAlloc && !customMem.customFree)
     {
         ctx = (ZSTD_CCtx*) calloc(1, sizeof(ZSTD_CCtx));
         if (!ctx) return NULL;
@@ -138,6 +138,9 @@ ZSTD_CCtx* ZSTD_createCCtx_advanced(ZSTD_customMem customMem)
         return ctx;
     }
 
+    if (!customMem.customAlloc || !customMem.customFree)
+        return NULL;
+
     ctx = (ZSTD_CCtx*) customMem.customAlloc(sizeof(ZSTD_CCtx));
     if (!ctx) return NULL;
 

lib/decompress/zbuff_decompress.c

@@ -82,15 +82,41 @@ struct ZBUFF_DCtx_s {
     size_t blockSize;
     BYTE headerBuffer[ZSTD_FRAMEHEADERSIZE_MAX];
     size_t lhSize;
+    ZSTD_allocFunction customAlloc;
+    ZSTD_freeFunction customFree;
 };   /* typedef'd to ZBUFF_DCtx within "zstd_buffered.h" */
 
 
 ZBUFF_DCtx* ZBUFF_createDCtx(void)
 {
-    ZBUFF_DCtx* zbd = (ZBUFF_DCtx*)malloc(sizeof(ZBUFF_DCtx));
+    ZSTD_customMem customMem = { NULL, NULL };
+    return ZBUFF_createDCtx_advanced(customMem);
+}
+
+ZBUFF_DCtx* ZBUFF_createDCtx_advanced(ZSTD_customMem customMem)
+{
+    ZBUFF_DCtx* zbd;
+
+    if (!customMem.customAlloc && !customMem.customFree)
+    {
+        zbd = (ZBUFF_DCtx*)calloc(1, sizeof(ZBUFF_DCtx));
+        if (zbd==NULL) return NULL;
+        zbd->customAlloc = malloc;
+        zbd->customFree = free;
+        zbd->zd = ZSTD_createDCtx();
+        zbd->stage = ZBUFFds_init;
+        return zbd;
+    }
+
+    if (!customMem.customAlloc || !customMem.customFree)
+        return NULL;
+
+    zbd = (ZBUFF_DCtx*)customMem.customAlloc(sizeof(ZBUFF_DCtx));
     if (zbd==NULL) return NULL;
-    memset(zbd, 0, sizeof(*zbd));
+    memset(zbd, 0, sizeof(ZBUFF_DCtx));
-    zbd->zd = ZSTD_createDCtx();
+    zbd->customAlloc = customMem.customAlloc; 
+    zbd->customFree = customMem.customFree;
+    zbd->zd = ZSTD_createDCtx_advanced(customMem);
     zbd->stage = ZBUFFds_init;
     return zbd;
 }
@@ -99,9 +125,9 @@ size_t ZBUFF_freeDCtx(ZBUFF_DCtx* zbd)
 {
     if (zbd==NULL) return 0;   /* support free on null */
     ZSTD_freeDCtx(zbd->zd);
-    free(zbd->inBuff);
+    zbd->customFree(zbd->inBuff);
-    free(zbd->outBuff);
+    zbd->customFree(zbd->outBuff);
-    free(zbd);
+    zbd->customFree(zbd);
     return 0;
 }
 
@@ -170,16 +196,16 @@ size_t ZBUFF_decompressContinue(ZBUFF_DCtx* zbd,
             {   size_t const blockSize = MIN(1 << zbd->fParams.windowLog, ZSTD_BLOCKSIZE_MAX);
                 zbd->blockSize = blockSize;
                 if (zbd->inBuffSize < blockSize) {
-                    free(zbd->inBuff);
+                    zbd->customFree(zbd->inBuff);
                     zbd->inBuffSize = blockSize;
-                    zbd->inBuff = (char*)malloc(blockSize);
+                    zbd->inBuff = (char*)zbd->customAlloc(blockSize);
                     if (zbd->inBuff == NULL) return ERROR(memory_allocation);
                 }
                 {   size_t const neededOutSize = ((size_t)1 << zbd->fParams.windowLog) + blockSize;
                     if (zbd->outBuffSize < neededOutSize) {
-                        free(zbd->outBuff);
+                        zbd->customFree(zbd->outBuff);
                         zbd->outBuffSize = neededOutSize;
-                        zbd->outBuff = (char*)malloc(neededOutSize);
+                        zbd->outBuff = (char*)zbd->customAlloc(neededOutSize);
                         if (zbd->outBuff == NULL) return ERROR(memory_allocation);
             }   }   }
             zbd->stage = ZBUFFds_read;

lib/decompress/zstd_decompress.c

@@ -153,7 +153,7 @@ ZSTD_DCtx* ZSTD_createDCtx_advanced(ZSTD_customMem customMem)
 {
     ZSTD_DCtx* dctx;
 
-    if (!customMem.customAlloc || !customMem.customFree)
+    if (!customMem.customAlloc && !customMem.customFree)
     {
         dctx = (ZSTD_DCtx*) malloc(sizeof(ZSTD_DCtx));
         if (!dctx) return NULL;
@@ -164,6 +164,9 @@ ZSTD_DCtx* ZSTD_createDCtx_advanced(ZSTD_customMem customMem)
         return dctx;
     }
 
+    if (!customMem.customAlloc || !customMem.customFree)
+        return NULL;
+
     dctx = (ZSTD_DCtx*) customMem.customAlloc(sizeof(ZSTD_DCtx));
     if (!dctx) return NULL;
     dctx->customAlloc = customMem.customAlloc;