Encoder-Decoder/Algo/huffman.c
/*
    Huffman Implementation
    
    Author      => Abdallah Mohamed Elsharif
    Date        => 29-1-2024
*/

#include "huffman.h"

HuffmanTreeNode *huffman_calcfreq(unsigned char *data, size_t ulSize, __uint16_t *unpNumberOfNodes)
{
    HuffmanTreeNode *pNodes;
    Hash *htable;
    Key *pKey;
    Item *pItem;
    size_t ulCtr, ulIdx;

    if ( ulSize == 0 )
        return NULL;

    if ( !(htable = hash_new(256, REPLACEMENT)) )
        return NULL;

    // Iterate over all bytes to be inserted into a hash table with their frequency
    while ( ulSize-- )
    {
        if ( pKey = hash_getkey(htable, data, sizeof(unsigned char)) )
        {
            // Check if the frequency reached the maximum ( freq will be stored in only 4 bytes within the table )
            if ( (ulCtr = *(__uint32_t *) hash_get(htable, pKey)) == 0xffffffff )
            {
                hash_destroy( &htable );
                return NULL;
            }

            ulCtr++;
            pItem = item_new(&ulCtr, sizeof(__uint32_t), malloc, free, NULL, NULL);
            hash_update(htable, pKey, pItem);
        } 
        
        else {
            ulCtr = 1;
            pKey = key_new(data, sizeof(unsigned char), NULL, ByteCmp);
            pItem = item_new(&ulCtr, sizeof(__uint32_t), malloc, free, NULL, NULL);
            hash_insert(htable, pKey, pItem);
        }
        
        data++;
    }

    // Set the number of nodes
    ( *unpNumberOfNodes ) = hash_size( htable );

    // Allocate memory for the nodes
    if ( !(pNodes = (HuffmanTreeNode *) malloc(*unpNumberOfNodes * sizeof(HuffmanTreeNode))) )
    {
        hash_destroy( &htable );
        return NULL;
    }

    ulIdx = ulCtr = 0;

    while ( hash_next(htable, &ulCtr, &pKey, &pItem) )
    {
        ( pNodes + ulIdx )->data = *(unsigned char *) pKey->data;
        ( pNodes + ulIdx )->ulFrequency = *(__uint32_t *) hash_get( htable, pKey );
        ( pNodes + ulIdx )->left = NULL;
        ( pNodes + ulIdx )->right = NULL;
        ulIdx++;
    }
    
    hash_destroy( &htable );

    return pNodes;
}

HuffmanTreeNode *huffman_parsefreq(unsigned char *data, __uint16_t *unpNumberOfNodes)
{
    HuffmanTreeNode *pNodes;

    // Freq table size ( We start counting from 0, so we have to add one )
    *unpNumberOfNodes = (__uint16_t) ( *data++ ) + 1;

    if ( !(pNodes = malloc(*unpNumberOfNodes * sizeof(HuffmanTreeNode))) )
        return NULL;

    for ( __uint16_t unIdx = 0; unIdx < *unpNumberOfNodes; unIdx++ )
    {
        pNodes[ unIdx ].data = ((HuffmanItem *) data)->data;
        pNodes[ unIdx ].ulFrequency = ((HuffmanItem *) data)->unFrequency;
        pNodes[ unIdx ].left = NULL;
        pNodes[ unIdx ].right = NULL;

        // Next freq entry
        data += sizeof( HuffmanItem );
    }

    return pNodes;
}

int huffman_cmpnode(void *x, void *y)
{
    return ULongCmp(
        &( (HuffmanTreeNode *) x )->ulFrequency,
        &( (HuffmanTreeNode *) y )->ulFrequency
    );
}

HuffmanTreeNode *huffman_newnode(HuffmanTreeNode *node)
{
    HuffmanTreeNode *hpNode;

    if ( hpNode = malloc(sizeof(HuffmanTreeNode)) )
    {
        hpNode->data = node->data;
        hpNode->ulFrequency = node->ulFrequency;
        hpNode->left = node->left;
        hpNode->right = node->right;
        return hpNode;    
    }

    return NULL;
}

HuffmanTreeNode *huffman_build(HuffmanTreeNode *hNodes, __uint16_t unSize)
{
    Heap *htree;
    HuffmanTreeNode *hpLeftTreeNode, *hpRightTreeNode, *hpTreeNode;
    HuffmanTreeNode hLeftNode, hRightNode, hNode;

    if ( !(htree = heap_init(unSize, sizeof(HuffmanTreeNode), MINHEAP, malloc, free, NULL, huffman_cmpnode)) )
        return NULL;
    
    if ( ! heap_build(htree, hNodes, unSize) )
    {
        heap_cleanup( &htree );
        return NULL;
    }
    
    while ( heap_size(htree) > 1 )
    {
        hLeftNode = *(HuffmanTreeNode *) heap_root(htree)->data;
        hpLeftTreeNode = huffman_newnode( &hLeftNode );
        heap_delete(htree, &hLeftNode);
        hRightNode = *(HuffmanTreeNode *) heap_root(htree)->data;
        hpRightTreeNode = huffman_newnode( &hRightNode );
        heap_delete(htree, &hRightNode);
        hNode.ulFrequency = hLeftNode.ulFrequency + hRightNode.ulFrequency;
        hNode.left = hpLeftTreeNode;
        hNode.right = hpRightTreeNode;
        hNode.data = 0x00;
        heap_insert(htree, &hNode);
    }

    hNode = *(HuffmanTreeNode *) heap_root( htree )->data;
    hpTreeNode = huffman_newnode( &hNode );
    heap_delete(htree, &hNode);
    heap_cleanup( &htree );

    return hpTreeNode;
}

