/*---------------------------------------------------------------------------*
  Project:     Compress/uncompress library
  File:        CXUncompression.h
  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 <dolphin/types.h>
#include <dolphin/os.h>
#include <revolution/cx/CXUncompression.h>
#include <revolution/cx/CXSecureUncompression.h>
#include "CXUtil.h"


//======================================================================
//          Expanding compressed data
//======================================================================

/*---------------------------------------------------------------------------*
  Name:         CXSecureUncompressAny

  Description:  Determines the compression type from the data header and performs the appropriate decompression.
                 
                Since decompression processes for all types of compression are linked, it may be better to execute separate functions for each for each compression type when not using formats other than the specific compression formats.
                
                

  Arguments:    srcp    Source address
                srcSize Source size
                destp   Destination address

  Returns:      Returns 0 when conversion succeeds
                Returns a negative error code if failed
 *---------------------------------------------------------------------------*/
s32 CXSecureUncompressAny( const void* srcp, u32 srcSize, void* destp )
{
    switch ( CXGetCompressionType( srcp ) )
    {
    // Run-length compressed data
    case CX_COMPRESSION_RL:
        return CXSecureUncompressRL( srcp, srcSize, destp );
    // LZ77 compressed data
    case CX_COMPRESSION_LZ:
        return CXSecureUncompressLZ( srcp, srcSize, destp );
    // Huffman compressed data
    case CX_COMPRESSION_HUFFMAN:
        return CXSecureUncompressHuffman( srcp, srcSize, destp );
    // Difference filter
    case CX_COMPRESSION_DIFF:
        return CXSecureUnfilterDiff( srcp, srcSize, destp );
    default:
        return CX_ERR_UNSUPPORTED;
    }
}


/*---------------------------------------------------------------------------*
  Name:         CXSecureUncompressRL

  Description:  8-bit decompression of run-length compressed data

  - Decompresses run-length compressed data, writing in 8 bit units.
  - Use 4 byte alignment for the source address.

  - Data header
      u32 :4                  Reserved
          compType:4          Compression type( = 3)
          destSize:24         Data size after decompression

  - Flag data format
      u8  length:7            Decompressed data length - 1 (When not compressed)
                              Decompressed data length - 3 (only compress when the contiguous length is 3 bytes or greater)
          flag:1              (0, 1) = (not compressed, compressed)

  Arguments:    *srcp   Source address
                srcSize Source size
                *destp  Destination address

  Returns:      Returns 0 when conversion succeeds
                Returns a negative error code if failed.
 *---------------------------------------------------------------------------*/
s32 CXSecureUncompressRL( const void *srcp, u32 srcSize, void *destp )
{
    const u8 *pSrc  = srcp;
    u8       *pDst  = destp;
    u8       compType  = (u8)( CXiConvertEndian_( *(u32*)pSrc ) & 0xFF );
    u32      destCount = CXiConvertEndian_( *(u32*)pSrc ) >> 8;
    s32      srcCount  = (s32)srcSize;
    
    if ( (compType & CX_COMPRESSION_TYPE_MASK) != CX_COMPRESSION_RL )
    {
        return CX_ERR_UNSUPPORTED;
    }
    if ( (compType & 0xF) != 0 )
    {
        return CX_ERR_UNSUPPORTED;
    }
    if ( srcSize <= 4 )
    {
        return CX_ERR_SRC_SHORTAGE;
    }
    
    pSrc     += 4;
    srcCount -= 4;
    
    if ( destCount == 0 )
    {
        if ( srcCount < 4 )
        {
            return CX_ERR_SRC_SHORTAGE;
        }
        destCount = CXiConvertEndian_( *(u32*)pSrc );
        pSrc     += 4;
        srcCount -= 4;
    }
    
    while ( destCount > 0 )
    {
        u8  flags  = *pSrc++;
        s32 length = flags & 0x7f;
        if ( --srcCount < 0 )
        {
            return CX_ERR_SRC_SHORTAGE;
        }
        
        if ( !(flags & 0x80) )
        {
            length++;
            if ( length > destCount )
            // Measures for buffer overrun when invalid data is decompressed.
            {
                return CX_ERR_DEST_OVERRUN;
            }
            srcCount -= length;
            if ( srcCount < 0 )
            {
                return CX_ERR_SRC_SHORTAGE;
            }
            
            destCount -= length;
            do
            {
                *pDst++ = *pSrc++;
            } while ( --length > 0 );
        }
        else
        {
            u8 srcTmp;
            
            length    += 3;
            if ( length > destCount )
            // Measures for buffer overrun when invalid data is decompressed.
            {
                return CX_ERR_DEST_OVERRUN;
            }
            
            destCount -= length;
            srcTmp    = *pSrc++;
            if ( --srcCount < 0 )
            {
                return CX_ERR_SRC_SHORTAGE;
            }
            do
            {
                *pDst++ =  srcTmp;
            } while ( --length > 0 );
        }
    }
    
    if ( srcCount > 32 )
    {
        return CX_ERR_SRC_REMAINDER;
    }
    
    return CX_ERR_SUCCESS;
}


