/*---------------------------------------------------------------------------*
  Project:     Compress/uncompress library
  File:        CXStreamingUncompression.c
  Programmer:  Makoto Takano

  Copyright 2005 Nintendo.  All rights reserved.

  These coded instructions, statements, and computer programs contain
  proprietary information of Nintendo of America Inc. and/or Nintendo
  Company Ltd., and are protected by Federal copyright law.  They may
  not be disclosed to third parties or copied or duplicated in any form,
  in whole or in part, without the prior written consent of Nintendo.
 *---------------------------------------------------------------------------*/

#include <revolution/cx/CXUncompression.h>
#include <revolution/cx/CXStreamingUncompression.h>
#include <revolution/cx/CXSecureUncompression.h>
#include "CXUtil.h"


/*---------------------------------------------------------------------------*
  Name        : CXInitUncompContextRL

  Description : Initializes the streaming decompression context for run-length compressed data.
                

  Arguments   : context     Pointer to the run-length uncompressed context
                dest        Destination address for uncompressed data

  Returns     : Can get the data size after decompression.

 *---------------------------------------------------------------------------*/
void CXInitUncompContextRL( CXUncompContextRL *context, void* dest )
{
    context->destp      = (u8*)dest;
    context->destCount  = 0;
    context->flags      = 0;
    context->length     = 0;
    context->headerSize = 8;
    context->forceDestCount = 0;
}


/*---------------------------------------------------------------------------*
  Name        : CXInitUncompContextLZ

  Description : Initializes the streaming decompression context for LZ compressed data.

  Arguments   : context     Pointer to the LZ uncompressed context
                dest        Destination address for uncompressed data
                header      Pointer to the start data for the compressed data

 *---------------------------------------------------------------------------*/
void CXInitUncompContextLZ( CXUncompContextLZ *context, void* dest )
{
    context->destp      = (u8*)dest;
    context->destCount  = 0;
    context->flags      = 0;
    context->flagIndex  = 0;
    context->length     = 0;
    context->lengthFlg  = 3;
    context->headerSize = 8;
    context->exFormat   = 0;
    context->forceDestCount = 0;
}


/*---------------------------------------------------------------------------*
  Name        : CXInitUncompContextHuffman

  Description : Initializes the streaming decompression context for Huffman compressed data.

  Arguments   : context     Pointer to the Huffman uncompressed context
                dest        Destination address for uncompressed data
                header      Pointer to the start data for the compressed data

 *---------------------------------------------------------------------------*/
void CXInitUncompContextHuffman( CXUncompContextHuffman *context, void* dest )
{
    context->destp      = (u8*)dest;
    context->destCount  = 0;
    context->bitSize    = 0;
    context->treeSize   = -1;
    context->treep      = &context->tree[ 0 ];
    context->destTmp    = 0;
    context->destTmpCnt = 0;
    context->srcTmp     = 0;
    context->srcTmpCnt  = 0;
    context->headerSize = 8;
    context->forceDestCount = 0;
}

/*---------------------------------------------------------------------------*
  Name:         CXiReadHeader

  Description:  Header analysis

  Arguments:    headerSize    Pointer to the remaining size of the header to be read
                *destCount    Pointer to the uncompressed data size
                srcp          Pointer to a buffer containing the header information
                srcSize       Pointer to the size of the buffer containing the header information
                forceDestSize Uncompressed data size (if 0, use the binary header information as is)

  Returns:      The size of the source data read
 *---------------------------------------------------------------------------*/
static inline u32 CXiReadHeader( u8* headerSize, s32 *destCount, const u8* srcp, u32 srcSize, s32 forceDestSize )
{
    u32 readLen = 0;
    while ( *headerSize > 0 )
    {
        --(*headerSize);
        if ( *headerSize <= 3 )
        {
            *destCount |= (*srcp << ((3 - *headerSize) * 8));
        }
        else if ( *headerSize <= 6 )
        {
            *destCount |= (*srcp << ((6 - *headerSize) * 8));
        }
        ++srcp;
        ++readLen;
        if ( *headerSize == 4 && *destCount > 0 )
        {
            *headerSize = 0;
        }
        if ( --srcSize == 0 && *headerSize > 0 )
        {
            return readLen;
        }
    }
    
    if ( ( forceDestSize > 0          ) &&
         ( forceDestSize < *destCount ) )
    {
        *destCount = forceDestSize;
    }
    return readLen;
}


