static size_t Process(ZSTD_CStream *stream, ZSTD_inBuffer *input,
ZSTD_outBuffer *output, bool flush)
{
if (flush)
return ZSTD_flushStream(stream, output);
else
return ZSTD_compressStream(stream, output, input);
}
std::unique_ptr<IOBuf> ZSTDCodec::doCompress(const IOBuf* data) {
// Support earlier versions of the codec (working with a single IOBuf,
// and using ZSTD_decompress which requires ZSTD frame to contain size,
// which isn't populated by streaming API).
if (!data->isChained()) {
auto out = IOBuf::createCombined(ZSTD_compressBound(data->length()));
const auto rc = ZSTD_compress(
out->writableData(),
out->capacity(),
data->data(),
data->length(),
level_);
zstdThrowIfError(rc);
out->append(rc);
return out;
}
auto zcs = ZSTD_createCStream();
SCOPE_EXIT {
ZSTD_freeCStream(zcs);
};
auto rc = ZSTD_initCStream(zcs, level_);
zstdThrowIfError(rc);
Cursor cursor(data);
auto result = IOBuf::createCombined(ZSTD_compressBound(cursor.totalLength()));
ZSTD_outBuffer out;
out.dst = result->writableTail();
out.size = result->capacity();
out.pos = 0;
for (auto buffer = cursor.peekBytes(); !buffer.empty();) {
ZSTD_inBuffer in;
in.src = buffer.data();
in.size = buffer.size();
for (in.pos = 0; in.pos != in.size;) {
rc = ZSTD_compressStream(zcs, &out, &in);
zstdThrowIfError(rc);
}
cursor.skip(in.size);
buffer = cursor.peekBytes();
}
rc = ZSTD_endStream(zcs, &out);
zstdThrowIfError(rc);
CHECK_EQ(rc, 0);
result->append(out.pos);
return result;
}
static void write_zstd(int f, int fd, const char *arg)
{
ZSTD_inBuffer zin;
ZSTD_outBuffer zout;
size_t const inbufsz = ZSTD_CStreamInSize();
zin.src = malloc(inbufsz);
zout.size = ZSTD_CStreamOutSize();
zout.dst = malloc(zout.size);
if (!zin.src || !zout.dst)
goto zstd_w_no_stream;
ZSTD_CStream* const stream = ZSTD_createCStream();
if (!stream)
goto zstd_w_no_stream;
if (ZSTD_isError(ZSTD_initCStream(stream, 3)))
goto zstd_w_error;
size_t s;
while ((s = read(fd, (void*)zin.src, inbufsz)) > 0)
{
zin.size = s;
zin.pos = 0;
while (zin.pos < zin.size)
{
zout.pos = 0;
size_t w = ZSTD_compressStream(stream, &zout, &zin);
if (ZSTD_isError(w))
goto zstd_w_error;
if (write(f, zout.dst, zout.pos) != (ssize_t)zout.pos)
goto zstd_w_error;
}
}
zout.pos = 0;
ZSTD_endStream(stream, &zout);
// no way to handle an error here
write(f, zout.dst, zout.pos);
zstd_w_error:
ZSTD_freeCStream(stream);
zstd_w_no_stream:
free((void*)zin.src);
free(zout.dst);
close(f);
close(fd);
}
size_t ZBUFF_compressContinue(ZBUFF_CCtx* zbc,
void* dst, size_t* dstCapacityPtr,
const void* src, size_t* srcSizePtr)
{
size_t result;
ZSTD_outBuffer outBuff;
ZSTD_inBuffer inBuff;
outBuff.dst = dst;
outBuff.pos = 0;
outBuff.size = *dstCapacityPtr;
inBuff.src = src;
inBuff.pos = 0;
inBuff.size = *srcSizePtr;
result = ZSTD_compressStream(zbc, &outBuff, &inBuff);
*dstCapacityPtr = outBuff.pos;
*srcSizePtr = inBuff.pos;
return result;
}
size_t ZSTD_seekable_compressStream(ZSTD_seekable_CStream* zcs, ZSTD_outBuffer* output, ZSTD_inBuffer* input)
{
const BYTE* const inBase = (const BYTE*) input->src + input->pos;
size_t inLen = input->size - input->pos;
inLen = MIN(inLen, (size_t)(zcs->maxFrameSize - zcs->frameDSize));
/* if we haven't finished flushing the last frame, don't start writing a new one */
if (inLen > 0) {
ZSTD_inBuffer inTmp = { inBase, inLen, 0 };
size_t const prevOutPos = output->pos;
size_t const ret = ZSTD_compressStream(zcs->cstream, output, &inTmp);
if (zcs->framelog.checksumFlag) {
XXH64_update(&zcs->xxhState, inBase, inTmp.pos);
}
zcs->frameCSize += output->pos - prevOutPos;
zcs->frameDSize += inTmp.pos;
input->pos += inTmp.pos;
if (ZSTD_isError(ret)) return ret;
}
if (zcs->maxFrameSize == zcs->frameDSize) {
/* log the frame and start over */
size_t const ret = ZSTD_seekable_endFrame(zcs, output);
if (ZSTD_isError(ret)) return ret;
/* get the client ready for the next frame */
return (size_t)zcs->maxFrameSize;
}
return (size_t)(zcs->maxFrameSize - zcs->frameDSize);
}
static void compressFile_orDie(const char* fname, const char* outName, int cLevel)
{
FILE* const fin = fopen_orDie(fname, "rb");
FILE* const fout = fopen_orDie(outName, "wb");
size_t const buffInSize = ZSTD_CStreamInSize(); /* can always read one full block */
void* const buffIn = malloc_orDie(buffInSize);
size_t const buffOutSize = ZSTD_CStreamOutSize(); /* can always flush a full block */
void* const buffOut = malloc_orDie(buffOutSize);
ZSTD_CStream* const cstream = ZSTD_createCStream();
if (cstream==NULL) { fprintf(stderr, "ZSTD_createCStream() error \n"); exit(10); }
size_t const initResult = ZSTD_initCStream(cstream, cLevel);
if (ZSTD_isError(initResult)) { fprintf(stderr, "ZSTD_initCStream() error : %s \n", ZSTD_getErrorName(initResult)); exit(11); }
size_t read, toRead = buffInSize;
while( (read = fread_orDie(buffIn, toRead, fin)) ) {
ZSTD_inBuffer input = { buffIn, read, 0 };
while (input.pos < input.size) {
ZSTD_outBuffer output = { buffOut, buffOutSize, 0 };
toRead = ZSTD_compressStream(cstream, &output , &input); /* toRead is guaranteed to be <= ZSTD_CStreamInSize() */
if (ZSTD_isError(toRead)) { fprintf(stderr, "ZSTD_compressStream() error : %s \n", ZSTD_getErrorName(toRead)); exit(12); }
if (toRead > buffInSize) toRead = buffInSize; /* Safely handle when `buffInSize` is manually changed to a smaller value */
fwrite_orDie(buffOut, output.pos, fout);
}
}
ZSTD_outBuffer output = { buffOut, buffOutSize, 0 };
size_t const remainingToFlush = ZSTD_endStream(cstream, &output); /* close frame */
if (remainingToFlush) { fprintf(stderr, "not fully flushed"); exit(13); }
fwrite_orDie(buffOut, output.pos, fout);
ZSTD_freeCStream(cstream);
fclose_orDie(fout);
fclose_orDie(fin);
free(buffIn);
free(buffOut);
}
size_t ZSTDMT_compressStream(ZSTDMT_CCtx* zcs, ZSTD_outBuffer* output, ZSTD_inBuffer* input)
{
size_t const newJobThreshold = zcs->dictSize + zcs->targetSectionSize + zcs->marginSize;
if (zcs->frameEnded) return ERROR(stage_wrong); /* current frame being ended. Only flush is allowed. Restart with init */
if (zcs->nbThreads==1) return ZSTD_compressStream(zcs->cstream, output, input);
/* fill input buffer */
{ size_t const toLoad = MIN(input->size - input->pos, zcs->inBuffSize - zcs->inBuff.filled);
memcpy((char*)zcs->inBuff.buffer.start + zcs->inBuff.filled, input->src, toLoad);
input->pos += toLoad;
zcs->inBuff.filled += toLoad;
}
if ( (zcs->inBuff.filled >= newJobThreshold) /* filled enough : let's compress */
&& (zcs->nextJobID <= zcs->doneJobID + zcs->jobIDMask) ) { /* avoid overwriting job round buffer */
CHECK_F( ZSTDMT_createCompressionJob(zcs, zcs->targetSectionSize, 0) );
}
/* check for data to flush */
CHECK_F( ZSTDMT_flushNextJob(zcs, output, (zcs->inBuff.filled == zcs->inBuffSize)) ); /* block if it wasn't possible to create new job due to saturation */
/* recommended next input size : fill current input buffer */
return zcs->inBuffSize - zcs->inBuff.filled; /* note : could be zero when input buffer is fully filled and no more availability to create new job */
}
ByteArray CompressFile(const FilePath& path, const int32 compressionLevel)
{
BinaryReader reader(path);
if (!reader)
{
return ByteArray();
}
const size_t inputBufferSize = ZSTD_CStreamInSize();
const auto pInputBuffer = std::make_unique<Byte[]>(inputBufferSize);
const size_t outputBufferSize = ZSTD_CStreamOutSize();
const auto pOutputBuffer = std::make_unique<Byte[]>(outputBufferSize);
ZSTD_CStream* const cStream = ZSTD_createCStream();
if (!cStream)
{
return ByteArray();
}
const size_t initResult = ZSTD_initCStream(cStream, compressionLevel);
if (ZSTD_isError(initResult))
{
ZSTD_freeCStream(cStream);
return ByteArray();
}
size_t toRead = inputBufferSize;
Array<Byte> buffer;
while (const size_t read = static_cast<size_t>(reader.read(pInputBuffer.get(), toRead)))
{
ZSTD_inBuffer input = { pInputBuffer.get(), read, 0 };
while (input.pos < input.size)
{
ZSTD_outBuffer output = { pOutputBuffer.get(), outputBufferSize, 0 };
toRead = ZSTD_compressStream(cStream, &output, &input);
if (ZSTD_isError(toRead))
{
ZSTD_freeCStream(cStream);
return ByteArray();
}
if (toRead > inputBufferSize)
{
toRead = inputBufferSize;
}
buffer.insert(buffer.end(), pOutputBuffer.get(), pOutputBuffer.get() + output.pos);
}
}
ZSTD_outBuffer output = { pOutputBuffer.get(), outputBufferSize, 0 };
const size_t remainingToFlush = ZSTD_endStream(cStream, &output);
ZSTD_freeCStream(cStream);
if (remainingToFlush)
{
return ByteArray();
}
buffer.insert(buffer.end(), pOutputBuffer.get(), pOutputBuffer.get() + output.pos);
return ByteArray(std::move(buffer));
}
bool CompressFileToFile(const FilePath& inputPath, const FilePath& outputPath, const int32 compressionLevel)
{
BinaryReader reader(inputPath);
if (!reader)
{
return false;
}
const size_t inputBufferSize = ZSTD_CStreamInSize();
const auto pInputBuffer = std::make_unique<Byte[]>(inputBufferSize);
const size_t outputBufferSize = ZSTD_CStreamOutSize();
const auto pOutputBuffer = std::make_unique<Byte[]>(outputBufferSize);
ZSTD_CStream* const cStream = ZSTD_createCStream();
if (!cStream)
{
return false;
}
const size_t initResult = ZSTD_initCStream(cStream, compressionLevel);
if (ZSTD_isError(initResult))
{
ZSTD_freeCStream(cStream);
return false;
}
size_t toRead = inputBufferSize;
BinaryWriter writer(outputPath);
if (!writer)
{
ZSTD_freeCStream(cStream);
return false;
}
while (const size_t read = static_cast<size_t>(reader.read(pInputBuffer.get(), toRead)))
{
ZSTD_inBuffer input = { pInputBuffer.get(), read, 0 };
while (input.pos < input.size)
{
ZSTD_outBuffer output = { pOutputBuffer.get(), outputBufferSize, 0 };
toRead = ZSTD_compressStream(cStream, &output, &input);
if (ZSTD_isError(toRead))
{
writer.clear();
ZSTD_freeCStream(cStream);
return false;
}
if (toRead > inputBufferSize)
{
toRead = inputBufferSize;
}
writer.write(pOutputBuffer.get(), output.pos);
}
}
ZSTD_outBuffer output = { pOutputBuffer.get(), outputBufferSize, 0 };
const size_t remainingToFlush = ZSTD_endStream(cStream, &output);
ZSTD_freeCStream(cStream);
if (remainingToFlush)
{
writer.clear();
return false;
}
writer.write(pOutputBuffer.get(), output.pos);
return true;
}