// zdeflate.cpp - written and placed in the public domain by Wei Dai

// Many of the algorithms and tables used here came from the deflate implementation
// by Jean-loup Gailly, which was included in Crypto++ 4.0 and earlier. I completely
// rewrote it in order to fix a bug that I could not figure out. This code
// is less clever, but hopefully more understandable and maintainable.

#include "pch.h"
#include "zdeflate.h"
#include <functional>

NAMESPACE_BEGIN(CryptoPP)

using namespace std;

LowFirstBitWriter::LowFirstBitWriter(BufferedTransformation *attachment)
	: Filter(attachment), m_counting(false), m_buffer(0), m_bitsBuffered(0), m_bytesBuffered(0)
{
}

void LowFirstBitWriter::StartCounting()
{
	assert(!m_counting);
	m_counting = true;
	m_bitCount = 0;
}

unsigned long LowFirstBitWriter::FinishCounting()
{
	assert(m_counting);
	m_counting = false;
	return m_bitCount;
}

void LowFirstBitWriter::PutBits(unsigned long value, unsigned int length)
{
	if (m_counting)
		m_bitCount += length;
	else
	{
		m_buffer |= value << m_bitsBuffered;
		m_bitsBuffered += length;
		assert(m_bitsBuffered <= sizeof(unsigned long)*8);
		while (m_bitsBuffered >= 8)
		{
			m_outputBuffer[m_bytesBuffered++] = (byte)m_buffer;
			if (m_bytesBuffered == m_outputBuffer.size())
			{
				AttachedTransformation()->PutModifiable(m_outputBuffer, m_bytesBuffered);
				m_bytesBuffered = 0;
			}
			m_buffer >>= 8;
			m_bitsBuffered -= 8;
		}
	}
}

void LowFirstBitWriter::FlushBitBuffer()
{
	if (m_counting)
		m_bitCount += 8*(m_bitsBuffered > 0);
	else
	{
		if (m_bytesBuffered > 0)
		{
			AttachedTransformation()->PutModifiable(m_outputBuffer, m_bytesBuffered);
			m_bytesBuffered = 0;
		}
		if (m_bitsBuffered > 0)
		{
			AttachedTransformation()->Put((byte)m_buffer);
			m_buffer = 0;
			m_bitsBuffered = 0;
		}
	}
}

void LowFirstBitWriter::ClearBitBuffer()
{
	m_buffer = 0;
	m_bytesBuffered = 0;
	m_bitsBuffered = 0;
}

HuffmanEncoder::HuffmanEncoder(const unsigned int *codeBits, unsigned int nCodes)
{
	Initialize(codeBits, nCodes);
}

struct HuffmanNode
{
	unsigned int symbol;
	union {unsigned int parent, depth, freq;};
};

struct FreqLessThan
{
	inline bool operator()(unsigned int lhs, const HuffmanNode &rhs) {return lhs < rhs.freq;}
	inline bool operator()(const HuffmanNode &lhs, const HuffmanNode &rhs) const {return lhs.freq < rhs.freq;}
};

void HuffmanEncoder::GenerateCodeLengths(unsigned int *codeBits, unsigned int maxCodeBits, const unsigned int *codeCounts, unsigned int nCodes)
{
	assert(nCodes > 0);
	assert(nCodes <= (unsigned int)(1 << maxCodeBits));

	unsigned int i;
	SecBlockWithHint<HuffmanNode, 2*286> tree(nCodes);
	for (i=0; i<nCodes; i++)
	{
		tree[i].symbol = i;
		tree[i].freq = codeCounts[i];
	}
	sort(tree.begin(), tree.end(), FreqLessThan());
	unsigned int treeBegin = upper_bound(tree.begin(), tree.end(), 0, FreqLessThan()) - tree.begin();
	if (treeBegin == nCodes)
	{	// special case for no codes
		fill(codeBits, codeBits+nCodes, 0);
		return;
	}
	tree.resize(nCodes + nCodes - treeBegin - 1);

	unsigned int leastLeaf = treeBegin, leastInterior = nCodes;
	for (i=nCodes; i<tree.size(); i++)
	{
		unsigned int least;
		least = (leastLeaf == nCodes || (leastInterior < i && tree[leastInterior].freq < tree[leastLeaf].freq)) ? leastInterior++ : leastLeaf++;
		tree[i].freq = tree[least].freq;
		tree[least].parent = i;
		least = (leastLeaf == nCodes || (leastInterior < i && tree[leastInterior].freq < tree[leastLeaf].freq)) ? leastInterior++ : leastLeaf++;
		tree[i].freq += tree[least].freq;
		tree[least].parent = i;
	}

	tree[tree.size()-1].depth = 0;
	if (tree.size() >= 2)
		for (i=tree.size()-2; i>=nCodes; i--)
			tree[i].depth = tree[tree[i].parent].depth + 1;
	unsigned int sum = 0;
	SecBlockWithHint<unsigned int, 15+1> blCount(maxCodeBits+1);
	fill(blCount.begin(), blCount.end(), 0);
	for (i=treeBegin; i<nCodes; i++)
	{
		unsigned int depth = STDMIN(maxCodeBits, tree[tree[i].parent].depth + 1);
		blCount[depth]++;
		sum += 1 << (maxCodeBits - depth);
	}

	unsigned int overflow = sum > (unsigned int)(1 << maxCodeBits) ? sum - (1 << maxCodeBits) : 0;

	while (overflow--)
	{
		unsigned int bits = maxCodeBits-1;
		while (blCount[bits] == 0)
			bits--;
		blCount[bits]--;
		blCount[bits+1] += 2;
		assert(blCount[maxCodeBits] > 0);
		blCount[maxCodeBits]--;
	}

	for (i=0; i<treeBegin; i++)
		codeBits[tree[i].symbol] = 0;
	unsigned int bits = maxCodeBits;
	for (i=treeBegin; i<nCodes; i++)
	{
		while (blCount[bits] == 0)
			bits--;
		codeBits[tree[i].symbol] = bits;
		blCount[bits]--;
	}
	assert(blCount[bits] == 0);
}

