Only Re-Alloc / Reset When Needed When Switching Between Regular and High Compression Modes

1 files changed, 27 insertions, 13 deletions

diff --git a/lib/lz4frame.c b/lib/lz4frame.c
index bf70b5c..cd61925 100644
--- a/lib/lz4frame.c
+++ b/lib/lz4frame.c
@@ -188,7 +188,15 @@ typedef struct LZ4F_cctx_s
     U64    totalInSize;
     XXH32_state_t xxh;
     void*  lz4CtxPtr;
-    U32    lz4CtxLevel;   /* 0: unallocated;  1: LZ4_stream_t;  3: LZ4_streamHC_t */
+    /* lz4CtxLevel records both what size of memory has been allocated into the
+     * ctx pointer field and what it's being used as.
+     * - The low bit (& 1) value indicates whether the ctx is being used for
+     *   hc (1) or not (0).
+     * - The next bit (& 2) indicates whether the allocated memory is big
+     *   enough for a non-hc context.
+     * - The next bit (& 4) indicates whether the allocated memory is big
+     *   enough for an hc context. */
+    U32    lz4CtxLevel;
 } LZ4F_cctx_t;
 
 
@@ -356,7 +364,7 @@ size_t LZ4F_compressFrame_usingCDict(void* dstBuffer, size_t dstCapacity,
     if (prefs.compressionLevel < LZ4HC_CLEVEL_MIN) {
         LZ4_resetStream(&lz4ctx);
         cctxI.lz4CtxPtr = &lz4ctx;
-        cctxI.lz4CtxLevel = 1;
+        cctxI.lz4CtxLevel = 2;
     }  /* fast compression context pre-created on stack */
 
     memset(&options, 0, sizeof(options));
@@ -513,21 +521,27 @@ size_t LZ4F_compressBegin_usingCDict(LZ4F_cctx* cctxPtr,
     cctxPtr->prefs = *preferencesPtr;
 
     /* Ctx Management */
-    {   U32 const ctxTypeID = (cctxPtr->prefs.compressionLevel < LZ4HC_CLEVEL_MIN) ? 1 : 2;  /* 0:nothing ; 1:LZ4 table ; 2:HC tables */
-        if (cctxPtr->lz4CtxLevel < ctxTypeID) {
+    {   U32 const ctxSufficientAllocBits = (cctxPtr->prefs.compressionLevel < LZ4HC_CLEVEL_MIN) ? 2 : 6;
+        U32 const ctxTypeBit = cctxPtr->prefs.compressionLevel >= LZ4HC_CLEVEL_MIN;
+        if ((cctxPtr->lz4CtxLevel & ctxSufficientAllocBits) != ctxSufficientAllocBits) {
             FREEMEM(cctxPtr->lz4CtxPtr);
-            if (cctxPtr->prefs.compressionLevel < LZ4HC_CLEVEL_MIN)
+            if (cctxPtr->prefs.compressionLevel < LZ4HC_CLEVEL_MIN) {
                 cctxPtr->lz4CtxPtr = (void*)LZ4_createStream();
-            else
+            } else {
                 cctxPtr->lz4CtxPtr = (void*)LZ4_createStreamHC();
+            }
             if (cctxPtr->lz4CtxPtr == NULL) return err0r(LZ4F_ERROR_allocation_failed);
-            cctxPtr->lz4CtxLevel = ctxTypeID;
-        } else if (cctxPtr->lz4CtxLevel != ctxTypeID) {
-            /* otherwise, we must be transitioning from HC -> LZ4.
-             * In that case, avoid reallocating, since a LZ4 ctx
-             * fits in an HC ctx. Just reset. */
-            LZ4_resetStream((LZ4_stream_t *) cctxPtr->lz4CtxPtr);
-            cctxPtr->lz4CtxLevel = ctxTypeID;
+            cctxPtr->lz4CtxLevel = ctxSufficientAllocBits | ctxTypeBit;
+        } else if ((cctxPtr->lz4CtxLevel & 1) != ctxTypeBit) {
+            /* otherwise, a sufficient buffer is allocated, but we need to
+             * reset it to the correct context type */
+            if (cctxPtr->prefs.compressionLevel < LZ4HC_CLEVEL_MIN) {
+                LZ4_resetStream((LZ4_stream_t *) cctxPtr->lz4CtxPtr);
+                cctxPtr->lz4CtxLevel &= ~1;
+            } else {
+                LZ4_resetStreamHC((LZ4_streamHC_t *) cctxPtr->lz4CtxPtr, cctxPtr->prefs.compressionLevel);
+                cctxPtr->lz4CtxLevel |= 1;
+            }
         }
     }