AuroraRuntime/Source/Compression/BlockDecompressor.cpp

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

    File: BlockDecompressor.cpp
    Date: 2021-6-17
    Author: Reece
***/
#include <RuntimeInternal.hpp>
#include "Compression.hpp"
#include "BlockDecompressor.hpp"

#include "bzlib.h"
#include "zstd.h"
#include "zlib.h"
#include "lz4.h"

namespace Aurora::Compression
{
    bool BaseStream::ReadByInflatedN(void * buffer, AuUInt32 minimumInflated, AuStreamReadWrittenPair_t &pair, bool ingestUntilEOS)
    {
        AuUInt32 read {}, len {};

        if (ingestUntilEOS)
        {
            while (this->_outbuffer.RemainingBytes() < minimumInflated)
            {
                if (Ingest(4096).second == 0)
                {
                    if (!this->_outbuffer.RemainingBytes())
                    {
                        return false;
                    }

                    break;
                }

                read += 4096;
            }
        }

        len  = this->_outbuffer.Read(buffer, len, buffer == nullptr);
        pair = {read, len};
        return len != 0;
    }

    bool BaseStream::GoBackByInflatedN(AuUInt32 offset)
    {
        return this->_outbuffer.ReaderTryGoBack(offset);
    }

    bool BaseStream::GoForwardByInflatedN(AuUInt32 offset)
    {
        return this->_outbuffer.ReaderTryGoForward(offset);
    }

    class ZSTDInflate : public BaseStream
    {
    public:
        ZSTDInflate() : BaseStream()
        {

        }

        ~ZSTDInflate()
        {
            if (auto dctx = std::exchange(dctx_, {}))
            {
                ZSTD_freeDCtx(dctx);
            }
        }

        bool Init(Aurora::IO::IStreamReader *reader)
        {
            this->reader_ = reader;
            this->dctx_ = ZSTD_createDCtx();
            
            if (!this->dctx_)
            {
                SysPushErrorGen("Couldn't create decompressor");
                return false;
            }

            return true;
        }

        AuStreamReadWrittenPair_t Ingest(AuUInt32 input) override
        {
            AuUInt32 length = ZSTD_DStreamInSize();
            void *din = alloca(length);
            auto outFrameLength = ZSTD_DStreamOutSize();
            void *dout = alloca(outFrameLength);
            AuUInt32 done{}, read{};

            while (read != input)
            {
                AuUInt32 request = std::min(input, length);
                if (this->reader_->Read(din, request) != IO::EStreamError::eErrorNone)
                {
                    return AuMakePair(read, done);
                }
                read += request;

                ZSTD_inBuffer input = { din, request, 0 };
                while (input.pos < input.size)
                {
                    ZSTD_outBuffer output = { dout, outFrameLength, 0 };
                    
                    auto ret = ZSTD_decompressStream(this->dctx_, &output, &input);
                    if (ZSTD_isError(ret))
                    {
                        SysPushErrorIO("Compression error: {}", ret);
                        return AuMakePair(read, 0);
                    }

                    done += output.pos;
                    this->_outbuffer.Write(reinterpret_cast<const AuUInt8 *>(output.dst),
                                           output.pos);
                }
            }
            
            return AuMakePair(read, done);
        }
  
    private:
      
        Aurora::IO::IStreamReader *reader_;
        ZSTD_DCtx *dctx_;
    };

    class ZIPInflate : public BaseStream
    {
    public:
        ~ZIPInflate()
        {
            if (auto ctx = std::exchange(this->init_, {}))
            {
                inflateEnd(&this->ctx_);
            }
        }

        bool Init(Aurora::IO::IStreamReader *reader)
        {
            this->reader_ = reader;
            
            auto ret = inflateInit(&this->ctx_);
            if (ret != Z_OK)
            {
                SysPushErrorMem("Error: {}", ret);
                return false;
            }

            this->init_ = true;
            return true;
        }

