HuffmanZipperByTsion

#include <assert.h>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <vector>
#include <stack>

#include "Huffman.h"

using namespace std;
typedef char byte;

class HNode
{
private:
  bool leaf;
  byte letter;
  int frequency;
  HNode *left;
  HNode *right;
public:
  HNode(byte letter, int frequency, bool leaf, HNode *left, HNode *right)
  {
    this->letter = letter;
    this->frequency = frequency;
    this->leaf = leaf;
    this->left = left;
    this->right = right;
  }

  HNode(const HNode *&hnode)
  {
    this->letter = hnode->letter;
    this->frequency = hnode->frequency;
    this->leaf = hnode->leaf;
    this->left = hnode->left;
    this->right = hnode->right;
  }

  HNode &operator=(const HNode *&hnode)
  {
    this->letter = hnode->letter;
    this->frequency = hnode->frequency;
    this->leaf = hnode->leaf;
    this->left = hnode->left;
    this->right = hnode->right;
    return *this;
  }

  bool IsLeaf()
  {
    return this->leaf;
  }

  byte &GetLetter()
  {
    return this->letter;
  }

  int GetFrequency()
  {
    return this->frequency;
  }

  HNode *&GetLeft()
  {
    return this->left;
  }

  HNode *&GetRight()
  {
    return this->right;
  }
};


struct HNodeComparator
{
  bool operator()(HNode *arg0, HNode *arg1)
  {
    return arg0->GetFrequency() > arg1->GetFrequency();
  }
};


class HuffmanTree
{
public:
  vector<byte> dataforzip;
  vector<bool> packedtree;
private:
  HNode *root;
  map<byte, int> freqcounter;
  priority_queue<HNode *, vector<HNode *>, HNodeComparator> heap;
public:
  HuffmanTree(const vector<byte> &inputstreamdata)
  {
    dataforzip = vector<byte>(inputstreamdata.begin(), inputstreamdata.end());
    this->calculateFrequency_();
    this->buildTree_();
    root = heap.top();
  }

  ~HuffmanTree()
  {
    if (!dataforzip.empty())
      dataforzip.clear();
    if (!freqcounter.empty())
      freqcounter.clear();
    while (!heap.empty())
      heap.pop();
    clear_(root);
  }

  HNode *&GetRoot()
  {
    return this->root;
  }

  void CompleteCodeTable(map<byte, vector<bool> > &encodetable, vector<bool> &path)
  {
    completeCodeTable_(encodetable, path, root);
  }

  void GetByte(const vector<bool> &code, byte &letter)
  {
    getByte_(code, letter, this->root);
  }

  void GetPackedTree(vector<bool> &buffer)
  {
    if (packedtree.size() == 0)
      packTree_(root);
    getPackedTreeSize_(buffer);
    for (int i = 0; i < packedtree.size(); i++)
      buffer.push_back(packedtree[i]);
  }

private:
  void calculateFrequency_()
  {
    for (int i = 0; i < dataforzip.size(); i++)
      freqcounter[dataforzip[i]]++;
  }

  void buildTree_()
  {
    HNode *prevnode = nullptr;
    for (map<byte, int>::iterator it = freqcounter.begin(); it != freqcounter.end(); it++)
    {
      prevnode = new HNode(it->first, it->second, true, nullptr, nullptr);
      heap.push(prevnode);
    }
    HNode *currnode = nullptr;
    while (heap.size() > 1)
    {
      prevnode = heap.top();
      heap.pop();
      currnode = new HNode('\0', prevnode->GetFrequency() + heap.top()->GetFrequency(), false, prevnode, heap.top());
      heap.pop();
      heap.push(currnode);
    }
  }

  void clear_(HNode *&hnode)
  {
    if (hnode == nullptr)
      return;
    clear_(hnode->GetLeft());
    clear_(hnode->GetRight());
    delete hnode;
  }

  void completeCodeTable_(map<byte, vector<bool> > &encodetable, vector<bool> &path, HNode *&hnode)
  {
    if (hnode == nullptr)
      path.pop_back();
    else
    {
      if (hnode->IsLeaf())
      {
        encodetable[hnode->GetLetter()] = vector<bool>(path.begin(), path.end());
        path.pop_back();
      }
      else
      {
        path.push_back(false);
        completeCodeTable_(encodetable, path, hnode->GetLeft());
        path.push_back(true);
        completeCodeTable_(encodetable, path, hnode->GetRight());
        if (path.size() == 0)
          return;
        else
          path.pop_back();
      }
    }
  }

  void getByte_(const vector<bool> &code, byte &letter, HNode *&hnode, int bit = 0)
  {
    if (bit == code.size())
    {
      assert(hnode->IsLeaf());
      letter = hnode->GetLetter();
    }
    else
    {
      if (!code[bit])
        getByte_(code, letter, hnode->GetLeft(), bit + 1);
      else
        getByte_(code, letter, hnode->GetRight(), bit + 1);
      return;
    }
  }

  void packByte_(byte letter)
  {
    for (int i = 7; i >= 0; i--)
      packedtree.push_back((letter >> i) & 1);
  }

  void getPackedTreeSize_(vector<bool> &buffer)
  {
    int packedsize = packedtree.size();
    for (int i = 31; i >= 0; i--)
      buffer.push_back((packedsize >> i) & 1);
  }

  void packTree_(HNode *&hnode)
  {
    if (!hnode)
      return;
    packTree_(hnode->GetLeft());
    packTree_(hnode->GetRight());
    if (!hnode->IsLeaf())
      packedtree.push_back(true);
    else
      packedtree.push_back(false), packByte_(hnode->GetLetter());
  }
};