void huffman_traverse2(HuffmanTreeNode *pRoot, Hash *pHuffmanTable, HuffmanCode c)
{
    if ( !pRoot->left && !pRoot->right )
    {
        c.unFrequency = pRoot->ulFrequency;

        hash_insert(
            pHuffmanTable, 
            key_new( &pRoot->data, 1, NULL, ByteCmp ), 
            item_new( &c, sizeof(HuffmanCode), malloc, free, NULL, NULL )
        );
    }

    else {
        c.ulLength++;

        if ( pRoot->left )
        {
            c.unCode <<= 1;
            huffman_traverse2(pRoot->left, pHuffmanTable, c);
        }

        if ( pRoot->right )
        {
            c.unCode += 1;
            huffman_traverse2(pRoot->right, pHuffmanTable, c);
        }
    }
    
}

Hash *huffman_traverse(HuffmanTreeNode *pRoot, __uint16_t unNumberOfNodes)
{
    Hash *pHuffmanTable;
    HuffmanCode c = { 0 };

    if ( pHuffmanTable = hash_new(256, REPLACEMENT) )
    {
        huffman_traverse2(pRoot, pHuffmanTable, c);
        return pHuffmanTable;
    }

    return NULL;
}

// This function calculates the encoded data length from the table
size_t huffman_calcenclen(Hash *pHuffmanTable)
{
    Item *pItem;
    HuffmanCode *pCode;
    size_t ulBitsSize, ulSize, ulPos;

    ulBitsSize = ulSize = ulPos = 0;

    while ( hash_next(pHuffmanTable, &ulPos, NULL, &pItem) )
    {
        pCode = (HuffmanCode *) pItem->data;
        ulBitsSize += pCode->ulLength * pCode->unFrequency;
    }

    // Padding the bits to be convertible to bytes
    if ( ulBitsSize % 8 != 0 )
        ulBitsSize += 8 - ( ulBitsSize % 8 );

    // Bits length in bytes
    ulBitsSize /= 8;

    // frequency table length
    ulSize += 1;

    // frequency table entries size
    ulSize += hash_size(pHuffmanTable) * sizeof(HuffmanItem);

    // Size of the codes 
    ulSize += ulBitsSize;

    return ulSize;
}

// This function calculates the original data length from the table
size_t huffman_calcdeclen(Hash *pHuffmanTable)
{
    HuffmanCode *pCode;
    Item *pItem;
    size_t ulSize, ulPos;

    ulSize = ulPos = 0;

    while ( hash_next(pHuffmanTable, &ulPos, NULL, &pItem) )
    {
        pCode = (HuffmanCode *) pItem->data;
        ulSize += pCode->unFrequency;
    }

    return ulSize;
}

unsigned char *huffman_encode(unsigned char *data, size_t ulSize, bool bForceEncoding, size_t *ulpResultSize)
{
    HuffmanTreeNode *pHuffmanNodes  = NULL;
    HuffmanTreeNode *pHuffmanRoot   = NULL;
    Hash *pHuffmanTable             = NULL;
    unsigned char *pResult          = NULL;
    unsigned char *pTempResult;
    Key *pKey;
    Item *pItem;
    HuffmanCode *pCode;
    __uint16_t unActive;
    __uint8_t unBitPos = 8;
    size_t ulBitsLen, ulPos = 0;

    if ( !(pHuffmanNodes = huffman_calcfreq(data, ulSize, &unActive)) )
        goto LEAVE;

    if ( !(pHuffmanRoot = huffman_build(pHuffmanNodes, unActive)) )
        goto LEAVE;

    if ( !(pHuffmanTable = huffman_traverse(pHuffmanRoot, unActive)) )
        goto LEAVE;

    if ( (*ulpResultSize = huffman_calcenclen(pHuffmanTable)) >= ulSize && ! bForceEncoding )
        goto LEAVE; // We f*cked up, the result size is greater than the orignal size itself.

    if ( !(pTempResult = pResult = (unsigned char *) calloc(*ulpResultSize, sizeof(unsigned char))) )
        goto LEAVE;

    // Build table headers (-1 because we store number of actives in one byte, if all bytes were active then one byte not enough and we will need extra space)
    *pTempResult++ = (__uint8_t) unActive - 1;

    while ( hash_next(pHuffmanTable, &ulPos, &pKey, &pItem) )
    {
        *pTempResult++ = *(unsigned char *) pKey->data;
        *((__uint32_t *) pTempResult) = ((HuffmanCode *) pItem->data)->unFrequency;
        pTempResult += sizeof(__uint32_t);
    }
    
    // Build bits table
    while ( ulSize-- )
    {
        // Retrieve huffman codes of current byte from the codes table 
        pKey = hash_getkey(pHuffmanTable, data, 1);
        pCode = (HuffmanCode *) hash_get(pHuffmanTable, pKey);
        ulBitsLen = pCode->ulLength;

        // Iterate over huffman codes of current byte
        while ( ulBitsLen-- )
        {
            unBitPos--;

            if ( (pCode->unCode >> ulBitsLen) & 1 ) 
                *pTempResult |= ( 1 << unBitPos );

            // Jump on the next byte within encoding buffer if the last bit in current byte had been set
            if ( unBitPos == 0 )
            {
                unBitPos = 8;
                pTempResult++;
            }

        }
        
        data++;
    }


LEAVE:
    if ( pHuffmanNodes )
        free( pHuffmanNodes );

    if ( pHuffmanRoot )
        huffman_cleanup( &pHuffmanRoot );

    if ( pHuffmanTable )
        hash_destroy( &pHuffmanTable );

    return pResult;
}

