//////////////////////////////////////////////////////////////////////////////
//
// Asura_Huffman.cpp
//
// Copyright (c) 2005 Rebellion. All rights reserved.
//
//////////////////////////////////////////////////////////////////////////////

// I N C L U D E S ///////////////////////////////////////////////////////////

#include "Core/Asura.h"

// This file's header.
#include "ChunkSystem/Asura_Chunk_Stream.h"
#include "ChunkSystem/Asura_Huffman.h"

// Engine independent.
#ifdef ASURA_TOOLS
#include "Core/Asura_Endian.h"
#endif
#include "Core/Asura_Maths.h"
#include "Core/Memory/Asura_Mem.h"
#include "FileAccess/Asura_FileAccess.h"

#include <stdlib.h>
#include <string.h>

#ifdef __SNC__
#pragma diag_suppress=129 // loop is not reachable from preceding code
#endif


// T Y P E S /////////////////////////////////////////////////////////////////

struct HuffItem
{
    int   Symbol;
    u_int Count;
};

struct HuffNode
{
    HuffNode* zero;   // points to the "zero" branch or...
    HuffNode* one;    // points to the "one" branch or...
    HuffNode* parent; // the THIRD four bytes, parent node
    u_int     bits;   // the bit pattern of this end node
};

struct HuffEncode
{
    int wid;
    int bits;
};

// C O N S T A N T S /////////////////////////////////////////////////////////

static const u_int uEDXMASK = ((((1u << (ASURA_HUFFMAN_MAX_DEPTH + 1)) - 1) ^ 1) ^ -1);

// P R O T O T Y P E S   A N D   I N L I N E   F U N C T I O N S /////////////

// Compression helper functions.
static void Encode_PerformSymbolCensus( HuffItem* pxCensus, const u_char* pucSrc, u_int uNumBytes );
static int Encode_SymbolSortCallback( const void* pvCmp1, const void* pvCmp2 );
static void Encode_MakeTreeFromCensus( HuffNode* pxTreeNodes, const HuffItem* pxSymbolCensus );
static void Encode_MakeCodeLengthsFromHuffTree( int* piDepths, const HuffNode* pxSymbolCensus );
static void Encode_AdjustDepths( int* piDepths );
static void Encode_MakeEncodeTable( HuffEncode* pxEncodingTable, const HuffItem* item, const int* piDepths );
static u_int Encode_EncodeBytes( int* piDest, const u_char* pucSrc, u_int uNumBytes, const HuffEncode* pxEncodingTable );

// Decompression helper functions.
static void Decode_MakeDecodeTable( const int* piDepths, const u_char* pucByteAssignments, int* piOut );

// C L A S S   I M P L E M E N T A T I O N S /////////////////////////////////

//
// Asura_Huffman_Encoder.
//

bool Asura_Huffman_Encoder::AllocateAndCompressForASR( const void* pvSrc, u_int uNumBytes, Asura_Huffman_CompressedFileHdr** ppxHdr, u_int* puTotalSize, bool bSwapEndianness )
{
    const u_int uMaxCompressedDataSize  = sizeof(Asura_Huffman_CompressedFileHdr) + uNumBytes;

    void* pvBfr = Asura_Allocate( uMaxCompressedDataSize );
    if( !pvBfr )
    {
        return false;
    }

    Asura_Chunk_Memory_Stream xStream( uMaxCompressedDataSize, pvBfr );

#ifdef ASURA_TOOLS
    xStream.SetSwapEndian( bSwapEndianness );
#else
    Asura_Assert( !bSwapEndianness, "Not supported outside tools builds." );
#endif

	xStream.Write( ASURA_COMPRESSED_FILE_ID_ACTUAL_CHARS, ASURA_COMPRESSED_FILE_ID );

    if( CompressToBfr(pvSrc, uNumBytes, puTotalSize, xStream) )
    {
        *ppxHdr = static_cast<Asura_Huffman_CompressedFileHdr*>( pvBfr );
        return true;
    }
    else
    {
        Asura_Deallocate( pvBfr );
        return false;
    }
}