/*---------------------------------------------------------------------------*
  Name        : CXReadUncompRL

  Description : This function performs streaming decompression of run-length compressed data.
                Data is written in units of 8 bits. Data cannot be directly uncompressed to VRAM.

  Arguments   : context Pointer to the run-length uncompressed context
                data    Pointer to the next data
                len     Data size

  Returns     : Size of remaining uncompressed data.
                Returns a negative error code if failed.

 *---------------------------------------------------------------------------*/
s32 CXReadUncompRL( CXUncompContextRL *context, const void* data, u32 len )
{
    const u8* srcp = (const u8*)data;
    u8  srcTmp;
    
    // Header analysis
    if ( context->headerSize > 0 )
    {
        u32 read_len;
        if ( context->headerSize == 8 )
        {
            if ( (*srcp & CX_COMPRESSION_TYPE_MASK) != CX_COMPRESSION_RL )
            {
                return CX_ERR_UNSUPPORTED;
            }
            if ( (*srcp & 0x0F ) != 0 )
            {
                return CX_ERR_UNSUPPORTED;
            }
        }
        read_len = CXiReadHeader( &context->headerSize, &context->destCount, srcp, len, context->forceDestCount );
        srcp += read_len;
        len  -= read_len;
        if ( len == 0 )
        {
            return (context->headerSize == 0)? context->destCount : -1;
        }
    }
    
    while ( context->destCount > 0 )
    {
        // Process when length is greater than 0.
        if ( ! (context->flags & 0x80) )
        // Uncompressed data has a length not equal to 0
        {
            while ( context->length > 0 )
            {
                *context->destp++ = *srcp++;
                context->length--;
                context->destCount--;
                len--;
                // End when the prepared buffer has been read in full
                if ( len == 0 )
                {
                    return context->destCount;
                }
            }
        }
        else if ( context->length > 0 )
        // Compressed data has a length not equal to 0
        {
            srcTmp = *srcp++;
            len--;
            while ( context->length > 0 )
            {
                *context->destp++ = srcTmp;
                context->length--;
                context->destCount--;
            }
            if ( len == 0 )
            {
                return context->destCount;
            }
        }
        
        // Reading the flag byte
        context->flags = *srcp++;
        len--;
        context->length = (u16)(context->flags & 0x7F);
        if ( context->flags & 0x80 )
        {
            context->length += 3;
        }
        else
        {
            context->length += 1;
        }
        
        if ( context->length > context->destCount )
        // Measures for buffer overrun when invalid data is decompressed.
        {
            if ( context->forceDestCount == 0 )
            {
                return CX_ERR_DEST_OVERRUN;
            }
            context->length = (u16)context->destCount;
        }
        if ( len == 0 )
        {
            return context->destCount;
        }
    }
    
    // Processing to perform in the event that (context->destCount  == 0)
    if ( (context->forceDestCount == 0) && (len > 32) )
    {
        return CX_ERR_SRC_REMAINDER;
    }
    return 0;
}


/*---------------------------------------------------------------------------*
  Name        : CXReadUncompLZ

  Description : This function performs streaming decompression of LZ compressed data.
                Data is written in units of 8 bits. Data cannot be directly uncompressed to VRAM.

  Arguments   : context Pointer to the LZ uncompressed context
                data    Pointer to the next data
                len     Data size

  Returns     : Size of remaining uncompressed data.
                Returns a negative error code if failed.

 *---------------------------------------------------------------------------*/
