From cb3276329aae840245fbef0b8f6dce141ea37da4 Mon Sep 17 00:00:00 2001
From: Yann Collet <cyan@fb.com>
Date: Tue, 23 Aug 2016 00:30:31 +0200
Subject: [PATCH] added sizeof CStream and DStream

---
 lib/compress/zstd_compress.c     |  4 ++++
 lib/decompress/zstd_decompress.c | 18 +++++++++---------
 lib/zstd.h                       |  4 +++-
 tests/fuzzer.c                   |  5 ++++-
 tests/zstreamtest.c              | 20 +++++++++++++++++---
 5 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c
index 00b1408c..bc74da09 100644
--- a/lib/compress/zstd_compress.c
+++ b/lib/compress/zstd_compress.c
@@ -2867,6 +2867,10 @@ size_t ZSTD_initCStream(ZSTD_CStream* zcs, int compressionLevel)
     return ZSTD_initCStream_usingDict(zcs, NULL, 0, compressionLevel);
 }
 
+size_t ZSTD_sizeofCStream(const ZSTD_CStream* zcs)
+{
+    return sizeof(zcs) + ZSTD_sizeofCCtx(zcs->zc) + zcs->outBuffSize + zcs->inBuffSize;
+}
 
 /*======   Compression   ======*/
 
diff --git a/lib/decompress/zstd_decompress.c b/lib/decompress/zstd_decompress.c
index 2311719f..c62ef199 100644
--- a/lib/decompress/zstd_decompress.c
+++ b/lib/decompress/zstd_decompress.c
@@ -980,9 +980,10 @@ size_t ZSTD_decompress(void* dst, size_t dstCapacity, const void* src, size_t sr
 }
 
 
-/*-**********************************
-*   Streaming Decompression API
-************************************/
+/*-**************************************
+*   Advanced Streaming Decompression API
+*   Bufferless and synchronous
+****************************************/
 size_t ZSTD_nextSrcSizeToDecompress(ZSTD_DCtx* dctx) { return dctx->expected; }
 
 ZSTD_nextInputType_e ZSTD_nextInputType(ZSTD_DCtx* dctx) {
@@ -1382,6 +1383,11 @@ size_t ZSTD_initDStream(ZSTD_DStream* zds)
     return ZSTD_initDStream_usingDict(zds, NULL, 0);
 }
 
+size_t ZSTD_sizeofDStream(const ZSTD_DStream* zds)
+{
+    return sizeof(*zds) + ZSTD_sizeofDCtx(zds->zd) + zds->inBuffSize + zds->outBuffSize;
+}
+
 
 /* *** Decompression *** */
 
@@ -1531,9 +1537,3 @@ size_t ZSTD_decompressStream(ZSTD_DStream* zds, ZSTD_outBuffer* output, ZSTD_inB
         return nextSrcSizeHint;
     }
 }
-
-
-
-
-
-
diff --git a/lib/zstd.h b/lib/zstd.h
index d07cc3c7..ab5b4f58 100644
--- a/lib/zstd.h
+++ b/lib/zstd.h
@@ -431,13 +431,15 @@ ZSTD_CStream* ZSTD_createCStream_advanced(ZSTD_customMem customMem);
 size_t ZSTD_initCStream_usingDict(ZSTD_CStream* zcs, const void* dict, size_t dictSize, int compressionLevel);
 size_t ZSTD_initCStream_advanced(ZSTD_CStream* zcs, const void* dict, size_t dictSize,
                                  ZSTD_parameters params, unsigned long long pledgedSrcSize);
+size_t ZSTD_sizeofCStream(const ZSTD_CStream* zcs);
+
 
 /*======   decompression   ======*/
 
 /* advanced */
 ZSTD_DStream* ZSTD_createDStream_advanced(ZSTD_customMem customMem);
 size_t ZSTD_initDStream_usingDict(ZSTD_DStream* zds, const void* dict, size_t dictSize);