void HuffmanEncoder::Initialize(const unsigned int *codeBits, unsigned int nCodes)
{
	assert(nCodes > 0);
	unsigned int maxCodeBits = *max_element(codeBits, codeBits+nCodes);
	if (maxCodeBits == 0)
		return;		// assume this object won't be used

	SecBlockWithHint<unsigned int, 15+1> blCount(maxCodeBits+1);
	fill(blCount.begin(), blCount.end(), 0);
	unsigned int i;
	for (i=0; i<nCodes; i++)
		blCount[codeBits[i]]++;

	code_t code = 0;
	SecBlockWithHint<code_t, 15+1> nextCode(maxCodeBits+1);
	nextCode[1] = 0;
	for (i=2; i<=maxCodeBits; i++)
	{
		code = (code + blCount[i-1]) << 1;
		nextCode[i] = code;
	}
	assert(maxCodeBits == 1 || code == (1 << maxCodeBits) - blCount[maxCodeBits]);

	m_valueToCode.resize(nCodes);
	for (i=0; i<nCodes; i++)
	{
		unsigned int len = m_valueToCode[i].len = codeBits[i];
		if (len != 0)
			m_valueToCode[i].code = BitReverse(nextCode[len]++) >> (8*sizeof(code_t)-len);
	}
}

inline void HuffmanEncoder::Encode(LowFirstBitWriter &writer, value_t value) const
{
	assert(m_valueToCode[value].len > 0);
	writer.PutBits(m_valueToCode[value].code, m_valueToCode[value].len);
}

Deflator::Deflator(BufferedTransformation *attachment, int deflateLevel, int log2WindowSize)
	: LowFirstBitWriter(attachment)
{
	InitializeStaticEncoders();
	IsolatedInitialize(MakeParameters("DeflateLevel", deflateLevel)("Log2WindowSize", log2WindowSize));
}

Deflator::Deflator(const NameValuePairs &parameters, BufferedTransformation *attachment)
	: LowFirstBitWriter(attachment)
{
	InitializeStaticEncoders();
	IsolatedInitialize(parameters);
}

void Deflator::InitializeStaticEncoders()
{
	unsigned int codeLengths[288];
	fill(codeLengths + 0, codeLengths + 144, 8);
	fill(codeLengths + 144, codeLengths + 256, 9);
	fill(codeLengths + 256, codeLengths + 280, 7);
	fill(codeLengths + 280, codeLengths + 288, 8);
	m_staticLiteralEncoder.Initialize(codeLengths, 288);
	fill(codeLengths + 0, codeLengths + 32, 5);
	m_staticDistanceEncoder.Initialize(codeLengths, 32);
}

void Deflator::IsolatedInitialize(const NameValuePairs &parameters)
{
	int log2WindowSize = parameters.GetIntValueWithDefault("Log2WindowSize", DEFAULT_LOG2_WINDOW_SIZE);

	if (!(MIN_LOG2_WINDOW_SIZE <= log2WindowSize && log2WindowSize <= MAX_LOG2_WINDOW_SIZE))
		throw InvalidArgument("Deflator: " + IntToString(log2WindowSize) + " is an invalid window size");

	m_log2WindowSize = log2WindowSize;
	DSIZE = 1 << m_log2WindowSize;
	DMASK = DSIZE - 1;
	HSIZE = 1 << m_log2WindowSize;
	HMASK = HSIZE - 1;
	m_byteBuffer.New(2*DSIZE);
	m_head.New(HSIZE);
	m_prev.New(DSIZE);
	m_matchBuffer.New(DSIZE/2);

	SetDeflateLevel(parameters.GetIntValueWithDefault("DeflateLevel", DEFAULT_DEFLATE_LEVEL));
	Reset(true);
}