bool Asura_Huffman_Encoder::CompressToBfr( const void* pvSrc, u_int uNumBytes, u_int* puTotalPackageSize, Asura_Chunk_Memory_Stream& xStream )
{
    HuffItem*   pxSymbolCensus  = new HuffItem[257];
    HuffNode*   pxTreeNodes     = new HuffNode[2 * 257];
    int*        piDepths        = new int[ASURA_HUFFMAN_MAX_DEPTH + 1];
    HuffEncode* pxEncodingTable = new HuffEncode[257];

    bool bOK = ( pxSymbolCensus && pxTreeNodes && piDepths && pxEncodingTable );

    if( bOK )
    {
        const u_char* pucSrcBytes = static_cast<const u_char*>( pvSrc );

        Encode_PerformSymbolCensus( pxSymbolCensus, pucSrcBytes, uNumBytes );
        qsort( pxSymbolCensus, 257, sizeof(HuffItem), Encode_SymbolSortCallback );

        Encode_MakeTreeFromCensus( pxTreeNodes, pxSymbolCensus );

        Encode_MakeCodeLengthsFromHuffTree( piDepths, pxTreeNodes );   
        Encode_AdjustDepths( piDepths );

        Encode_MakeEncodeTable( pxEncodingTable, &pxSymbolCensus[256], piDepths );

		Asura_Huffman_Package*	pxPackage		= static_cast<Asura_Huffman_Package*>(xStream.GetAddress());
		int*					piOutput		= reinterpret_cast<int*>(pxPackage + 1);
		u_int					uCompressedSize	= Encode_EncodeBytes (piOutput, pucSrcBytes, uNumBytes, pxEncodingTable);

		xStream << uCompressedSize;
		xStream << uNumBytes;

        for( int i = 0; i < ASURA_HUFFMAN_MAX_DEPTH; ++i )
        {
            xStream << piDepths [i + 1];
        }

        for( int i = 0; i < 256; ++i )
        {
            xStream << static_cast<u_char>(pxSymbolCensus [i + 1].Symbol);
        }

#ifdef ASURA_TOOLS
		if (xStream.ShouldSwapEndian())
		{
			for (u_int u = 0; u < uCompressedSize; u += sizeof(int), piOutput++)
			{
				*piOutput	= Asura_Endian::SwapBytes (*piOutput);
			}
		}
#endif

        if( puTotalPackageSize )
        {
            *puTotalPackageSize = sizeof(Asura_Huffman_CompressedFileHdr) + uCompressedSize;
        }
    }

    delete[] pxEncodingTable;
    delete[] piDepths;   
    delete[] pxTreeNodes;
    delete[] pxSymbolCensus;

    return bOK;
}

//
// Asura_Huffman_Decoder.
//

u_int Asura_Huffman_Decoder::AllocateAndDecompressAllFromASR( const Asura_Huffman_CompressedFileHdr* pxCompFileHdr, void** ppvBfr )
{
    void* pvBfr = Asura_Allocate( pxCompFileHdr->m_xHuffmanPackage.m_uUncompressedDataSize );
    *ppvBfr = pvBfr;
    return( pvBfr ? DecompressAll(&pxCompFileHdr->m_xHuffmanPackage, pvBfr) : 0 );
}

u_int Asura_Huffman_Decoder::DecompressAll( const Asura_Huffman_Package* pxPackage, void* pvDest )
{
    Asura_Huffman_Decoder xDecoder;
    xDecoder.Begin( pxPackage );
    u_int uNumBytes = xDecoder.Decompress( pvDest, xDecoder.GetDecompressedSize() );
    xDecoder.End();
    return uNumBytes;
}

Asura_Huffman_Decoder::Asura_Huffman_Decoder( void )
{
#ifdef ASURA_PLATFORM_CAN_SUPPLY_FILE_DECOMPRESSOR_MEM
    m_piTable = NULL;
#else
    m_piTable = new int[ASURA_HUFFMAN_TABLE_SIZE];
#endif
}

Asura_Huffman_Decoder::~Asura_Huffman_Decoder( void )
{
#ifndef ASURA_PLATFORM_CAN_SUPPLY_FILE_DECOMPRESSOR_MEM
    delete[] m_piTable;
#endif
}