/*---------------------------------------------------------------------------*
  Name:         CXSecureUncompressLZ

  Description:  8-bit decompression of LZ77 compressed data

  - Decompresses LZ77 compressed data, writing in 8 bit units.
  - Use 4 byte alignment for the source address.

  - Data header
      u32 :4                  Reserved
          compType:4          Compression type( = 1)
          destSize:24         Data size after decompression

  - Flag data format
      u8  flags               Compression/no compression flag
                              (0, 1) = (not compressed, compressed)
  - Code data format (Big Endian)
      u16 length:4            Decompressed data length - 3 (only compress when the match length is 3 bytes or greater)
          offset:12           Match data offset - 1

  Arguments:    *srcp   Source address
                srcSize Source size
                *destp  Destination address

  Returns:      Returns 0 when conversion succeeds
                Returns a negative error code if failed.
 *---------------------------------------------------------------------------*/
s32 CXSecureUncompressLZ( const void *srcp, u32 srcSize, void *destp )
{
    const u8* pSrc      = srcp;
    u8*       pDst      = destp;
    u8        compType  = (u8)( CXiConvertEndian_( *(u32*)pSrc ) & 0xFF );
    u32       destCount = CXiConvertEndian_( *(u32*)pSrc ) >> 8;
    s32       srcCount  = (s32)srcSize;
    BOOL      exFormat  = (*pSrc & 0x0F)? TRUE : FALSE;
    
    if ( (compType & CX_COMPRESSION_TYPE_MASK) != CX_COMPRESSION_LZ )
    {
        return CX_ERR_UNSUPPORTED;
    }
    if ( ((compType & 0xF) != 0x0) && ((compType & 0xF) != 0x1) )
    {
        return CX_ERR_UNSUPPORTED;
    }
    if ( srcSize <= 4 )
    {
        return CX_ERR_SRC_SHORTAGE;
    }
    
    pSrc += 4;
    srcCount -= 4;
    
    if ( destCount == 0 )
    {
        if ( srcCount < 4 )
        {
            return CX_ERR_SRC_SHORTAGE;
        }
        
        destCount = CXiConvertEndian_( *(u32*)pSrc );
        pSrc += 4;
        srcCount -= 4;
    }
    
    while ( destCount > 0 )
    {
        u32 i;
        u32 flags = *pSrc++;
        if ( --srcCount < 0 )
        {
            return CX_ERR_SRC_SHORTAGE;
        }
        
        for ( i = 0; i < 8; ++i )
        {
            if ( !(flags & 0x80) )
            {
                *pDst++ = *pSrc++;
                if ( --srcCount < 0 )
                {
                    return CX_ERR_SRC_SHORTAGE;
                }
                destCount--;
            }
            else
            {
                s32 length = (*pSrc >> 4);
                s32 offset;
                
                if ( ! exFormat )
                {
                    length += 3;
                }
                else
                {
                    // LZ77 extended format
                    if ( length == 1 )
                    {
                        length =  (*pSrc++ & 0x0F) << 12;
                        length |= (*pSrc++) << 4;
                        length |= (*pSrc >> 4);
                        length += 0xFF + 0xF + 3;
                        srcCount -= 2;
                    }
                    else if ( length == 0 )
                    {
                        length =  (*pSrc++ & 0x0F) << 4;
                        length |= (*pSrc >> 4);
                        length += 0xF + 2;
                        srcCount -= 1;
                    }
                    else
                    {
                        length += 1;
                    }
                }
                offset  = (*pSrc++ & 0x0f) << 8;
                offset = (offset | *pSrc++) + 1;
                
                srcCount -= 2;
                if ( srcCount < 0 )
                {
                    return CX_ERR_SRC_SHORTAGE;
                }
                
                // Measures for buffer overrun when invalid data is decompressed.
                if ( length > destCount )
                {
                    return CX_ERR_DEST_OVERRUN;
                }
                if ( &pDst[ -offset ] < destp )
                {
                    return CX_ERR_DEST_OVERRUN;
                }
                
                destCount -= length;
                do
                {
                    *pDst++ = pDst[ -offset ];
                } while ( --length > 0 );
            }
            if ( destCount <= 0 )
            {
                break;
            }
            flags <<= 1;
        }
    }
    
    if ( srcCount > 32 )
    {
        return CX_ERR_SRC_REMAINDER;
    }
    
    return CX_ERR_SUCCESS;
}