void Deflator::Reset(bool forceReset)
{
	if (forceReset)
		ClearBitBuffer();
	else
		assert(m_bitsBuffered == 0);

	m_headerWritten = false;
	m_matchAvailable = false;
	m_dictionaryEnd = 0;
	m_stringStart = 0;
	m_lookahead = 0;
	m_minLookahead = MAX_MATCH;
	m_previousMatch = 0;
	m_previousLength = 0;
	m_matchBufferEnd = 0;
	m_blockStart = 0;
	m_blockLength = 0;

	// m_prev will be initialized automaticly in InsertString
	fill(m_head.begin(), m_head.end(), 0);

	fill(m_literalCounts.begin(), m_literalCounts.end(), 0);
	fill(m_distanceCounts.begin(), m_distanceCounts.end(), 0);
}

void Deflator::SetDeflateLevel(int deflateLevel)
{
	if (!(MIN_DEFLATE_LEVEL <= deflateLevel && deflateLevel <= MAX_DEFLATE_LEVEL))
		throw InvalidArgument("Deflator: " + IntToString(deflateLevel) + " is an invalid deflate level");

	unsigned int configurationTable[10][4] = {
		/*      good lazy nice chain */
		/* 0 */ {0,    0,  0,    0},  /* store only */
		/* 1 */ {4,    3,  8,    4},  /* maximum speed, no lazy matches */
		/* 2 */ {4,    3, 16,    8},
		/* 3 */ {4,    3, 32,   32},
		/* 4 */ {4,    4, 16,   16},  /* lazy matches */
		/* 5 */ {8,   16, 32,   32},
		/* 6 */ {8,   16, 128, 128},
		/* 7 */ {8,   32, 128, 256},
		/* 8 */ {32, 128, 258, 1024},
		/* 9 */ {32, 258, 258, 4096}}; /* maximum compression */

	GOOD_MATCH = configurationTable[deflateLevel][0];
	MAX_LAZYLENGTH = configurationTable[deflateLevel][1];
	MAX_CHAIN_LENGTH = configurationTable[deflateLevel][3];

	m_deflateLevel = deflateLevel;
}

unsigned int Deflator::FillWindow(const byte *str, unsigned int length)
{
	unsigned int accepted = STDMIN(length, 2*DSIZE-(m_stringStart+m_lookahead));

	if (m_stringStart >= 2*DSIZE - MAX_MATCH)
	{
		if (m_blockStart < DSIZE)
			EndBlock(false);

		memcpy(m_byteBuffer, m_byteBuffer + DSIZE, DSIZE);

		m_dictionaryEnd = m_dictionaryEnd < DSIZE ? 0 : m_dictionaryEnd-DSIZE;
		assert(m_stringStart >= DSIZE);
		m_stringStart -= DSIZE;
		assert(m_previousMatch >= DSIZE || m_previousLength < MIN_MATCH);
		m_previousMatch -= DSIZE;
		assert(m_blockStart >= DSIZE);
		m_blockStart -= DSIZE;

		unsigned int i;

		for (i=0; i<HSIZE; i++)
			m_head[i] = SaturatingSubtract(m_head[i], DSIZE);

		for (i=0; i<DSIZE; i++)
			m_prev[i] = SaturatingSubtract(m_prev[i], DSIZE);

		accepted = STDMIN(accepted + DSIZE, length);
	}
	assert(accepted > 0);

	memcpy(m_byteBuffer + m_stringStart + m_lookahead, str, accepted);
	m_lookahead += accepted;
	return accepted;
}

inline unsigned int Deflator::ComputeHash(const byte *str) const
{
	assert(str+3 <= m_byteBuffer + m_stringStart + m_lookahead);
	return ((str[0] << 10) ^ (str[1] << 5) ^ str[2]) & HMASK;
}

unsigned int Deflator::LongestMatch(unsigned int &bestMatch) const
{
	assert(m_previousLength < MAX_MATCH);

	bestMatch = 0;
	unsigned int bestLength = STDMAX(m_previousLength, (unsigned int)MIN_MATCH-1);
	if (m_lookahead <= bestLength)
		return 0;

	const byte *scan = m_byteBuffer + m_stringStart, *scanEnd = scan + STDMIN((unsigned int)MAX_MATCH, m_lookahead);
	unsigned int limit = m_stringStart > (DSIZE-MAX_MATCH) ? m_stringStart - (DSIZE-MAX_MATCH) : 0;
	unsigned int current = m_head[ComputeHash(scan)];

	unsigned int chainLength = MAX_CHAIN_LENGTH;
	if (m_previousLength >= GOOD_MATCH)
		chainLength >>= 2;