AuroraRuntime/Source/Compression/Compressors/ZSTDCompressor.hpp
2023-10-17 09:28:58 +01:00

277 lines
8.7 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 : BaseStream
{
CompressInfo meta;
ZSTDDeflate(const CompressInfo &meta) : meta(meta), BaseStream(meta.uInternalStreamSize)
{}
~ZSTDDeflate()
{
if (auto cctx = AuExchange(this->cctx_, {}))
{
ZSTD_freeCCtx(cctx);
}
}
bool Init(const AuSPtr<IO::IStreamReader> &pReader) override
{
AuUInt uRet;
if (!this->IsValid())
{
SysPushErrorMem();
return false;
}
this->pReader_ = pReader;
this->cctx_ = ZSTD_createCCtx();
if (!this->cctx_)
{
SysPushErrorGen("Couldn't create compressor");
return false;
}
uRet = ZSTD_CCtx_setParameter(this->cctx_, ZSTD_c_compressionLevel, meta.uCompressionLevel);
if (ZSTD_isError(uRet))
{
SysPushErrorArg("Invalid compression level");
this->SetLastError(uRet, ZSTD_getErrorName(uRet));
return false;
}
uRet = ZSTD_CCtx_setParameter(this->cctx_, ZSTD_c_checksumFlag, meta.bErrorCheck ? 1 : 0);
if (ZSTD_isError(uRet))
{
SysPushErrorArg("Invalid option");
this->SetLastError(uRet, ZSTD_getErrorName(uRet));
return false;
}
if (meta.uOptQuality)
{
uRet = ZSTD_CCtx_setParameter(this->cctx_, ZSTD_c_strategy, meta.uOptQuality.value());
if (ZSTD_isError(uRet))
{
SysPushErrorArg("ZSTD_c_strategy");
this->SetLastError(uRet, ZSTD_getErrorName(uRet));
return false;
}
}
if (auto uThreadsMin = AuHwInfo::GetCPUInfo().uThreads)
{
meta.uThreads = AuMin(uThreadsMin, meta.uThreads);
}
if (meta.uThreads > 1)
{
uRet = ZSTD_CCtx_setParameter(this->cctx_, ZSTD_c_nbWorkers, AuMax(meta.uThreads, AuUInt8(1u)));
if (ZSTD_isError(uRet))
{
this->SetLastError(uRet, ZSTD_getErrorName(uRet));
return false;
}
}
for (const auto &[a, b] : meta.options)
{
uRet = ZSTD_CCtx_setParameter(this->cctx_, (ZSTD_cParameter)a, b);
if (ZSTD_isError(uRet))
{
SysPushErrorArg("Compressor argument assignment {} = {}", a, b);
this->SetLastError(uRet, ZSTD_getErrorName(uRet));
return false;
}
}
this->pIterator_ = this->din_;
this->uAvailableIn_ = 0;
this->SetArray(this->din_);
this->SetOutArray(this->dout_);
return true;
}
AuStreamReadWrittenPair_t Ingest_s(AuUInt32 input) override
{
AuUInt32 uLength = AuUInt32(ZSTD_DStreamInSize());
AuUInt32 uOutFrameLength = AuUInt32(ZSTD_DStreamOutSize());
AuUInt32 done {}, read {};
if (!this->pReader_)
{
return {};
}
while (read < input)
{
read += IngestForInPointer<char, AuUInt32>(this->pReader_, this->pIterator_, this->uAvailableIn_, input - read, this);
if (!this->uAvailableIn_)
{
Flush();
return {read, done};
}
this->input_ = ZSTD_inBuffer {this->pIterator_, this->uAvailableIn_, 0};
size_t uRet {};
bool bLastFrame = input - read == 0;
while ((this->input_.pos != this->input_.size) ||
(/*bLastFrame && uRet*/ false))
{
auto [pMainDOut, uMainDOutLength] = this->GetDOutPair();
ZSTD_outBuffer output = { pMainDOut, uMainDOutLength, 0 };
ZSTD_EndDirective mode = ZSTD_e_continue;// bLastFrame ? ZSTD_e_flush : ZSTD_e_continue;
uRet = ZSTD_compressStream2(this->cctx_, &output, &this->input_, mode);
if (ZSTD_isError(uRet))
{
this->SetLastError(uRet, ZSTD_getErrorName(uRet));
this->uAvailableIn_ -= AuUInt32(output.pos);
this->pReader_.reset();
return AuMakePair(read, 0);
}
this->bContd = true;
this->bLastFrameHasNotFinished_ = bLastFrame && uRet;
if (!output.pos)
{
continue;
}
done += AuUInt32(output.pos);
if (!Write2(reinterpret_cast<const AuUInt8 *>(pMainDOut),
AuUInt32(output.pos)))
{
SysPushErrorIO("Compression Out of Overhead");
this->pReader_.reset();
return AuMakePair(read, 0);
}
}
this->pIterator_ += this->input_.pos;
this->uAvailableIn_ -= AuUInt32(this->input_.pos);
}
return {read, done};
}
bool Flush() override
{
AuUInt32 uLength = AuUInt32(ZSTD_DStreamInSize());
AuUInt32 uOutFrameLength = AuUInt32(ZSTD_DStreamOutSize());
this->input_ = ZSTD_inBuffer { this->pIterator_, this->uAvailableIn_, 0 };
AuUInt uRet {};
while ((this->input_.pos < this->input_.size) ||
(this->bLastFrameHasNotFinished_) ||
(uRet))
{
auto [pMainDOut, uMainDOutLength] = this->GetDOutPair();
ZSTD_outBuffer output = { pMainDOut, uMainDOutLength, 0 };
this->bLastFrameHasNotFinished_ = false;
uRet = ZSTD_compressStream2(this->cctx_, &output, &this->input_, ZSTD_e_flush);
if (ZSTD_isError(uRet))
{
this->SetLastError(uRet, ZSTD_getErrorName(uRet));
this->uAvailableIn_ -= AuUInt32(output.pos);
return {};
}
if (!output.pos)
{
continue;
}
if (!Write2(reinterpret_cast<const AuUInt8 *>(pMainDOut),
AuUInt32(output.pos)))
{
SysPushErrorIO("Compression Out of Overhead");
this->pReader_.reset();
return false;
}
}
this->pIterator_ += this->input_.pos;
this->uAvailableIn_ -= this->input_.pos;
this->bContd = false;
return true;// RunFlush(ZSTD_e_continue);
}
bool Finish() override
{
AuUInt32 uOutFrameLength = AuUInt32(ZSTD_DStreamOutSize());
ZSTD_outBuffer output = { this->dout_, uOutFrameLength, 0 };
if (meta.uThreads == 1)
{
Flush();
}
else
{
if (!this->bContd)
{
return true;
}
}
size_t uRet;
do
{
uRet = ZSTD_endStream(this->cctx_, &output);
if (ZSTD_isError(uRet))
{
this->SetLastError(uRet, ZSTD_getErrorName(uRet));
return {};
}
if (!Write(reinterpret_cast<const AuUInt8 *>(this->dout_),
AuUInt32(output.pos)))
{
SysPushErrorIO("Compression Out of Overhead");
return false;
}
}
while (uRet);
this->bContd = false;
return true;
}
private:
AuSPtr<IO::IStreamReader> pReader_;
ZSTD_CCtx *cctx_ {};
char din_[ZSTD_BLOCKSIZE_MAX];
char dout_[ZSTD_COMPRESSBOUND(ZSTD_BLOCKSIZE_MAX) + 3 +/*ZSTD_BLOCKHEADERSIZE*/ + 4 /*32bit hash*/];
char *pIterator_ {};
AuUInt32 uAvailableIn_ {};
ZSTD_inBuffer input_ {};
bool bLastFrameHasNotFinished_ {};
bool bContd {};
};
}