void Asura_Huffman_Decoder::Begin( Asura_Handle iFileHandle )
{
    m_puSrcData   = NULL;
    m_iFileHandle = iFileHandle;

    Asura_FileAccess::Platform_ReadData( m_iFileHandle, &m_xPackage, sizeof(Asura_Huffman_Package) );

    SetupCommonState();
}

void Asura_Huffman_Decoder::Begin( const Asura_Huffman_Package* pxPackage )
{
    m_puSrcData   = reinterpret_cast<const u_int*>( pxPackage + 1 );
    m_iFileHandle = ASURA_INVALID_HANDLE;

    Asura_Mem::Copy( &m_xPackage, pxPackage, sizeof(Asura_Huffman_Package) );

    SetupCommonState();
}

void Asura_Huffman_Decoder::SetupCommonState( void )
{
#ifdef ASURA_PLATFORM_CAN_SUPPLY_FILE_DECOMPRESSOR_MEM
    Platform_AllocTableBfr();
#endif

	#ifdef ASURA_TOOLS
	m_bSwapEndian = false;
	//Determine if the compression is endian swapped
	//Can only do this when reading from a file , since we need access to the file size
	//for the endian detecting mechanism (TM)
	//RAR 22/08/2007
	if(m_iFileHandle != ASURA_INVALID_HANDLE)
	{
		u_int uFileSize = Asura_FileAccess::GetFileSize(m_iFileHandle);
		u_int uExpectedFileSize = sizeof(Asura_Huffman_CompressedFileHdr) + m_xPackage.m_uCompressedDataSize;

		if(uFileSize != uExpectedFileSize)
		{
			//The file size doesn't add up , so see if it does when endian swapped
			uExpectedFileSize = sizeof(Asura_Huffman_CompressedFileHdr) + Asura_Endian::SwapBytes(m_xPackage.m_uCompressedDataSize);
			Asura_Assert(uFileSize == uExpectedFileSize,"Can't determine whether compressian is endian swapped , since the sizes are wrong either way");

			if(uFileSize == uExpectedFileSize)
			{
				m_bSwapEndian = true;
			}
		}
	}

	if(m_bSwapEndian)
	{
		//need to swap everything but the byte assignments
		m_xPackage.m_uCompressedDataSize = Asura_Endian::SwapBytes(m_xPackage.m_uCompressedDataSize);
		m_xPackage.m_uUncompressedDataSize = Asura_Endian::SwapBytes(m_xPackage.m_uUncompressedDataSize);
		for(int i=0; i<ASURA_HUFFMAN_MAX_DEPTH; i++)
		{
			m_xPackage.m_aiCodelengthCount[i] = Asura_Endian::SwapBytes(m_xPackage.m_aiCodelengthCount[i]);
		}
	}
	#endif //ASURA_TOOLS

    Decode_MakeDecodeTable( m_xPackage.m_aiCodelengthCount, m_xPackage.m_aucByteAssignment, m_piTable );

    m_iAvailable    = 0;
    m_iReserve      = 0;
    m_iWid          = 0;
    m_uBits         = 0;
    m_uResBits      = 0;
    m_iFill         = 0;
    m_uBytesRead    = 0;
    m_bJumpIntoLoop = false;
}

void Asura_Huffman_Decoder::End( void )
{
#ifdef ASURA_PLATFORM_CAN_SUPPLY_FILE_DECOMPRESSOR_MEM
    Platform_DeallocTableBfr();
#endif
}