s32 CXReadUncompLZ( CXUncompContextLZ *context, const void* data, u32 len )
{
    const u8* srcp = (const u8*)data;
    s32 offset;
    
    // Header analysis
    if ( context->headerSize > 0 )
    {
        u32 read_len;
        // Process the first byte
        if ( context->headerSize == 8 )
        {
            if ( ( *srcp & CX_COMPRESSION_TYPE_MASK ) != CX_COMPRESSION_LZ )
            {
                return CX_ERR_UNSUPPORTED;
            }
            // Record as an LZ compression parameter
            context->exFormat = (u8)( *srcp & 0x0F );
            if ( (context->exFormat != 0x0) && (context->exFormat != 0x1) )
            {
                return CX_ERR_UNSUPPORTED;
            }
        }
        read_len = CXiReadHeader( &context->headerSize, &context->destCount, srcp, len, context->forceDestCount );
        srcp += read_len;
        len  -= read_len;
        if ( len == 0 )
        {
            return (context->headerSize == 0)? context->destCount : -1;
        }
    }
    
    while ( context->destCount > 0 )
    {
        while ( context->flagIndex > 0 )
        {
            if ( len == 0 )
            {
                return context->destCount;
            }
            
            if ( ! (context->flags & 0x80) )
            // Process for non-compressed data
            {
                *context->destp++ = *srcp++;
                context->destCount--;
                len--;
            }
            else
            // Process for compressed data
            {
                while ( context->lengthFlg > 0 )
                {
                    --context->lengthFlg;
                    if ( ! context->exFormat )
                    {
                        context->length  = *srcp++;
                        context->length += (3 << 4);
                        context->lengthFlg = 0;
                    }
                    else
                    {
                        switch ( context->lengthFlg )
                        {
                        case 2:
                            {
                                context->length = *srcp++;
                                if ( (context->length >> 4) == 1 )
                                {
                                    // Read two more bytes
                                    context->length =  (context->length & 0x0F) << 16;
                                    context->length += ( (0xFF + 0xF + 3) << 4 );
                                }
                                else if ( (context->length >> 4) == 0 )
                                {
                                    // Read one more byte
                                    context->length =  (context->length & 0x0F) << 8;
                                    context->length += ( (0xF + 2) << 4 );
                                    context->lengthFlg = 1;
                                }
                                else
                                {
                                    context->length += (1 << 4);
                                    context->lengthFlg = 0;
                                }
                            }
                            break;
                        case 1:
                            {
                                context->length += (*srcp++ << 8);
                            }
                            break;
                        case 0:
                            {
                                context->length += *srcp++;
                            }
                            break;
                        }
                    }
                    if ( --len == 0 )
                    {
                        return context->destCount;
                    }
                }
                
                offset = (context->length & 0xF) << 8;
                context->length = context->length >> 4;
                offset = (offset | *srcp++) + 1;
                len--;
                context->lengthFlg = 3;
                
                // Measures for buffer overrun when invalid data is decompressed.
                if ( context->length > context->destCount )
                {
                    if ( context->forceDestCount == 0 )
                    {
                        return CX_ERR_DEST_OVERRUN;
                    }
                    context->length = context->destCount;
                }
                // Copy a length amount of data located at the offset position
                while ( context->length > 0 )
                {
                    *context->destp = context->destp[ -offset ];
                    context->destp++;
                    context->destCount--;
                    context->length--;
                }
            }
            
            if ( context->destCount == 0 )
            {
                goto out;
            }
            context->flags <<= 1;
            context->flagIndex--;
        }
        
        if ( len == 0 )
        {
            return context->destCount;
        }
        // Read a new flag
        context->flags     = *srcp++;
        context->flagIndex = 8;
        len--;
    }
    
out:
    // Processing to perform in the event that (context->destCount  == 0)
    if ( (context->forceDestCount == 0) && (len > 32) )
    {
        return CX_ERR_SRC_REMAINDER;
    }
    return 0;
}


// Get the next node in the Huffman signed table
static inline u8* GetNextNode( const u8* pTree, u32 select )
{
    return (u8*)( ((u32)pTree & ~0x1) + ( ( (*pTree & 0x3F) + 1 ) * 2 ) + select );
}

extern BOOL CXiVerifyHuffmanTable_( const void* pTable, u8 bit );