-
+size_t ZSTD_sizeofDStream(const ZSTD_DStream* zds);
 
 
 /* ******************************************************************
diff --git a/tests/fuzzer.c b/tests/fuzzer.c
index 87647b9d..bed86f59 100644
--- a/tests/fuzzer.c
+++ b/tests/fuzzer.c
@@ -100,7 +100,6 @@ static unsigned FUZ_rand(unsigned* src)
     return rand32 >> 5;
 }
 
-
 static unsigned FUZ_highbit32(U32 v32)
 {
     unsigned nbBits = 0;
@@ -110,6 +109,10 @@ static unsigned FUZ_highbit32(U32 v32)
 }
 
 
+/*=============================================
+*   Basic Unit tests
+=============================================*/
+
 #define CHECK_V(var, fn)  size_t const var = fn; if (ZSTD_isError(var)) goto _output_error
 #define CHECK(fn)  { CHECK_V(err, fn); }
 #define CHECKPLUS(var, fn, more)  { CHECK_V(var, fn); more; }
diff --git a/tests/zstreamtest.c b/tests/zstreamtest.c
index 7f136b1b..4448f23c 100644
--- a/tests/zstreamtest.c
+++ b/tests/zstreamtest.c
@@ -61,7 +61,6 @@ static const U32 prime1 = 2654435761U;
 static const U32 prime2 = 2246822519U;
 
 
-
 /*-************************************
 *  Display Macros
 **************************************/
@@ -93,7 +92,6 @@ static U32 FUZ_GetMilliStart(void)
     return nCount;
 }
 
-
 static U32 FUZ_GetMilliSpan(U32 nTimeStart)
 {
     U32 const nCurrent = FUZ_GetMilliStart();
@@ -117,7 +115,6 @@ unsigned int FUZ_rand(unsigned int* seedPtr)
     return rand32 >> 5;
 }
 
-
 /*
 static unsigned FUZ_highbit32(U32 v32)
 {
@@ -141,6 +138,11 @@ static void freeFunction(void* opaque, void* address)
     free(address);
 }
 
+
+/*======================================================
+*   Basic Unit tests
+======================================================*/
+
 static int basicUnitTests(U32 seed, double compressibility, ZSTD_customMem customMem)
 {
     int testResult = 0;
@@ -187,6 +189,12 @@ static int basicUnitTests(U32 seed, double compressibility, ZSTD_customMem custo
     cSize += outBuff.pos;
     DISPLAYLEVEL(4, "OK (%u bytes : %.2f%%)\n", (U32)cSize, (double)cSize/COMPRESSIBLE_NOISE_LENGTH*100);
 
+    DISPLAYLEVEL(4, "test%3i : check CStream size : ", testNb++);
+    { size_t const s = ZSTD_sizeofCStream(zc);
+      if (ZSTD_isError(s)) goto _output_error;
+      DISPLAYLEVEL(4, "OK (%u bytes) \n", (U32)s);
+    }
+
     /* skippable frame test */
     DISPLAYLEVEL(4, "test%3i : decompress skippable frame : ", testNb++);
     ZSTD_initDStream_usingDict(zd, CNBuffer, 128 KB);
@@ -218,6 +226,12 @@ static int basicUnitTests(U32 seed, double compressibility, ZSTD_customMem custo
     }   }
     DISPLAYLEVEL(4, "OK \n");
 
+    DISPLAYLEVEL(4, "test%3i : check DStream size : ", testNb++);
+    { size_t const s = ZSTD_sizeofDStream(zd);
+      if (ZSTD_isError(s)) goto _output_error;
+      DISPLAYLEVEL(4, "OK (%u bytes) \n", (U32)s);
+    }
+
     /* Byte-by-byte decompression test */
     DISPLAYLEVEL(4, "test%3i : decompress byte-by-byte : ", testNb++);
     {   size_t r = 1;