174 lines
4.8 KiB
C++
174 lines
4.8 KiB
C++
|
/***
|
||
|
Copyright (C) 2022 J Reece Wilson (a/k/a "Reece"). All rights reserved.
|
||
|
|
||
|
File: ZSTDCompressor.hpp
|
||
|
Date: 2022-2-15
|
||
|
Author: Reece
|
||
|
***/
|
||
|
#pragma once
|
||
|
|
||
|
#include "zstd.h"
|
||
|
|
||
|
namespace Aurora::Compression
|
||
|
{
|
||
|
struct ZSTDDeflate : public BaseStream
|
||
|
{
|
||
|
CompressionInfo meta;
|
||
|
|
||
|
ZSTDDeflate(const CompressionInfo &meta) : meta(meta), BaseStream(meta.internalStreamSize)
|
||
|
{}
|
||
|
|
||
|
~ZSTDDeflate()
|
||
|
{
|
||
|
if (auto dctx = AuExchange(cctx_, {}))
|
||
|
{
|
||
|
ZSTD_freeCCtx(dctx);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
bool Init(const AuSPtr<IO::IStreamReader> &reader)
|
||
|
{
|
||
|
size_t ret;
|
||
|
|
||
|
this->reader_ = reader;
|
||
|
this->cctx_ = ZSTD_createCCtx();
|
||
|
|
||
|
if (!this->cctx_)
|
||
|
{
|
||
|
SysPushErrorGen("Couldn't create compressor");
|
||
|
}
|
||
|
|
||
|
ret = ZSTD_CCtx_setParameter(this->cctx_, ZSTD_c_compressionLevel, meta.compressionLevel);
|
||
|
if (ZSTD_isError(ret))
|
||
|
{
|
||
|
SysPushErrorGen("Invalid compression level");
|
||
|
return false;
|
||
|
}
|
||
|
|
||
|
ret = ZSTD_CCtx_setParameter(this->cctx_, ZSTD_c_checksumFlag, 1);
|
||
|
if (ZSTD_isError(ret))
|
||
|
{
|
||
|
SysPushErrorGen("Invalid option");
|
||
|
return false;
|
||
|
}
|
||
|
|
||
|
ret = ZSTD_CCtx_setParameter(this->cctx_, ZSTD_c_nbWorkers, AuMax(meta.threads, AuUInt8(1u)));
|
||
|
if (ZSTD_isError(ret))
|
||
|
{
|
||
|
SysPushErrorGen();
|
||
|
return false;
|
||
|
}
|
||
|
|
||
|
if (!this->_outbuffer)
|
||
|
{
|
||
|
SysPushErrorMem();
|
||
|
return false;
|
||
|
}
|
||
|
|
||
|
this->curPtr_ = this->din_;
|
||
|
this->availIn_ = 0;
|
||
|
SetArray(this->din_);
|
||
|
return true;
|
||
|
}
|
||
|
|
||
|
AuStreamReadWrittenPair_t Ingest_s(AuUInt32 input) override
|
||
|
{
|
||
|
AuUInt32 length = ZSTD_DStreamInSize();
|
||
|
AuUInt32 outFrameLength = ZSTD_DStreamOutSize();
|
||
|
AuUInt32 done {}, read {};
|
||
|
|
||
|
while (read < input)
|
||
|
{
|
||
|
read += IngestForInPointer<char, AuUInt32>(this->reader_, this->curPtr_, this->availIn_, input - read);
|
||
|
|
||
|
if (!this->availIn_)
|
||
|
{
|
||
|
return {read, done};
|
||
|
}
|
||
|
|
||
|
this->input_ = ZSTD_inBuffer {this->curPtr_, this->availIn_, 0};
|
||
|
|
||
|
while (this->input_.pos < this->input_.size)
|
||
|
{
|
||
|
ZSTD_outBuffer output = {this->dout_, outFrameLength, 0};
|
||
|
|
||
|
auto ret = ZSTD_compressStream(this->cctx_, &output, &this->input_);
|
||
|
if (ZSTD_isError(ret))
|
||
|
{
|
||
|
SysPushErrorIO("Compression error: {}", ZSTD_getErrorName(ret));
|
||
|
this->availIn_ -= output.pos;
|
||
|
return AuMakePair(read, 0);
|
||
|
}
|
||
|
|
||
|
done += output.pos;
|
||
|
|
||
|
if (!Write(reinterpret_cast<const AuUInt8 *>(output.dst),
|
||
|
output.pos))
|
||
|
{
|
||
|
return AuMakePair(read, 0);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
this->availIn_ -= this->input_.pos;
|
||
|
}
|
||
|
|
||
|
return {read, done};
|
||
|
}
|
||
|
|
||
|
bool Flush() override
|
||
|
{
|
||
|
return RunFlush(ZSTD_e_continue);
|
||
|
}
|
||
|
|
||
|
bool Finish() override
|
||
|
{
|
||
|
return RunFlush(ZSTD_e_end);
|
||
|
}
|
||
|
|
||
|
bool RunFlush(ZSTD_EndDirective type)
|
||
|
{
|
||
|
AuUInt32 length = ZSTD_DStreamInSize();
|
||
|
AuUInt32 outFrameLength = ZSTD_DStreamOutSize();
|
||
|
|
||
|
if (!this->availIn_)
|
||
|
{
|
||
|
return {};
|
||
|
}
|
||
|
|
||
|
this->input_ = ZSTD_inBuffer {this->curPtr_, this->availIn_, 0};
|
||
|
|
||
|
while (this->input_.pos < this->input_.size)
|
||
|
{
|
||
|
ZSTD_outBuffer output = {this->dout_, outFrameLength, 0};
|
||
|
|
||
|
auto ret = ZSTD_compressStream(this->cctx_, &output, &this->input_);
|
||
|
if (ZSTD_isError(ret))
|
||
|
{
|
||
|
SysPushErrorIO("Compression error: {}", ZSTD_getErrorName(ret));
|
||
|
this->availIn_ -= output.pos;
|
||
|
return {};
|
||
|
}
|
||
|
|
||
|
if (!Write(reinterpret_cast<const AuUInt8 *>(output.dst),
|
||
|
output.pos))
|
||
|
{
|
||
|
return false;
|
||
|
}
|
||
|
}
|
||
|
|
||
|
this->availIn_ -= this->input_.pos;
|
||
|
|
||
|
return true;
|
||
|
}
|
||
|
|
||
|
private:
|
||
|
|
||
|
AuSPtr<IO::IStreamReader> reader_;
|
||
|
ZSTD_CCtx *cctx_;
|
||
|
char din_[ZSTD_BLOCKSIZE_MAX];
|
||
|
char dout_[ZSTD_BLOCKSIZE_MAX + 3 /*ZSTD_BLOCKHEADERSIZE*/];
|
||
|
char *curPtr_;
|
||
|
AuUInt32 availIn_;
|
||
|
ZSTD_inBuffer input_;
|
||
|
};
|
||
|
}
|