unsigned char *huffman_decode(unsigned char *data, size_t ulSize, bool bForceDecoding, size_t *ulpResultSize)
{
    HuffmanTreeNode *pHuffmanNodes  = NULL;
    HuffmanTreeNode *pHuffmanRoot   = NULL;
    Hash *pHuffmanTable             = NULL;
    unsigned char *pResult          = NULL;
    __uint8_t *pBitsTable           = NULL;
    unsigned char *pTempResult;
    __uint8_t *pTempBitsTable;
    size_t ulHeaderSize, ulBitsLen;
    __uint16_t unActive;

    if ( !(pHuffmanNodes = huffman_parsefreq(data, &unActive)) )
        goto LEAVE;

    // The table is invalid if the freq table size is greater than the full data
    if ( unActive * sizeof(HuffmanItem) > ulSize )
        goto LEAVE;

    if ( !(pHuffmanRoot = huffman_build(pHuffmanNodes, unActive)) )
        goto LEAVE;

    if ( !(pHuffmanTable = huffman_traverse(pHuffmanRoot, unActive)) )
        goto LEAVE;

    if ( (*ulpResultSize = huffman_calcdeclen(pHuffmanTable)) < ulSize && ! bForceDecoding )
        goto LEAVE; // Why is the encoded data size greater than the original, motherf*cker !?!?
    
    // Calculate table headers size
    ulHeaderSize = 1 + ( unActive * sizeof(HuffmanItem) );

    // Calculate the length of bits
    ulBitsLen = ( ulSize - ulHeaderSize ) * 8;

    if ( !(pTempBitsTable = pBitsTable = calloc(ulBitsLen, sizeof(__uint8_t))) )
        goto LEAVE;

    if ( !(pTempResult = pResult = malloc(*ulpResultSize)) )
        goto LEAVE;

    // Pointing now to the huffman codes
    data += ulHeaderSize;

    // Read bits from the table
    for ( size_t ulBitPos = ulBitsLen; ulBitPos--; )
    {
        // Check if the current bit is set
        if ( *data & (1 << ulBitPos % 8) )
            (*pTempBitsTable)++;

        // The end of current byte ( jump on the next one ) 
        if ( ulBitPos % 8 == 0 )
            data++;

        pTempBitsTable++;
    }

    // Back again to pointing to the beginning of the table
    pTempBitsTable = pBitsTable;

    // Decompress all bytes one by one
    for ( size_t ulCtr = 0; ulCtr < *ulpResultSize; ulCtr++ )
    {
        // Check if we deal with invalid table
        if ( ! (
            pTempBitsTable < pBitsTable + ulBitsLen &&
            huffman_lookup(pHuffmanRoot, pTempResult++, &pTempBitsTable) 
        ) )
        {
            free( pResult );
            pResult = NULL;
            break;
        }
    }

LEAVE:
    if ( pHuffmanNodes )
        free( pHuffmanNodes );

    if ( pHuffmanRoot )
        huffman_cleanup( &pHuffmanRoot );

    if ( pHuffmanTable )
        hash_destroy( &pHuffmanTable );

    if ( pBitsTable )
        free( pBitsTable );

    return pResult;
}

bool huffman_lookup(HuffmanTreeNode *pRoot, unsigned char *pResult, __uint8_t **pBitsTable)
{
    if ( ! pRoot->right && ! pRoot->left )
    {
        *pResult = pRoot->data;
        return true;
    }
    
    if ( *(*pBitsTable)++ )
    {
        if ( ! pRoot->right )
            return false;

        huffman_lookup( pRoot->right, pResult, pBitsTable );
    }

    else {
        if ( ! pRoot->left )
            return false;

        huffman_lookup( pRoot->left, pResult, pBitsTable );
    }
    
}

void huffman_cleanup(HuffmanTreeNode **root)
{
    if ( (*root)->left )
        huffman_cleanup( &(*root)->left );

    if ( (*root)->right )
        huffman_cleanup( &(*root)->right );

    if ( !(*root)->left && !(*root)->right )
    {
        free( *root );
        ( *root ) = NULL;
    }
}