class HuffmanZipper
{
private:
  HuffmanTree *htree;
  map<byte, vector<bool> > encodetable;
public:
  HuffmanZipper(const vector<byte> &inputstreamdata)
  {
    htree = new HuffmanTree(inputstreamdata);
    vector<bool> path;
    htree->CompleteCodeTable(encodetable, path);
    if (!path.empty())
      path.clear();
  }

  ~HuffmanZipper()
  {
    delete htree;
    if (!encodetable.empty())
      encodetable.clear();
  }

  void GetCode(vector<bool> &buffer, const byte &letter)
  {
    buffer.clear();
    buffer = vector<bool>(encodetable[letter].begin(), encodetable[letter].end());
  }

  void GetByte(const vector<bool> &code, byte &letter)
  {
    this->htree->GetByte(code, letter);
  }

  void GetPackedTree(vector<bool> &buffer)
  {
    this->htree->GetPackedTree(buffer);
  }

  void PutCompressedTree(vector<byte> &compresseddata)
  {
    vector<bool> buffer;
    htree->GetPackedTree(buffer);
    byte outputvalue = 0;
    for (int i = 0; i < buffer.size(); i++)
    {
      outputvalue |= buffer[i] << (7 - (i % 8));
      if (i % 8 == 7)
      {
        compresseddata.push_back(outputvalue);
        outputvalue = 0;
      }
    }
    if (outputvalue != 0)
      compresseddata.push_back(outputvalue);
    buffer.clear();
  }

  void PutCompressedText(vector<byte> &compresseddata)
  {
    unsigned char bitindex = 0;
    vector<bool> buffer;
    byte outputvalue = 0;
    int unusedcount = 0;
    for (int i = 0; i < 8; i++)
      if (!(compresseddata[compresseddata.size() - 1] & (1 << i)))
        unusedcount++;
      else
        break;
    int unusedCopy = unusedcount;
    int treeSize = compresseddata.size();
    compresseddata.push_back(0);
    for (int i = 0; i < htree->dataforzip.size(); i++)
    {
      GetCode(buffer, htree->dataforzip[i]);
      for (int j = 0; j < buffer.size(); j++)
      {
        outputvalue = (outputvalue << 1) | buffer[j];
        bitindex++;
        if (unusedcount != 0 && bitindex == unusedcount)
        {
          bitindex = 0;
          for (int z = 0; z < unusedcount; z++)
            compresseddata[compresseddata.size() - 1] |= outputvalue & (1 << z);
          outputvalue = 0;
          unusedcount = 0;
        }
        if (bitindex == 8)
        {
          bitindex = 0;
          compresseddata.push_back(outputvalue);
          outputvalue = 0;
        }
      }
      buffer.clear();
    }
    compresseddata[treeSize - 1] |= (8 - bitindex) >> (8 - unusedCopy);
    compresseddata[treeSize] |= (8 - bitindex) << unusedCopy;
    if (bitindex != 0)
      compresseddata.push_back(outputvalue << (8 - bitindex));
  }
};


void Encode(IInputStream &original, IOutputStream &compressed)
{
  vector<byte> inputstreamdata;
  byte inputvalue = 0;
  while (original.Read(inputvalue))
    inputstreamdata.push_back(inputvalue);
  HuffmanZipper *hzipper = new HuffmanZipper(inputstreamdata);
  inputstreamdata.clear();
  vector<byte> outputstreamdata;
  hzipper->PutCompressedTree(outputstreamdata);
  hzipper->PutCompressedText(outputstreamdata);
  for (int i = 0; i < outputstreamdata.size(); i++)
    compressed.Write(outputstreamdata[i]);
  delete hzipper;
  outputstreamdata.clear();
}


void Decode(IInputStream &compressed, IOutputStream &original)
{
  unsigned int treesize = 0;
  byte inputvalue = 0;
  for (int i = 0; i < 4; i++)
  {
    compressed.Read(inputvalue);
    treesize <<= 8;
    treesize |= (unsigned char) inputvalue;
  }
  byte tmp = 0;
  string s;
  stack<HNode *> nodes;
  vector<bool> bynaryString;
  while (compressed.Read(tmp))
    s.push_back(tmp);
  for (int i = 0; i < s.size(); i++)
  {
    for (int j = 7; j >= 0; j--)
    {
      bynaryString.push_back((s[i] >> j) & 1);
    }
  }
  int scannedBits = 0;
  int c = 0;
  for (int i = 0; i < treesize; i++)
  {
    if (scannedBits != 9)
    {
      if (bynaryString[i] && scannedBits == 0)
      {
        HNode *right = nodes.top();
        nodes.pop();
        HNode *left = nodes.top();
        nodes.pop();
        HNode *node = new HNode(0, 0, false, left, right);
        nodes.push(node);
      }
      else
      {
        if (scannedBits != 0)
        {
          int D = (bynaryString[i] ? 1 : 0);
          int U = 8 - scannedBits;
          c |= D << U;
        }
        scannedBits++;
      }
    }
    else
    {
      scannedBits = 0;
      nodes.push(new HNode(c, 0, true, nullptr, nullptr));
      c = 0;
      i--;
    }
  }
  int b = nodes.size();
  HNode *treeRoot = nodes.top();
  byte fakeBits = 0;
  for (int i = treesize; i < treesize + 8; i++)
  {
    fakeBits |= bynaryString[i] << (7 - (i - treesize));
  }
  HNode *candidate = treeRoot;
  for (int i = treesize + 8; i < bynaryString.size() - fakeBits; i++)
  {
    if (!bynaryString[i])
    {
      candidate = candidate->GetLeft();
    }
    else
    {
      candidate = candidate->GetRight();
    }
    if (candidate && candidate->IsLeaf())
    {
      original.Write(candidate->GetLetter());
      candidate = treeRoot;
    }
  }
}