        AuStreamReadWrittenPair_t Ingest(AuUInt32 input) override
        {
            int ret;

            AuUInt32 done{}, read{};
            
            while (read < input)
            {
                AuUInt32 request = std::min(input, AuUInt32(AuArraySize(din_)));
                if (this->reader_->Read(din_, request) != IO::EStreamError::eErrorNone)
                {
                    return AuMakePair(read, done);
                }

                read += request;

                this->ctx_.avail_in = request;
                this->ctx_.next_in  = reinterpret_cast<unsigned char *>(din_);

                do
                {
                    this->ctx_.avail_out = AuArraySize(dout_);
                    this->ctx_.next_out  = dout_;

                    if (!this->ctx_.avail_out)
                    {
                        break;
                    }

                    ret = inflate(&this->ctx_, Z_NO_FLUSH);
                    if (ret != Z_OK)
                    {
                        SysPushErrorIO("Error: {}", ret);
                        return AuMakePair(read, 0);
                    }

                    auto have = AuArraySize(dout_) - this->ctx_.avail_out;
                    done += have;

                    this->_outbuffer.Write(reinterpret_cast<const AuUInt8 *>(dout_),
                                           have);

                } while (this->ctx_.avail_out == 0);
                SysAssert(this->ctx_.avail_in == 0);
            }
            
            return AuMakePair(read, done);
        }

    private:

        Aurora::IO::IStreamReader *reader_;
        z_stream ctx_ {};
        bool init_ {};
        unsigned char din_[4096];
        unsigned char dout_[4096];
    };

    class BZIPInflate : public BaseStream
    {
    public:
        ~BZIPInflate()
        {
            if (auto ctx = std::exchange(this->init_, {}))
            {
                BZ2_bzDecompressEnd(&this->ctx_);
            }
        }

        bool Init(Aurora::IO::IStreamReader *reader)
        {
            this->reader_ = reader;
            
            auto ret = BZ2_bzDecompressInit(&this->ctx_, 0, 0);
            if (ret != Z_OK)
            {
                SysPushErrorMem("Error: {}", ret);
                return false;
            }

            this->init_ = true;
            return true;
        }

        AuStreamReadWrittenPair_t Ingest(AuUInt32 input) override
        {
            int ret;
            
            AuUInt32 done{}, read{};
            while (read < input)
            {
                AuUInt32 request = std::min(input, AuUInt32(AuArraySize(din_)));
                if (this->reader_->Read(din_, request) != IO::EStreamError::eErrorNone)
                {
                    return AuMakePair(read, done);
                }
                read += request;

                this->ctx_.avail_in = request;
                this->ctx_.next_in  = reinterpret_cast<char *>(din_);

                do
                {
                    this->ctx_.avail_out = AuArraySize(dout_);
                    this->ctx_.next_out  = dout_;

                    ret = BZ2_bzDecompress(&this->ctx_);
                    if (ret != Z_OK)
                    {
                        SysPushErrorIO("Error: {}", ret);
                        return AuMakePair(read, 0);
                    }

                    auto have = AuArraySize(dout_) - this->ctx_.avail_out;
                    done += have;

                    this->_outbuffer.Write(reinterpret_cast<const AuUInt8 *>(dout_),
                                           have);


                } while (this->ctx_.avail_out == 0);
                SysAssert(this->ctx_.avail_in == 0);
            }

            return AuMakePair(read, done);
        }

    private:

        Aurora::IO::IStreamReader *reader_;
        bz_stream ctx_ {};
        bool init_ {};
        char dout_[4096];
        char din_[4096];
    };

    class LZ4Inflate : public BaseStream
    {
    public:

        LZ4Inflate() : BaseStream(64 * 1024 * 2)
        {}

        ~LZ4Inflate()
        {
        }

        bool Init(Aurora::IO::IStreamReader *reader)
        {
            this->reader_ = reader;
            
            this->lz4Stream_ = LZ4_createStreamDecode();
            if (!this->lz4Stream_)
            {
                SysPushErrorMem();
                return false;
            }

            return true;
        }


        AuStreamReadWrittenPair_t Ingest(AuUInt32 input) override
        {
            return {};
        }

    private:

        Aurora::IO::IStreamReader *reader_;
        LZ4_streamDecode_t* lz4Stream_ {};
    };

    AUKN_SYM ICompressionStream *DecompressorNew(IO::IStreamReader *reader, ECompresionType type)
    {
        BaseStream * ret{};

        switch (type)
        {
        case ECompresionType::eZSTD:
            ret = new ZSTDInflate();
            break;
        case ECompresionType::eBZIP2:
            ret = new BZIPInflate();
            break;
        case ECompresionType::eLZ4:
            ret = new LZ4Inflate();
            break;
        case ECompresionType::eDeflate:
            ret = new ZIPInflate();
            break;
        default:
            ret = nullptr;
            break;
        }
        
        if (ret)
        {
            if (!ret->Init(reader))
            {
                delete ret;
                ret = nullptr;
            }
        }

        return ret;
    }

    AUKN_SYM void DecompressorRelease(ICompressionStream * stream)
    {
        SafeDelete<BaseStream *>(stream);
    }
}