extern BOOL CXiVerifyHuffmanTable_( const void* pTable, u8 bit );
/*---------------------------------------------------------------------------*
  Name:         CXiVerifyHuffmanTable_

  Description:  Huffman table integrity check

  Arguments:    Pointer to the Huffman table

  Returns:      TRUE when the table is valid
                FALSE when the table is invalid
 *---------------------------------------------------------------------------*/
BOOL
CXiVerifyHuffmanTable_( const void* pTable, u8 bit )
{
    const u32 FLAGS_ARRAY_NUM = 512 / 8; /* 64 Byte */
    u8* treep = (u8*)pTable;
    u8* treeStartp = treep + 1;
    u32 treeSize   = *treep;
    u8* treeEndp   = (u8*)pTable + (treeSize + 1) * 2;
    u32 i;
    u8  end_flags[ FLAGS_ARRAY_NUM ];
    u32 idx;
    
    for ( i = 0; i < FLAGS_ARRAY_NUM; i++ )
    {
        end_flags[ i ] = 0;
    }
    
    if ( bit == 4 )
    {
        if ( treeSize >= 0x10 )
        {
            return FALSE;
        }
    }
    
    idx = 1;
    treep = treeStartp;
    while ( treep < treeEndp )
    {
        if ( (end_flags[ idx / 8 ] & (1 << (idx % 8) )) == 0 )
        {
            u32  offset = (u32)( ( (*treep & 0x3F) + 1 ) << 1);
            u8*  nodep  = (u8*)( (((u32)treep >> 1) << 1) + offset );
            
            // Skip data added at the end for alignment.
            if ( *treep == 0 && idx >= (treeSize * 2) )
            {
                goto next;
            }
            
            if ( nodep >= treeEndp )
            {
                return FALSE;
            }
            if ( *treep & 0x80 )
            {
                u32 left = (idx & ~0x1) + offset;
                end_flags[ left / 8 ] |= (u8)( 1 << (left % 8) );
            }
            if ( *treep & 0x40 )
            {
                u32 right = (idx & ~0x1) + offset + 1;
                end_flags[ right / 8 ] |= (u8)( 1 << (right % 8) );
            }
        }
    next:
        ++idx;
        ++treep;
    }
    return TRUE;
}