u_int Asura_Huffman_Decoder::Decompress( void* pvDest, u_int uNumBytes )
{
    u_char* pucDest = static_cast<u_char*>( pvDest );
    u_char* const pucDestStart = pucDest;

    const short*  psTableStart = reinterpret_cast<const short*>( m_piTable );
    const u_char* p;

    if( m_bJumpIntoLoop ) { goto lpent; }

    m_bJumpIntoLoop = true;

    do 
    {
        m_iAvailable  += m_iWid;
        m_iFill        = 31 - m_iAvailable;                             
        m_uBits      <<= m_iFill;

        if( m_iFill > m_iReserve )
        {
            m_iFill      -= m_iReserve;
            m_iAvailable += m_iReserve;
            if( m_iReserve )
            {
                m_uBits = (m_uBits >> m_iReserve) | (m_uResBits << (32 - m_iReserve));
            }

            //Avoid reading out of range from the source buffer.
            //This happens on some occasions without this check,but in that case the bits read in
            //appear to be ignored anyway. RAR 03/04/03
            if( m_uBytesRead < m_xPackage.m_uCompressedDataSize )
            {
                if( m_puSrcData )
                {
                    m_uResBits = *m_puSrcData++;
                }
                else
                {
                    Asura_FileAccess::Platform_ReadData( m_iFileHandle, &m_uResBits, 4 );
                }
				#ifdef ASURA_TOOLS
				if(m_bSwapEndian)
				{
					m_uResBits = Asura_Endian::SwapBytes(m_uResBits);
				}
				#endif
            }
            else
            {
                Asura_Assert( m_uBytesRead == m_xPackage.m_uCompressedDataSize, "" );
                m_uResBits = 0;
            }
            m_uBytesRead += 4;
            m_iReserve    = 32;
        }

        m_uBits        = (m_uBits >> m_iFill) | (m_uResBits << (32 - m_iFill));
        m_uResBits   >>= m_iFill;
        m_iReserve    -= m_iFill;
        m_iAvailable   = 31;
        goto lpent;
        
        do
        {
            m_uBits    >>= m_iWid;
            *pucDest++   = p[1];

            // Done yet?
            {
                u_int uBytesDecoded = pucDest - pucDestStart;
                if( uBytesDecoded == uNumBytes )
                {
                    return uBytesDecoded;
                }
            }

lpent:      p = reinterpret_cast<const u_char*>( psTableStart + (m_uBits & ~uEDXMASK) );
        }
        while( (m_iAvailable -= (m_iWid = *p)) >= 0 );
    }
    while( m_iAvailable > -32 );
    
    return pucDest - pucDestStart;
}

// F U N C T I O N   I M P L E M E N T A T I O N S ///////////////////////////

static void Encode_PerformSymbolCensus( HuffItem* pxCensus, const u_char* pucSrc, u_int uNumBytes )
{
    // init array
    for( int i = 0; i < 257; ++i )
    {
        pxCensus[i].Symbol = i;
        pxCensus[i].Count  = 0;
    }

    // count 'em
    for( ; uNumBytes; --uNumBytes )
    {
        pxCensus[*pucSrc++].Count++;
    }
}

static int Encode_SymbolSortCallback( const void* pvCmp1, const void* pvCmp2 )
{
    const u_int uCount1 = (static_cast<const HuffItem*>(pvCmp1))->Count;
    const u_int uCount2 = (static_cast<const HuffItem*>(pvCmp2))->Count;

    if( uCount1 > uCount2 )
    {
        return  1;
    }
    else if( uCount1 < uCount2 )
    {
        return -1;
    }
    else
    {
        return 0;
    }
}

static void Encode_MakeTreeFromCensus( HuffNode* pxTreeNodes, const HuffItem* pxSymbolCensus )
{
    HuffNode *movdest, *temp;
    int upperlim, lowerlim, index;
    unsigned int sum;

    movdest = pxTreeNodes + 1;
    temp    = pxTreeNodes + 257;

    for( int i = 0; i < 257; ++i )
    {
        temp[i].zero   = NULL;
        temp[i].one    = NULL;
        temp[i].parent = NULL;
        temp[i].bits   = pxSymbolCensus[i].Count;
    }

    u_int uCount = 257;
    while ((upperlim = --uCount)!=0)
    {
        if (temp[0].zero)
            temp[0].zero->parent = temp[0].one->parent = movdest;
        if (temp[1].zero)
            temp[1].zero->parent = temp[1].one->parent = movdest + 1;
        movdest[0]   = *temp++;
        movdest[1]   = *temp;
        sum = movdest[0].bits + movdest[1].bits;
        lowerlim = 1;
        
        while (lowerlim != upperlim)
        {
            index = (lowerlim + upperlim) >> 1;
            if (sum >= temp[index].bits)
            {
                lowerlim = index + 1;
            }
            else
            {
                upperlim = index;
            }
        }
        index            = lowerlim - 1;
        memmove(temp, temp + 1, index * sizeof(HuffNode));
        temp[index].bits = sum;
        temp[index].zero = movdest;
        temp[index].one  = movdest + 1;
        movdest         += 2;
    }
    
    pxTreeNodes[0] = temp[0];
    if( pxTreeNodes[0].zero )
    {
        pxTreeNodes[0].zero->parent = pxTreeNodes[0].one->parent = pxTreeNodes;
    }
}

