HuffmanCode.cpp

#include <iostream>
#include <queue>
#include <unordered_map>
#include <fstream>
#include <string>
#include <vector>
#include <bitset>

using namespace std;

/*
 * HuffmanCode — lossless file compression/decompression via Huffman coding.
 *
 * Compression:  count character frequencies → build a min-heap → merge nodes
 *               into a tree → assign binary codes → write header + packed bitstream.
 * Decompression: read header to rebuild code table → decode bitstream bit by bit
 *                → stop when the embedded '~' EOF marker is reached.
 */
class HuffmanCode {
private:

    // ── HuffNode ──────────────────────────────────────────────────────────────
    // Leaf nodes store a real character and its frequency.
    // Interior nodes store only a combined frequency; 'c' is '\0' as a placeholder.
    class HuffNode {
    public:
        char c;
        int  freq;
        HuffNode* left;
        HuffNode* right;

        // Leaf node
        HuffNode(char c, int f)
            : c(c), freq(f), left(nullptr), right(nullptr) {}

        // Interior node — frequency is the sum of both children
        HuffNode(HuffNode* left, HuffNode* right)
            : c('\0'), freq(left->freq + right->freq),
              left(left), right(right) {}
    };

    // Min-heap comparator — lower frequency = higher priority.
    // No char tie-breaking: interior nodes all share '\0', making it meaningless.
    struct compareNodes {
        bool operator()(HuffNode* a, HuffNode* b) {
            return a->freq > b->freq;
        }
    };

    HuffNode* root;
    unordered_map<char, string> huffmanCodes; // char  → binary code  (compression)
    unordered_map<string, char> codeToChar;   // binary code → char   (decompression)


    // Post-order recursive deletion of the entire tree
    void deleteTree(HuffNode* node) {
        if (node == nullptr) return;
        deleteTree(node->left);
        deleteTree(node->right);
        delete node;
    }

    // Read input file and count occurrences of each byte
    unordered_map<char, int> buildFreqTable(const string& fileName) {
        unordered_map<char, int> freqMap;
        ifstream file(fileName, ios::binary);
        char c;
        while (file.get(c)) freqMap[c]++;
        return freqMap;
    }

    // Walk the completed tree; assign a code to each leaf by accumulating
    // '0' (left) or '1' (right) at every edge on the way down
    void generateCodes(HuffNode* node, const string& curCode) {
        if (node == nullptr) return;
        if (!node->left && !node->right) {
            huffmanCodes[node->c] = curCode;
            return;
        }
        generateCodes(node->left,  curCode + "0");
        generateCodes(node->right, curCode + "1");
    }

    // Write the code table as a text header so the decompressor can
    // reconstruct the mapping without the original file.
    // Format:  <ASCII decimal> <code>\n  ...  ====\n
    void writeHuffHeader(ofstream& outFile) {
        for (const auto& [c, code] : huffmanCodes)
            outFile << static_cast<int>(static_cast<unsigned char>(c))
                    << " " << code << "\n";
        outFile << "====\n"; // sentinel — encoded content follows
    }

    // Parse the header written by writeHuffHeader and populate codeToChar.
    // Reads until the "====" sentinel line is found.
    void readHuffHeader(ifstream& inFile) {
        string line;
        while (getline(inFile, line)) {
            if (line == "====") break;

            // size_t / string::npos used here — string::find returns an unsigned
            // type, so comparing to -1 would be undefined behaviour
            size_t delimiter = line.find(' ');
            if (delimiter == string::npos) continue;

            int  asciiVal    = stoi(line.substr(0, delimiter));
            char c           = static_cast<char>(asciiVal);
            string code      = line.substr(delimiter + 1);
            codeToChar[code] = c;
        }
    }

