pzstd: add --stream-size option for stdin compression

contrib/pzstd/Options.cpp

@@ -40,6 +40,16 @@ unsigned parseUnsigned(const char **arg) {
   return result;
 }
 
+unsigned long long parseUnsignedLL(const char **arg) {
+  unsigned long long result = 0;
+  while (**arg >= '0' && **arg <= '9') {
+    result *= 10;
+    result += **arg - '0';
+    ++(*arg);
+  }
+  return result;
+}
+
 const char *getArgument(const char *options, const char **argv, int &i,
                         int argc) {
   if (options[1] != 0) {
@@ -95,6 +105,7 @@ void usage() {
   std::fprintf(stderr, "  -C, --check            : integrity check (default)\n");
   std::fprintf(stderr, "      --no-check         : no integrity check\n");
   std::fprintf(stderr, "  -t, --test             : test compressed file integrity\n");
+  std::fprintf(stderr, "      --stream-size=#    : specify known total input size for stdin compression\n");
   std::fprintf(stderr, "  --                     : all arguments after \"--\" are treated as files\n");
 }
 } // anonymous namespace
@@ -103,7 +114,7 @@ Options::Options()
     : numThreads(defaultNumThreads()), maxWindowLog(23),
       compressionLevel(kDefaultCompressionLevel), decompress(false),
       overwrite(false), keepSource(true), writeMode(WriteMode::Auto),
-      checksum(true), verbosity(2) {}
+      checksum(true), verbosity(2), streamSrcSize(0) {}
 
 Options::Status Options::parse(int argc, const char **argv) {
   bool test = false;
@@ -135,6 +146,27 @@ Options::Status Options::parse(int argc, const char **argv) {
         maxWindowLog = 0;
       } else if (!std::strcmp(arg, "--no-check")) {
         checksum = false;
+      } else if (!std::strncmp(arg, "--stream-size", 13)) {
+        const char *val = nullptr;
+        if (arg[13] == '=') {
+          val = arg + 14;
+        } else if (arg[13] == 0) {
+          ++i;
+          if (i == argc) {
+            std::fprintf(stderr, "Option --stream-size requires an argument\n");
+            return Status::Failure;
+          }
+          val = argv[i];
+        } else {
+          isLongOption = false;
+        }
+        if (val != nullptr) {
+          if (*val < '0' || *val > '9') {
+            std::fprintf(stderr, "Option --stream-size expects a number, but '%s' provided\n", val);
+            return Status::Failure;
+          }
+          streamSrcSize = parseUnsignedLL(&val);
+        }
       } else if (!std::strcmp(arg, "--sparse")) {
         writeMode = WriteMode::Sparse;
         notSupported("Sparse mode");

contrib/pzstd/Options.h

@@ -35,6 +35,7 @@ struct Options {
   WriteMode writeMode;
   bool checksum;
   int verbosity;
+  unsigned long long streamSrcSize;
 
   enum class Status {
     Success, // Successfully parsed options
@@ -46,22 +47,24 @@ struct Options {
   Options(unsigned numThreads, unsigned maxWindowLog, unsigned compressionLevel,
           bool decompress, std::vector<std::string> inputFiles,
           std::string outputFile, bool overwrite, bool keepSource,
-          WriteMode writeMode, bool checksum, int verbosity)
+          WriteMode writeMode, bool checksum, int verbosity,
+          unsigned long long streamSrcSize = 0)
       : numThreads(numThreads), maxWindowLog(maxWindowLog),
         compressionLevel(compressionLevel), decompress(decompress),
         inputFiles(std::move(inputFiles)), outputFile(std::move(outputFile)),
         overwrite(overwrite), keepSource(keepSource), writeMode(writeMode),
-        checksum(checksum), verbosity(verbosity) {}
+        checksum(checksum), verbosity(verbosity),
+        streamSrcSize(streamSrcSize) {}
 
   Status parse(int argc, const char **argv);
 
   ZSTD_parameters determineParameters() const {
-    ZSTD_parameters params = ZSTD_getParams(compressionLevel, 0, 0);
-    params.fParams.contentSizeFlag = 0;
+    ZSTD_parameters params = ZSTD_getParams(compressionLevel, streamSrcSize, 0);
+    params.fParams.contentSizeFlag = (streamSrcSize > 0);
     params.fParams.checksumFlag = checksum;
     if (maxWindowLog != 0 && params.cParams.windowLog > maxWindowLog) {
       params.cParams.windowLog = maxWindowLog;
-      params.cParams = ZSTD_adjustCParams(params.cParams, 0, 0);
+      params.cParams = ZSTD_adjustCParams(params.cParams, streamSrcSize, 0);
     }
     return params;
   }

contrib/pzstd/Pzstd.cpp

@@ -56,6 +56,9 @@ static std::uint64_t handleOneInput(const Options &options,
                              FILE* outputFd,
                              SharedState& state) {
   auto inputSize = fileSizeOrZero(inputFile);
+  if (inputSize == 0 && options.streamSrcSize > 0) {
+    inputSize = options.streamSrcSize;
+  }
   // WorkQueue outlives ThreadPool so in the case of error we are certain
   // we don't accidentally try to call push() on it after it is destroyed
   WorkQueue<std::shared_ptr<BufferWorkQueue>> outs{options.numThreads + 1};