static void Encode_MakeCodeLengthsFromHuffTree( int* piDepths, const HuffNode* pxSymbolCensus )
{
    memset( piDepths, 0, (ASURA_HUFFMAN_MAX_DEPTH + 1) * sizeof(int) );

    const HuffNode* pxCurNode = pxSymbolCensus;
    const HuffNode* pxBack;
    int iDepth = 0;

    while( true )
    {
        while( pxCurNode->one )
        {
            pxCurNode = pxCurNode->one;
            iDepth++;
        }
        
        piDepths[Asura_Maths::Min(iDepth, ASURA_HUFFMAN_MAX_DEPTH)]++;

        do
        {
            pxBack = pxCurNode;
            pxCurNode = pxCurNode->parent;
            if( !iDepth-- )
                return;
        }
        while( pxBack == pxCurNode->zero );

        pxCurNode = pxCurNode->zero;
	    iDepth++;
    }
}

static void Encode_AdjustDepths( int* piDepths )
{
    u_int items, sum, goal, gain;
    u_int excess, hi;
	int m, n, busts,promotions;

    goal = 1 << ASURA_HUFFMAN_MAX_DEPTH;
    for (n = 0, sum = 0, items = 0; n <= ASURA_HUFFMAN_MAX_DEPTH; n++)
    {
        items += piDepths[n];
        sum   += (goal >> n) * piDepths[n];
    }

    // (MP - 02.03.2005) The old code ignored this function's return value
    // anyway so let's just assert about it here instead.
//  if (items > goal)
//      return -1;                              // failure
    Asura_Assert( items <= goal, "" );

    for (n = ASURA_HUFFMAN_MAX_DEPTH - 1; sum > goal; n--)
    {
        if( piDepths[n] )
        {
            gain                               = (1 << (ASURA_HUFFMAN_MAX_DEPTH - n)) - 1;
            busts                              = (sum - goal + gain - 1) / gain;
            busts                              = piDepths[n] < busts ? piDepths[n] : busts;
            piDepths[n]                       -= busts;
            piDepths[ASURA_HUFFMAN_MAX_DEPTH] += busts;
            sum                               -= busts * gain;
        }
    }

    excess = goal - sum;
    for (n = 0; excess; n++)
    {
        hi = 1 << (ASURA_HUFFMAN_MAX_DEPTH - n);
        for (m = n + 1; m <= ASURA_HUFFMAN_MAX_DEPTH; m++)
        {
            gain = hi - (1 << (ASURA_HUFFMAN_MAX_DEPTH - m));
            if (excess < gain)
                break;
            if( piDepths[m] )
            {
                promotions   = excess / gain;
                promotions   = piDepths[m] > promotions ? promotions : piDepths[m];
                piDepths[n] += promotions;
                piDepths[m] -= promotions;
                excess      -= promotions * gain;
            }
        }
    }

//  return 0;                           // success
}

static void Encode_MakeEncodeTable( HuffEncode* pxEncodingTable, const HuffItem* item, const int* piDepths )
{
    u_int d, bitwidth, depthbit, bt, cur;
    const int *dep;

    dep      = piDepths + 1;            // skip depth zero
    bitwidth = 0;                       // start from small bitwidths
    cur      = 0;                       // current bit pattern
    do
    {
        do
        {
            bitwidth++;                         // go deeper
            depthbit = 1 << (bitwidth - 1);     // keep depth marker
            d = *dep++;                         // get count here
        }
        while (!d);                           // until count non-zero
        while (d--)
        {                           // for all on this level
            pxEncodingTable[item->Symbol].wid  = bitwidth; // record width
            pxEncodingTable[item->Symbol].bits = cur;      // record bits
            item--;                             // count backwards an item
            bt = depthbit;                      // bt is a temp value
            while (1)
            {
                cur  ^= bt;                     // do an add modulo 1
                if ((cur & bt) || !bt)          // break if now a 1
                    break;                      // or out of bits
                bt  >>=  1;                     // do next bit position
            }
        }
    }
    while (cur);                              // until cur exhausted
}