/*---------------------------------------------------------------------------*
  Name:         CXSecureUncompressHuffman

  Description:  Decompression of Huffman compressed data

  - Decompresses Huffman compressed data, writing in 32 bit units.
  - Use 4 byte alignment for the source address.
  - Use 4 byte alignment for the destination address.
  - The target decompression buffer size must be prepared in 4 byte multiples.

  - Data header
      u32 bitSize:4           1 data bit size (Normally 4|8)
          compType:4          Compression type( = 2)
          destSize:24         Data size after decompression

  - Tree table
      u8           treeSize        Tree table size/2 - 1
      TreeNodeData nodeRoot        Root node

      TreeNodeData nodeLeft        Root left node
      TreeNodeData nodeRight       Root right node

      TreeNodeData nodeLeftLeft    Left left node
      TreeNodeData nodeLeftRight   Left right node

      TreeNodeData nodeRightLeft   Right left node
      TreeNodeData nodeRightRight  Right right node

              .
              .

    The compressed data itself follows

  - TreeNodeData structure
      u8  nodeNextOffset:6    Offset to the next node data - 1 (2 byte units)
          rightEndFlag:1      Right node end flag
          leftEndzflag:1      Left node end flag
                              When end flag is set
                              There is data in next node

  Arguments:    *srcp   Source address
                srcSize Source size
                *destp  Destination address

  Returns:      Returns 0 when conversion succeeds
                Returns a negative error code if failed.
 *---------------------------------------------------------------------------*/
s32 CXSecureUncompressHuffman( const void *srcp, u32 srcSize, void *destp )
{
#define TREE_END 0x80
    const u32 *pSrc          = srcp;
    u32       *pDst          = destp;
    u8        compType       = (u8)( CXiConvertEndian_( *(u32*)pSrc ) & 0xFF );
    s32       destCount      = (s32)( CXiConvertEndian_( *pSrc ) >> 8 );
    u8        *treep         = (destCount != 0)? ((u8*)pSrc + 4) : ((u8*)pSrc + 8);
    u8        *treeStartp    = treep + 1;
    u32       dataBit        = *(u8*)pSrc & 0x0FU;
    u32       destTmp        = 0;
    u32       destTmpCount   = 0;
    u32       destTmpDataNum = 4 + ( dataBit & 0x7 );
    s32       srcCount       = (s32)srcSize;
    u32       treeSize       = (u32)( (*treep + 1) << 1 );
    
    if ( (compType & CX_COMPRESSION_TYPE_MASK) != CX_COMPRESSION_HUFFMAN )
    {
        return CX_ERR_UNSUPPORTED;
    }
    if ( (dataBit != 4) && (dataBit != 8) )
    {
        return CX_ERR_UNSUPPORTED;
    }
    
    if ( destCount == 0 )
    {
        if ( srcSize < 8 + treeSize )
        {
            return CX_ERR_SRC_SHORTAGE;
        }
        destCount = (s32)( CXiConvertEndian_( *(pSrc + 1) ) );
    }
    else if ( srcSize < 4 + treeSize )
    {
        return CX_ERR_SRC_SHORTAGE;
    }
    
    if ( ! CXiVerifyHuffmanTable_(treep, (u8)dataBit) )
    {
        return CX_ERR_ILLEGAL_TABLE;
    }
    
    pSrc  = (u32*)( treep + treeSize );
    srcCount -= ( (u32)pSrc - (u32)srcp );
    
    if ( srcCount < 0 )
    {
        return CX_ERR_SRC_SHORTAGE;
    }
    
    treep = treeStartp;
    
    while ( destCount > 0 )
    {
        s32 srcTmpCount = 32;
        u32 srcTmp   = CXiConvertEndian_( *pSrc++ );      // Endian strategy
        srcCount -= 4;
        if ( srcCount < 0 )
        {
            return CX_ERR_SRC_SHORTAGE;
        }
        
        while ( --srcTmpCount >= 0 )
        {
            u32 treeShift = (srcTmp >> 31) & 0x1;
            u32 treeCheck = *treep;
            treeCheck <<= treeShift;
            treep = (u8*)( (((u32)treep >> 1) << 1) + (((*treep & 0x3f) + 1) << 1) + treeShift );
            if ( treeCheck & TREE_END )
            {
                destTmp >>= dataBit;
                destTmp |= *treep << (32 - dataBit);
                treep = treeStartp;
                ++destTmpCount;
                
                if ( destCount <= (destTmpCount * dataBit) / 8 )
                {
                    destTmp >>= (destTmpDataNum - destTmpCount) * dataBit;
                    destTmpCount = destTmpDataNum;
                }
                
                if ( destTmpCount == destTmpDataNum )
                {
                    // Over-access until the last 4-byte alignment of the decompression buffer
                    *pDst++ = CXiConvertEndian_(destTmp); // Endian strategy
                    destCount -= 4;
                    destTmpCount = 0;
                }
            }
            if ( destCount <= 0 )
            {
                break;
            }
            srcTmp <<= 1;
        }
    }
    if ( srcCount > 32 )
    {
        return CX_ERR_SRC_REMAINDER;
    }
    
    return CX_ERR_SUCCESS;
}