/*---------------------------------------------------------------------------*
  Name        : CXReadUncompHuffman

  Description : This function performs streaming decompression of Huffman compressed data.
                Returns a negative error code if failed.

  Arguments   : context Pointer to the Huffman uncompressed context
                data    Pointer to the next data
                len     Data size

  Returns     : Size of remaining uncompressed data.
                Returns a negative error code if failed.

 *---------------------------------------------------------------------------*/
s32 CXReadUncompHuffman( CXUncompContextHuffman *context, const void* data, u32 len )
{
#define TREE_END_MASK   0x80U
    const u8* srcp = (const u8*)data;
    u32  select;
    u32  endFlag;
    
    // Header analysis
    if ( context->headerSize > 0 )
    {
        u32 read_len;
        // Process the first byte
        if ( context->headerSize == 8 )
        {
            context->bitSize = (u8)(*srcp & 0xF);
            
            if ( ( *srcp & CX_COMPRESSION_TYPE_MASK ) != CX_COMPRESSION_HUFFMAN )
            {
                return CX_ERR_UNSUPPORTED;
            }
            if ( (context->bitSize != 4) && (context->bitSize != 8) )
            {
                return CX_ERR_UNSUPPORTED;
            }
        }
        read_len = CXiReadHeader( &context->headerSize, &context->destCount, srcp, len, context->forceDestCount );
        srcp += read_len;
        len  -= read_len;
        if ( len == 0 )
        {
            return (context->headerSize == 0)? context->destCount : -1;
        }
    }
    
    // treeSize is set to -1 in CXInitUncompContextHuffman.
    // When context->treeSize is negative, the data's beginning is used.
    if ( context->treeSize < 0 )
    {
        context->treeSize = (s16)( ( *srcp + 1 ) * 2 - 1 );
        *context->treep++ = *srcp++;
        len--;
    }
    
    // Load the Huffman signed table
    while ( context->treeSize > 0 )
    {
        if ( len == 0 )
        {
            return context->destCount;
        }
        *context->treep++ = *srcp++;
        context->treeSize--;
        len--;
        if ( context->treeSize == 0 )
        {
            context->treep = &context->tree[ 1 ];
            if ( ! CXiVerifyHuffmanTable_( &context->tree[ 0 ], context->bitSize ) )
            {
                return CX_ERR_ILLEGAL_TABLE;
            }
        }
    }
    
    // Decoding process
    while ( context->destCount > 0 )
    {
        // src data is read in 4-byte units
        while ( context->srcTmpCnt < 32 )
        {
            if ( len == 0 )
            {
                return context->destCount;
            }
            context->srcTmp |= (*srcp++) << context->srcTmpCnt;
            len--;
            context->srcTmpCnt += 8;
        }
        
        // The loaded 32 bits are decoded. After those 32 bits are processed, the next 4 bytes are read.
        while ( context->srcTmpCnt > 0 )
        {
            select  = context->srcTmp >> 31;
            endFlag = (*context->treep << select) & TREE_END_MASK;
            context->treep = GetNextNode( context->treep, select );
            context->srcTmp <<= 1;
            context->srcTmpCnt--;
            
            if ( ! endFlag )
            {
                continue;
            }
            
            // When the Huffman tree's terminal flag is set, data is stored at the end of the offset
            //  
            context->destTmp >>= context->bitSize;
            context->destTmp |= *context->treep << ( 32 - context->bitSize );
            context->treep = &context->tree[ 1 ];
            context->destTmpCnt += context->bitSize;
            
            if ( context->destCount <= (context->destTmpCnt / 8) )
            {
                context->destTmp    >>= (32 - context->destTmpCnt);
                context->destTmpCnt = 32;
            }
            
            // Write in 4 byte units
            if ( context->destTmpCnt == 32 )
            {
                *(u32*)context->destp = CXiConvertEndian_( context->destTmp );
                context->destp     += 4;
                context->destCount -= 4;
                context->destTmpCnt = 0;
                if ( context->destCount <= 0 )
                {
                    goto out;
                }
            }
        }
    }
    
out:
    // Processing to perform in the event that (context->destCount  == 0)
    if ( (context->forceDestCount == 0) && (len > 32) )
    {
        return CX_ERR_SRC_REMAINDER;
    }
    return 0;
}