    // Build the Huffman tree from a frequency map, then call generateCodes.
    //
    // Algorithm: seed a min-heap with one leaf per character, then repeatedly
    // pop the two lowest-frequency nodes, merge them under a new interior node,
    // and push the result back — until only the root remains.
    //
    // Edge case: a single unique character has no edges to traverse, so we
    // assign the code "0" manually instead of relying on generateCodes.
    void buildHuffTree(unordered_map<char, int> frequencyMap) {
        priority_queue<HuffNode*, vector<HuffNode*>, compareNodes> minHeap;

        for (const auto& [c, freq] : frequencyMap)
            minHeap.push(new HuffNode(c, freq));

        // Single-character file — assign "0" and return early
        if (minHeap.size() == 1) {
            root = minHeap.top();
            huffmanCodes[root->c] = "0";
            codeToChar["0"]       = root->c;
            return;
        }

        while (minHeap.size() > 1) {
            HuffNode* left  = minHeap.top(); minHeap.pop();
            HuffNode* right = minHeap.top(); minHeap.pop();
            minHeap.push(new HuffNode(left, right));
        }

        root = minHeap.top();
        generateCodes(root, "");
    }


public:

    HuffmanCode() : root(nullptr) {}
    ~HuffmanCode() { deleteTree(root); }

    // Compress inputFile and write the result to outputFile.
    //
    // Output layout:  [ header / code table ]  [ "====\n" ]  [ packed bitstream ]
    //
    // Bit-packing: each character's code is written bit by bit into a running byte.
    // When the byte is full it is flushed; any partial final byte is zero-padded.
    // The '~' EOF marker embedded at the end tells the decompressor where real
    // data ends so padding bits are never mistakenly decoded.
    bool compressFile(const string& inputFile,
                      const string& outputFile = "outCompressed.txt") {

        // Free any tree from a previous call — without this, reuse leaks memory
        deleteTree(root);
        root = nullptr;
        huffmanCodes.clear();
        codeToChar.clear();

        unordered_map<char, int> freqMap = buildFreqTable(inputFile);
        if (freqMap.empty()) return false;

        freqMap['~'] = 1; // inject EOF marker so it gets its own Huffman code

        buildHuffTree(freqMap);

        cout << "Huffman Codes:" << endl;
        for (const auto& [c, code] : huffmanCodes)
            cout << (int)(unsigned char)c << " : " << code << endl;

        ifstream inFile(inputFile, ios::binary);
        ofstream outFile(outputFile, ios::binary);
        if (!inFile.is_open() || !outFile.is_open()) return false;

        writeHuffHeader(outFile);

        // Pack bits into bytes: bitPos tracks the current bit position (7 = MSB).
        // When bitPos goes below 0 the byte is full — flush and reset.
        char currentByte = 0;
        int  bitPos      = 7;
        char byte;

        while (inFile.get(byte)) {
            const string& code = huffmanCodes[byte];
            for (char bit : code) {
                if (bit == '1') currentByte |= (1 << bitPos);
                if (--bitPos < 0) {
                    outFile.put(currentByte);
                    currentByte = 0;
                    bitPos      = 7;
                }
            }
        }

        // Write the EOF marker using the same bit-packing logic
        for (char bit : huffmanCodes['~']) {
            if (bit == '1') currentByte |= (1 << bitPos);
            if (--bitPos < 0) {
                outFile.put(currentByte);
                currentByte = 0;
                bitPos      = 7;
            }
        }

        if (bitPos != 7) outFile.put(currentByte); // flush any remaining partial byte

        return true;
    }

    // Decompress a file produced by compressFile.
    //
    // Reads the bitstream one bit at a time (MSB first), appending each to
    // curBits.  After every bit, curBits is checked against the code table.
    // Because Huffman codes are prefix-free, the first match is always correct.
    // Decoding stops as soon as the '~' EOF marker code is recognised.
    bool decompressFile(const string& inputFile,
                        const string& outputFile = "outDecompressed.txt") {

        ifstream inFile(inputFile, ios::binary);
        ofstream outFile(outputFile, ios::binary);
        if (!inFile.is_open() || !outFile.is_open()) return false;

        readHuffHeader(inFile);

        string curBits;
        char   byte;
        bool   done = false; // flag lets us exit both loops cleanly when '~' is hit

        while (!done && inFile.get(byte)) {
            for (int i = 7; i >= 0 && !done; --i) {
                // Extract bit i: shift right so it lands at position 0, then mask
                curBits += ((byte >> i) & 1) ? '1' : '0';

                auto it = codeToChar.find(curBits);
                if (it != codeToChar.end()) {
                    if (it->second == '~') {
                        done = true; // EOF marker — stop decoding, ignore padding bits
                    } else {
                        outFile.put(it->second);
                        curBits.clear();
                    }
                }
            }
        }

        return true;
    }
};