/*---------------------------------------------------------------------------*
  Name:         CXSecureUnfilterDiff

  Description:  8-bit decompression to restore differential filter conversion.

    - Restores a differential filter, writing in 8 bit units.
    - Use 4 byte alignment for the source address.

  Arguments:    *srcp   Source address
                *destp  Destination address

  Returns:      Returns 0 when conversion succeeds
                Returns a negative error code if failed.
 *---------------------------------------------------------------------------*/
s32 CXSecureUnfilterDiff( register const void *srcp, u32 srcSize, register void *destp )
{
    const u8* pSrc = srcp;
    u8*       pDst = destp;
    u32       bitSize    = *pSrc & 0xFU;
    u8        compType   = (u8)( CXiConvertEndian_( *(u32*)pSrc ) & 0xFF );
    s32       destCount  = (s32)( CXiConvertEndian_( *(u32*)pSrc ) >> 8 );
    u32       sum = 0;
    s32       srcCount  = (s32)srcSize;
    
    if ( (compType & CX_COMPRESSION_TYPE_MASK) != CX_COMPRESSION_DIFF )
    {
        return CX_ERR_UNSUPPORTED;
    }
    if ( (bitSize != 0) && (bitSize != 1) )
    {
        return CX_ERR_UNSUPPORTED;
    }
    if ( srcSize <= 4 )
    {
        return CX_ERR_SRC_SHORTAGE;
    }
    
    pSrc     += 4;
    srcCount -= 4;
    
    if ( bitSize != 1 )
    {
        // Difference calculation in units of 8 bits
        do 
        {
            u8 tmp = *(pSrc++);
            if ( --srcCount < 0 )
            {
                return CX_ERR_SRC_SHORTAGE;
            }
            destCount--;
            sum += tmp;
            *(pDst++) = (u8)sum;
        } while ( destCount > 0 );
    }
    else
    {
        // Difference calculation in units of 16 bits
        do 
        {
            u16 tmp = CXiConvertEndian16_( *(u16*)pSrc );
            pSrc += 2;
            srcCount -= 2;
            if ( srcCount < 0 )
            {
                return CX_ERR_SRC_SHORTAGE;
            }
            destCount -= 2;
            sum += tmp;
            *(u16*)pDst = CXiConvertEndian16_( (u16)sum );
            pDst += 2;
        } while ( destCount > 0 );
    }
    if ( srcCount > 32 )
    {
        return CX_ERR_SRC_REMAINDER;
    }
    
    return CX_ERR_SUCCESS;
}