From 2f907fc2513f392a82307e910666f8553348953a Mon Sep 17 00:00:00 2001 From: VirajMishra1 Date: Wed, 2 Sep 2026 20:51:50 -0400 Subject: [PATCH] pzstd: add --stream-size option for stdin compression Add --stream-size flag to pzstd, matching the existing flag in the main zstd CLI. When compressing from stdin/pipe, this tells the compressor the total input size, enabling it to write the content size in the frame header and optimize window size for better compression ratios. Closes #3171 --- contrib/pzstd/Options.cpp | 34 +++++++++++++++++++++++++++++++++- contrib/pzstd/Options.h | 13 ++++++++----- contrib/pzstd/Pzstd.cpp | 3 +++ 3 files changed, 44 insertions(+), 6 deletions(-) diff --git a/contrib/pzstd/Options.cpp b/contrib/pzstd/Options.cpp index 90841b9bb04..a5c33a15da9 100644 --- a/contrib/pzstd/Options.cpp +++ b/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"); diff --git a/contrib/pzstd/Options.h b/contrib/pzstd/Options.h index 92c18a350cf..e9a252169aa 100644 --- a/contrib/pzstd/Options.h +++ b/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 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; } diff --git a/contrib/pzstd/Pzstd.cpp b/contrib/pzstd/Pzstd.cpp index 048a006b3b3..807cefca8f2 100644 --- a/contrib/pzstd/Pzstd.cpp +++ b/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> outs{options.numThreads + 1};