static u_int Encode_EncodeBytes( int* piDest, const u_char* pucSrc, u_int uNumBytes, const HuffEncode* pxEncodingTable )
{
    int          *start;
    int wid, val, available;
    u_int  accum=0, bits;
    const u_char *sourcelim, *sourceend;

    if( !uNumBytes ) { return 0; }

    start = piDest;
    sourcelim = sourceend = pucSrc + uNumBytes;
    available = 32;
    if (sourcelim - 32 < sourcelim)
	{
        sourcelim -= 32;
	}
    else
	{
        sourcelim = pucSrc;
	}
    if( pucSrc < sourcelim )
    {
        do
        {
            goto lpstart;
            do
            {
                accum = (accum >> wid) | (bits << (32 - wid));
lpstart:        val  = *pucSrc++;
                wid  = pxEncodingTable[val].wid;
                bits = pxEncodingTable[val].bits;
            }
            while ((available -= wid) >= 0);
           
            wid       += available;
            if (wid) accum = (accum >> wid) | (bits << (32 - wid));
            *piDest++  = accum;
            wid       -= available;
            accum      = bits << (32 - wid);
            available += 32;
        }
        while( pucSrc < sourcelim );
    }
    while (1)
    {
        if( pucSrc < sourceend )
		{
            val = *pucSrc++;
        }
        else if( pucSrc == sourceend )
        {
            val = 0x100;                        // terminator
            pucSrc++;
        }
        else break;                             // done

        wid  = pxEncodingTable[val].wid;
        bits = pxEncodingTable[val].bits;

        if ((available -= wid) < 0)
        {
            wid       += available;
            if (wid)
                accum  = (accum >> wid) | (bits << (32 - wid));
            *piDest++  = accum;
            wid       -= available;
            accum      = bits << (32 - wid);
            available += 32;
        }
        else
		{
            accum = (accum >> wid) | (bits << (32 - wid));
		}
    }

    *piDest++ = accum >> available;
    return (piDest - start) * 4;
}

static void Decode_MakeDecodeTable( const int* piDepths, const u_char* pucByteAssignments, int* piOut )
{
    int thisdepth, depthbit, repcount, repspace, lenbits, temp, count;
	int *outp;
    int o = 0;
    const u_char *p;

    lenbits   = 0;
    repcount  = 1 << ASURA_HUFFMAN_MAX_DEPTH;
    repspace  = 1;
    thisdepth = 0;
    depthbit  = 4;
    p         = pucByteAssignments + 255;
    while (1)
    {
        do
        {
            lenbits++;
            depthbit <<= 1;
            repspace <<= 1;
            repcount >>= 1;
        }
        while ((thisdepth = *piDepths++)==0);

        do
        {
            if( p < pucByteAssignments )
			{
                temp = 0xff;
            }
            else 
            {
                temp = lenbits | (*p-- << 8);
            }
            outp  = piOut + (o >> 2);
            count = repcount;
            do
            {
				u_char*	pucOutput	= reinterpret_cast<u_char*>(outp);

				// ART - 13/07/05 - The old code was incorrect for big-endian
				// hardware, it was simply optimising the writing of 4 chars 
				// into 1 integer (notice the >> 2 above), unfortunately this
				// provides the wrong data on big-endian architecture and hence
				// we need to write one byte at a time.
				pucOutput [0]	= static_cast<u_char>(temp & 0xFF);
				pucOutput [1]	= static_cast<u_char>((temp >> 8) & 0xFF);
				pucOutput [2]	= static_cast<u_char>((temp >> 16) & 0xFF);
				pucOutput [3]	= static_cast<u_char>((temp >> 24) & 0xFF);
				// ART...*outp  = temp;
                outp  += repspace;
            }
            while (--count);
            temp = depthbit;
            do
            {
                temp >>= 1;
                if (temp & 3) return;
                o ^= temp;
            }
            while (!(o & temp));
        }
        while (--thisdepth);
    }
}

#ifdef __SNC__
#pragma diag_default=129 // loop is not reachable from preceding code
#endif
