stuff-from-scratch/src/compression/deflate/DeflateBlock.cpp

#include "DeflateBlock.h"

#include "ByteUtils.h"
#include "AbstractChecksumCalculator.h"

#include <algorithm>
#include <iostream>
#include <sstream>

DeflateBlock::DeflateBlock(BitStream* inputStream, BitStream* outputStream)
    : mInputStream(inputStream),
      mOutputStream(outputStream)
{

}

std::string DeflateBlock::getMetaData() const
{
    std::stringstream sstr;
    sstr << "DeflateBlock Metadata \n";

    sstr << "Final block: " << mInFinalBlock << '\n';
    sstr << "Compression method: " << Deflate::toString(mCompressionMethod) << '\n';
    sstr << "Uncompressed block length: " << mUncompressedBlockLength << '\n';
    return sstr.str();
}

void DeflateBlock::setIsFinalBlock(bool isFinal)
{
    mInFinalBlock = isFinal;
}

bool DeflateBlock::isFinalBlock() const
{
    return mInFinalBlock;
}

bool DeflateBlock::read()
{
    auto working_byte = *mInputStream->readNextByte();

    //std::cout << mInputStream->logNextNBytes(60);
    //std::cout << "DeflateBlock::read location " << mInputStream->logLocation();

    unsigned char final_block{0};
    mInputStream->readNextNBits(1, final_block);
    mInFinalBlock = bool(final_block);

    unsigned char compression_type{0};
    mInputStream->readNextNBits(2, compression_type);
    mCompressionMethod = static_cast<Deflate::CompressionMethod>(compression_type);

    if (mCompressionMethod == Deflate::CompressionMethod::NONE)
    {
        return readUncompressedStream();
    }
    else if(mCompressionMethod == Deflate::CompressionMethod::FIXED_HUFFMAN)
    {
        return readFixedHuffmanStream();
    }
    else if(mCompressionMethod == Deflate::CompressionMethod::DYNAMIC_HUFFMAN)
    {
        return readDynamicHuffmanStream();
    }
    return false;
}

bool DeflateBlock::readUncompressedStream()
{
    auto byte0 = *mInputStream->readNextByte();
    auto byte1 = *mInputStream->readNextByte();
    mUncompressedBlockLength = (byte0 << 8) | byte1;

    std::cout << "Check block 0: " << ByteUtils::toString(byte0) << std::endl;
    std::cout << "Check block 1: " << ByteUtils::toString(byte1) << std::endl;

    auto byte2 = *mInputStream->readNextByte();
    auto byte3 = *mInputStream->readNextByte();
    uint16_t len_check = (byte2 << 8) | byte3;

    //std::cout << "Check block 2: " << ByteUtils::toString(byte2) << std::endl;
    //std::cout << "Check block 3: " << ByteUtils::toString(byte3) << std::endl;
    //if (!(byte0 ==(~byte2) && byte1 ==(~byte3)))
    //{
        //std::cout << "Uncompressed block length check failed - aborting." << std::endl;
        //return false;
    //}
    //else
    //{
        for(unsigned idx=0; idx<mUncompressedBlockLength;idx++)
        {
            mOutputStream->writeByte(*mInputStream->readNextByte());
        }
    //}
    return true;
}

bool DeflateBlock::readFixedHuffmanStream()
{
    //std::cout << "Reading fixed huffman stream" << std::endl;
    mHuffmanStream = std::make_unique<HuffmanStream>(mInputStream, mOutputStream);

    mHuffmanStream->generateFixedCodeMapping();
    return mHuffmanStream->decode();
}

bool DeflateBlock::readDynamicHuffmanStream()
{
    mHuffmanStream = std::make_unique<HuffmanStream>(mInputStream, mOutputStream);
    return mHuffmanStream->decode();
}

void DeflateBlock::write(uint16_t datalength)
{
    mUncompressedBlockLength = datalength;

    unsigned char working_block{0};
    working_block |= static_cast<unsigned char>(mInFinalBlock);
    working_block |= (static_cast<unsigned char>(mCompressionMethod) << 1);

    if (mCompressionMethod == Deflate::CompressionMethod::NONE)
    {
        writeUncompressedStream(working_block, datalength);
    }
    else if (mCompressionMethod == Deflate::CompressionMethod::FIXED_HUFFMAN)
    {
        mOutputStream->writeNBits(working_block, 3);
        while(auto byte = mInputStream->readNextByte())
        {
            mOutputStream->writeByte(*byte);
        }

        if (const auto& remaining_bits = mInputStream->getRemainingBits(); remaining_bits.second > 0)
        {
            mOutputStream->writeNBits(remaining_bits.first, remaining_bits.second);
        }
    }
}

void DeflateBlock::writeUncompressedStream(unsigned char working_byte, uint16_t datalength)
{
    //std::cout << "Writing compression block header " << ByteUtils::toString(working_byte) << std::endl;
    mOutputStream->writeByte(working_byte);

    //std::cout << "Writing data length " << mUncompressedBlockLength << " " << ByteUtils::toString(mUncompressedBlockLength) << std::endl;
    mOutputStream->writeWord(datalength);

    //std::cout << "Writing iverse data length " << ~mUncompressedBlockLength << " " << ByteUtils::toString(~mUncompressedBlockLength) << std::endl;
    mOutputStream->writeWord(static_cast<uint16_t>(~mUncompressedBlockLength));

    for(unsigned idx=0; idx<mUncompressedBlockLength;idx++)
    {
        auto byte = *mInputStream->readNextByte();
        //std::cout << "Writing next byte " << static_cast<int>(byte) << std::endl;
        mOutputStream->writeByte(byte);
    }
}