zlibWrapper: ZWRAP_createCCtx and ZWRAP_freeCCtx use custom memory allocation functions

diff --git a/lib/common/zstd_internal.h b/lib/common/zstd_internal.h
index 8432c1dd15f1e8a39ae7fc5c2f9978ea59e09c8a..230ce79fcd7e3f3cd5405efe5de7519743b0e055 100644 (file)
--- a/lib/common/zstd_internal.h
+++ b/lib/common/zstd_internal.h
@@ -259,5 +259,5 @@ int ZSTD_isSkipFrame(ZSTD_DCtx* dctx);
 void* ZSTD_defaultAllocFunction(void* opaque, size_t size);
 void ZSTD_defaultFreeFunction(void* opaque, void* address);
 static ZSTD_customMem const defaultCustomMem = { ZSTD_defaultAllocFunction, ZSTD_defaultFreeFunction, NULL };
- 
+
 #endif   /* ZSTD_CCOMMON_H_MODULE */

diff --git a/zlibWrapper/zstd_zlibwrapper.c b/zlibWrapper/zstd_zlibwrapper.c
index a4e22a118870bd463d9ac5ca80f02f5e55f08684..937895cd20898a9d8133ffee0600986b3c1f3c70 100644 (file)
--- a/zlibWrapper/zstd_zlibwrapper.c
+++ b/zlibWrapper/zstd_zlibwrapper.c
@@ -38,6 +38,7 @@
 #include "zstd_static.h"      /* ZSTD_MAGICNUMBER */
 #include "zbuff.h"
 #include "zbuff_static.h"      /* ZBUFF_createCCtx_advanced */
+#include "zstd_internal.h"     /* defaultCustomMem */
 
 
 #define Z_INFLATE_SYNC              8
@@ -47,8 +48,6 @@
 #define LOG_WRAPPER(...)  // printf(__VA_ARGS__)
 
 
-#define MIN(a,b) ((a)<(b)?(a):(b))
-
 #define FINISH_WITH_ERR(msg) { \
     fprintf(stderr, "ERROR: %s\n", msg); \
     return Z_MEM_ERROR; \
@@ -99,22 +98,31 @@ typedef struct {
     ZBUFF_CCtx* zbc;
     size_t bytesLeft;
     int compressionLevel;
+    ZSTD_customMem customMem;
     z_stream allocFunc; /* copy of zalloc, zfree, opaque */
 } ZWRAP_CCtx;
 
 
 ZWRAP_CCtx* ZWRAP_createCCtx(z_streamp strm)
 {
-    ZWRAP_CCtx* zwc = (ZWRAP_CCtx*)malloc(sizeof(ZWRAP_CCtx));
-    if (zwc==NULL) return NULL;
-    memset(zwc, 0, sizeof(*zwc));
+    ZWRAP_CCtx* zwc;
+
     if (strm->zalloc && strm->zfree) {
-       ZSTD_customMem ZWRAP_customMem = { ZWRAP_allocFunction, ZWRAP_freeFunction, &zwc->allocFunc };
-       memcpy(&zwc->allocFunc, strm, sizeof(z_stream));
-       zwc->zbc = ZBUFF_createCCtx_advanced(ZWRAP_customMem);
+        zwc = (ZWRAP_CCtx*)strm->zalloc(strm->opaque, 1, sizeof(ZWRAP_CCtx));
+        if (zwc==NULL) return NULL;
+        memset(zwc, 0, sizeof(ZWRAP_CCtx));
+        memcpy(&zwc->allocFunc, strm, sizeof(z_stream));
+        { ZSTD_customMem ZWRAP_customMem = { ZWRAP_allocFunction, ZWRAP_freeFunction, &zwc->allocFunc };
+          memcpy(&zwc->customMem, &ZWRAP_customMem, sizeof(ZSTD_customMem));
+        }
+    } else {
+        zwc = (ZWRAP_CCtx*)defaultCustomMem.customAlloc(defaultCustomMem.opaque, sizeof(ZWRAP_CCtx));
+        if (zwc==NULL) return NULL;
+        memset(zwc, 0, sizeof(ZWRAP_CCtx));
+        memcpy(&zwc->customMem, &defaultCustomMem, sizeof(ZSTD_customMem));
     }
-    else
-       zwc->zbc = ZBUFF_createCCtx();
+
+    zwc->zbc = ZBUFF_createCCtx_advanced(zwc->customMem);
     return zwc;
 }
 
@@ -123,7 +131,7 @@ size_t ZWRAP_freeCCtx(ZWRAP_CCtx* zwc)
 {
     if (zwc==NULL) return 0;   /* support free on NULL */
     ZBUFF_freeCCtx(zwc->zbc);
-    free(zwc);
+    zwc->customMem.customFree(zwc->customMem.opaque, zwc);
     return 0;
 }