AuroraRuntime/Source/Compression/Compressors/ZSTDDecompressor.hpp

/***
    Copyright (C) 2022 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: ZSTDDecompressor.hpp
    Date: 2022-2-15
    Author: Reece
***/
#pragma once

#include "zstd.h"

namespace Aurora::Compression
{
    struct ZSTDInflate : BaseStream
    {
        DecompressInfo meta;

        ZSTDInflate(const DecompressInfo &meta) : meta(meta), BaseStream(meta.uInternalStreamSize)
        {}

        ~ZSTDInflate()
        {
            if (auto dctx = AuExchange(this->dctx_, {}))
            {
                ZSTD_freeDCtx(dctx);
            }
        }

        bool Init(const AuSPtr<IO::IStreamReader> &pReader) override
        {
            this->pReader_ = pReader;
            this->dctx_ = ZSTD_createDCtx();

            if (!this->IsValid())
            {
                SysPushErrorMem();
                return false;
            }


            if (!this->dctx_)
            {
                SysPushErrorGen("Couldn't create decompressor");
                return false;
            }

            this->pIterator_ = this->din_;
            this->uAvailableIn_ = 0;
            SetArray(this->din_);
            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);

                if (!this->uAvailableIn_)
                {
                    return {read, done};
                }

                this->input_ = ZSTD_inBuffer {this->pIterator_, this->uAvailableIn_, 0};

                while (this->input_.pos < this->input_.size)
                {
                    ZSTD_outBuffer output = {this->dout_, uOutFrameLength, 0};

                    auto ret = ZSTD_decompressStream(this->dctx_, &output, &this->input_);
                    if (ZSTD_isError(ret))
                    {
                        SysPushErrorIO("Decompression error: {}", ZSTD_getErrorName(ret));
                        this->pReader_.reset();
                        return AuMakePair(read, 0);
                    }

                    done += AuUInt32(output.pos);

                    if (!Write(reinterpret_cast<const AuUInt8 *>(output.dst),
                               AuUInt32(output.pos)))
                    {
                        this->pReader_.reset();
                        return AuMakePair(read, 0);
                    }
                }

                this->uAvailableIn_ = 0;
            }

            return {read, done};
        }

    private:

        AuSPtr<IO::IStreamReader> pReader_;
        ZSTD_DCtx *dctx_ {};
        char din_[ZSTD_BLOCKSIZE_MAX + 3 /*ZSTD_BLOCKHEADERSIZE*/];
        char dout_[ZSTD_BLOCKSIZE_MAX];
        char *pIterator_ {};
        AuUInt32 uAvailableIn_ {};
        ZSTD_inBuffer input_;
    };
}