From 19ca3fbc037e2708b8f604c13a6ed176480f3642 Mon Sep 17 00:00:00 2001
From: Nick Terrell <terrelln@fb.com>
Date: Fri, 15 Feb 2019 15:24:55 -0800
Subject: [PATCH] [zstdcli] Respect --[no-]compress-literals in benchmark mode

---
 programs/benchzstd.c | 4 +++-
 programs/benchzstd.h | 1 +
 programs/zstdcli.c   | 9 ++++++---
 tests/playTests.sh   | 2 ++
 4 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/programs/benchzstd.c b/programs/benchzstd.c
index 0be74a96..4bd42cfe 100644
--- a/programs/benchzstd.c
+++ b/programs/benchzstd.c
@@ -135,7 +135,8 @@ BMK_advancedParams_t BMK_initAdvancedParams(void) {
         0, /* ldmMinMatch */
         0, /* ldmHashLog */
         0, /* ldmBuckSizeLog */
-        0  /* ldmHashRateLog */
+        0,  /* ldmHashRateLog */
+        ZSTD_lcm_auto /* literalCompressionMode */
     };
     return res;
 }
@@ -174,6 +175,7 @@ static void BMK_initCCtx(ZSTD_CCtx* ctx,
     CHECK_Z(ZSTD_CCtx_setParameter(ctx, ZSTD_c_ldmHashLog, adv->ldmHashLog));
     CHECK_Z(ZSTD_CCtx_setParameter(ctx, ZSTD_c_ldmBucketSizeLog, adv->ldmBucketSizeLog));
     CHECK_Z(ZSTD_CCtx_setParameter(ctx, ZSTD_c_ldmHashRateLog, adv->ldmHashRateLog));
+    CHECK_Z(ZSTD_CCtx_setParameter(ctx, ZSTD_c_literalCompressionMode, (int)adv->literalCompressionMode));
     CHECK_Z(ZSTD_CCtx_setParameter(ctx, ZSTD_c_windowLog, comprParams->windowLog));
     CHECK_Z(ZSTD_CCtx_setParameter(ctx, ZSTD_c_hashLog, comprParams->hashLog));
     CHECK_Z(ZSTD_CCtx_setParameter(ctx, ZSTD_c_chainLog, comprParams->chainLog));
diff --git a/programs/benchzstd.h b/programs/benchzstd.h
index 3a8b893e..d6f5486a 100644
--- a/programs/benchzstd.h
+++ b/programs/benchzstd.h
@@ -116,6 +116,7 @@ typedef struct {
     unsigned ldmHashLog;
     unsigned ldmBucketSizeLog;
     unsigned ldmHashRateLog;
+    ZSTD_literalCompressionMode_e literalCompressionMode;
 } BMK_advancedParams_t;
 
 /* returns default parameters used by nonAdvanced functions */
diff --git a/programs/zstdcli.c b/programs/zstdcli.c
index 60035e56..f5782174 100644
--- a/programs/zstdcli.c
+++ b/programs/zstdcli.c
@@ -568,6 +568,7 @@ int main(int argCount, const char* argv[])
 #ifndef ZSTD_NOBENCH
     BMK_advancedParams_t benchParams = BMK_initAdvancedParams();
 #endif
+    ZSTD_literalCompressionMode_e literalCompressionMode = ZSTD_lcm_auto;
 
 
     /* init */
@@ -660,8 +661,8 @@ int main(int argCount, const char* argv[])
                     if (!strcmp(argument, "--format=lz4")) { suffix = LZ4_EXTENSION; FIO_setCompressionType(prefs, FIO_lz4Compression);  continue; }
 #endif
                     if (!strcmp(argument, "--rsyncable")) { rsyncable = 1; continue; }
-                    if (!strcmp(argument, "--compress-literals")) { FIO_setLiteralCompressionMode(prefs, ZSTD_lcm_huffman); continue; }
-                    if (!strcmp(argument, "--no-compress-literals")) { FIO_setLiteralCompressionMode(prefs, ZSTD_lcm_uncompressed); continue; }
+                    if (!strcmp(argument, "--compress-literals")) { literalCompressionMode = ZSTD_lcm_huffman; continue; }
+                    if (!strcmp(argument, "--no-compress-literals")) { literalCompressionMode = ZSTD_lcm_uncompressed; continue; }
                     if (!strcmp(argument, "--no-progress")) { FIO_setNoProgress(1); continue; }
 
                     /* long commands with arguments */
@@ -998,6 +999,7 @@ int main(int argCount, const char* argv[])
         if (g_ldmHashRateLog != LDM_PARAM_DEFAULT) {
             benchParams.ldmHashRateLog = g_ldmHashRateLog;
         }
+        benchParams.literalCompressionMode = literalCompressionMode;
 
         if (cLevel > ZSTD_maxCLevel()) cLevel = ZSTD_maxCLevel();
         if (cLevelLast > ZSTD_maxCLevel()) cLevelLast = ZSTD_maxCLevel();
@@ -1111,6 +1113,7 @@ int main(int argCount, const char* argv[])
         FIO_setAdaptMin(prefs, adaptMin);
         FIO_setAdaptMax(prefs, adaptMax);
         FIO_setRsyncable(prefs, rsyncable);
+        FIO_setLiteralCompressionMode(prefs, literalCompressionMode);
         if (adaptMin > cLevel) cLevel = adaptMin;
         if (adaptMax < cLevel) cLevel = adaptMax;
 
@@ -1119,7 +1122,7 @@ int main(int argCount, const char* argv[])
         else
           operationResult = FIO_compressMultipleFilenames(prefs, filenameTable, filenameIdx, outFileName, suffix, dictFileName, cLevel, compressionParams);
 #else
-        (void)suffix; (void)adapt; (void)rsyncable; (void)ultra; (void)cLevel; (void)ldmFlag; /* not used when ZSTD_NOCOMPRESS set */
+        (void)suffix; (void)adapt; (void)rsyncable; (void)ultra; (void)cLevel; (void)ldmFlag; (void)literalCompressionMode; /* not used when ZSTD_NOCOMPRESS set */
         DISPLAY("Compression not supported \n");
 #endif
     } else {  /* decompression or test */
diff --git a/tests/playTests.sh b/tests/playTests.sh
index 62b3057e..49354296 100755
--- a/tests/playTests.sh
+++ b/tests/playTests.sh
@@ -207,6 +207,8 @@ $ZSTD tmp -c --no-compress-literals -19      | $ZSTD -t
 $ZSTD tmp -c --compress-literals    -1       | $ZSTD -t
 $ZSTD tmp -c --compress-literals    --fast=1 | $ZSTD -t
 $ZSTD tmp -c --compress-literals    -19      | $ZSTD -t
+$ZSTD -b --fast=1 -i1e1 tmp --compress-literals
+$ZSTD -b --fast=1 -i1e1 tmp --no-compress-literals
 
 $ECHO "test : file removal"
 $ZSTD -f --rm tmp