//! Huffman encoding.

//! Build a Huffman tree from character frequencies, generate prefix codes,

//! encode a message, decode it, and verify round-trip correctness.

const MAX_SYMBOLS: u32 = 32;

const MAX_NODES: u32 = 63;

const MAX_BITS: u32 = 512;

/// A node in the Huffman tree.

union HNodeKind {

    /// Leaf node with a symbol index.

    Leaf(u32),

    /// Interior node (no symbol).

    Interior,

}

record HNode {

    freq: u32,

    kind: HNodeKind,

    left: u32,

    right: u32,

}

const NIL: u32 = 0xFFFFFFFF;

record HuffState {

    nodes: *mut [HNode],

    nodeCount: u32,

    heap: *mut [u32],

    heapSize: u32,

    codeBits: *mut [u32],

    codeLen: *mut [u32],

    bitstream: *mut [u8],

    bitCount: u32,

}

fn newLeaf(s: *mut HuffState, freq: u32, symbol: u32) -> u32 {

    let idx: u32 = s.nodeCount;

    s.nodes[idx] = HNode { freq, kind: HNodeKind::Leaf(symbol), left: NIL, right: NIL };

    s.nodeCount += 1;

    return idx;

}

fn newInterior(s: *mut HuffState, freq: u32, left: u32, right: u32) -> u32 {

    let idx: u32 = s.nodeCount;

    s.nodes[idx] = HNode { freq, kind: HNodeKind::Interior, left, right };

    s.nodeCount += 1;

    return idx;

}

/// Get the symbol from a node, or nil if it's an interior node.

fn nodeSymbol(node: *HNode) -> ?u32 {

    match node.kind {

        case HNodeKind::Leaf(sym) => {

            return sym;

        }

        case HNodeKind::Interior => {

            return nil;

        }

    }

}

fn heapSwap(s: *mut HuffState, i: u32, j: u32) {

    let tmp: u32 = s.heap[i];

    s.heap[i] = s.heap[j];

    s.heap[j] = tmp;

}

fn heapFreq(s: *HuffState, i: u32) -> u32 {

    return s.nodes[s.heap[i]].freq;

}

fn siftUp(s: *mut HuffState, pos: u32) {

    let mut i: u32 = pos;

    while i > 0 {

        let parent: u32 = (i - 1) / 2;

        if heapFreq(s, i) < heapFreq(s, parent) {

            heapSwap(s, i, parent);

            i = parent;

        } else {

            return;

        }

    }

}

fn siftDown(s: *mut HuffState, pos: u32) {

    let mut i: u32 = pos;

    while true {

        let left: u32 = 2 * i + 1;

        let right: u32 = 2 * i + 2;

        let mut smallest: u32 = i;

        if left < s.heapSize and heapFreq(s, left) < heapFreq(s, smallest) {

            smallest = left;

        }

        if right < s.heapSize and heapFreq(s, right) < heapFreq(s, smallest) {

            smallest = right;

        }

        if smallest == i {

            return;

        }

        heapSwap(s, i, smallest);

        i = smallest;

    }

}

fn heapPush(s: *mut HuffState, nodeIdx: u32) {

    s.heap[s.heapSize] = nodeIdx;

    s.heapSize += 1;

    siftUp(s, s.heapSize - 1);

}

fn heapPop(s: *mut HuffState) -> u32 {

    let result: u32 = s.heap[0];

    s.heapSize -= 1;

    s.heap[0] = s.heap[s.heapSize];

    if s.heapSize > 0 {

        siftDown(s, 0);

    }

    return result;

}

fn buildTree(s: *mut HuffState, freqs: *[u32]) -> u32 {

    s.nodeCount = 0;

    s.heapSize = 0;

    for freq, sym in freqs {

        if freq > 0 {

            let idx: u32 = newLeaf(s, freq, sym);

            heapPush(s, idx);

        }

    }

    while s.heapSize > 1 {

        let left: u32 = heapPop(s);

        let right: u32 = heapPop(s);

        let combinedFreq: u32 = s.nodes[left].freq + s.nodes[right].freq;

        let parent: u32 = newInterior(s, combinedFreq, left, right);

        heapPush(s, parent);

    }

    return heapPop(s);

}

fn generateCodes(s: *mut HuffState, nodeIdx: u32, code: u32, depth: u32) {

    match s.nodes[nodeIdx].kind {

        case HNodeKind::Leaf(sym) => {

            s.codeBits[sym] = code;

            s.codeLen[sym] = depth;

        }

        case HNodeKind::Interior => {

            if s.nodes[nodeIdx].left != NIL {

                generateCodes(s, s.nodes[nodeIdx].left, code << 1, depth + 1);

            }

            if s.nodes[nodeIdx].right != NIL {

                generateCodes(s, s.nodes[nodeIdx].right, (code << 1) | 1, depth + 1);

            }

        }

    }

}

fn writeBit(s: *mut HuffState, bit: u32) {

    let byteIdx: u32 = s.bitCount / 8;

    let bitIdx: u32 = 7 - (s.bitCount % 8);

    if bit == 1 {

        s.bitstream[byteIdx] |= (1 as u8 << bitIdx as u8);

    }

    s.bitCount += 1;

}

fn readBit(s: *HuffState, pos: u32) -> u32 {

    let byteIdx: u32 = pos / 8;

    let bitIdx: u32 = 7 - (pos % 8);

    return (s.bitstream[byteIdx] >> bitIdx as u8) as u32 & 1;

}

fn encode(s: *mut HuffState, msg: *[u32]) {

    s.bitCount = 0;

    let mut i: u32 = 0;

    while i < MAX_BITS {

        s.bitstream[i] = 0;

        i += 1;

    }

    for sym in msg {

        let bits: u32 = s.codeBits[sym];

        let len: u32 = s.codeLen[sym];

        let mut b: u32 = 0;

        while b < len {

            let bit: u32 = (bits >> (len - 1 - b)) & 1;

            writeBit(s, bit);

            b += 1;

        }

    }

}

fn decode(s: *HuffState, root: u32, numSymbols: u32, out: *mut [u32]) -> u32 {

    let mut bitPos: u32 = 0;

    let mut decoded: u32 = 0;

    while decoded < numSymbols {

        let mut cur: u32 = root;

        // Walk the tree until we find a leaf.

        while nodeSymbol(&s.nodes[cur]) == nil {

            let bit: u32 = readBit(s, bitPos);

            bitPos += 1;

            if bit == 0 {

                cur = s.nodes[cur].left;

            } else {

                cur = s.nodes[cur].right;

            }

        }

        let sym = nodeSymbol(&s.nodes[cur]) else {

            return decoded;

        };

        out[decoded] = sym;

        decoded += 1;

    }

    return decoded;

}

fn resetCodes(s: *mut HuffState) {

    let mut i: u32 = 0;

    while i < MAX_SYMBOLS {

        s.codeBits[i] = 0;

        s.codeLen[i] = 0;

        i += 1;

    }

}

fn testBasic(s: *mut HuffState) -> i32 {

    let freqs: [u32; 5] = [5, 9, 12, 13, 16];

    let root: u32 = buildTree(s, &freqs[..]);

    assert s.nodes[root].freq == 55;

    resetCodes(s);

    generateCodes(s, root, 0, 0);

    // All symbols should have non-zero code lengths.

    let mut i: u32 = 0;

    while i < 5 {

        assert s.codeLen[i] != 0;

        i += 1;

    }

    // No two symbols at the same depth should have the same code.

    let mut a: u32 = 0;

    while a < 5 {

        let mut b: u32 = a + 1;

        while b < 5 {

            if s.codeLen[a] == s.codeLen[b] {

                assert s.codeBits[a] != s.codeBits[b];

            }

            b += 1;

        }

        a += 1;

    }

    return 0;

}

fn testRoundTrip(s: *mut HuffState) -> i32 {

    let freqs: [u32; 5] = [5, 9, 12, 13, 16];

    let root: u32 = buildTree(s, &freqs[..]);

    resetCodes(s);

    generateCodes(s, root, 0, 0);

    let msg: [u32; 16] = [4, 3, 2, 1, 0, 0, 1, 2, 3, 4, 2, 2, 3, 3, 4, 4];

    encode(s, &msg[..]);

    let mut decoded: [u32; 16] = [0; 16];

    let count: u32 = decode(s, root, 16, &mut decoded[..]);

    assert count == 16;

    // Verify round-trip.

    for expected, i in msg {

        if decoded[i] != expected {

            return i as i32 + 2;

        }

    }

    return 0;

}

fn testSkewed(s: *mut HuffState) -> i32 {

    let freqs: [u32; 6] = [100, 1, 1, 1, 1, 1];

    let root: u32 = buildTree(s, &freqs[..]);

    resetCodes(s);

    generateCodes(s, root, 0, 0);

    // The most frequent symbol should have the shortest code.

    let mut shortestLen: u32 = 100;

    let mut shortestSym: u32 = 0;

    let mut i: u32 = 0;

    while i < 6 {

        if s.codeLen[i] > 0 and s.codeLen[i] < shortestLen {

            shortestLen = s.codeLen[i];

            shortestSym = i;

        }

        i += 1;

    }

    assert shortestSym == 0;

    let msg: [u32; 12] = [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 2];

    encode(s, &msg[..]);

    let mut decoded: [u32; 12] = [0; 12];

    let count: u32 = decode(s, root, 12, &mut decoded[..]);

    assert count == 12;

    for expected, i in msg {

        assert decoded[i] == expected;

    }

    return 0;

}

fn testUniform(s: *mut HuffState) -> i32 {

    let freqs: [u32; 8] = [10, 10, 10, 10, 10, 10, 10, 10];

    let root: u32 = buildTree(s, &freqs[..]);

    assert s.nodes[root].freq == 80;

    resetCodes(s);

    generateCodes(s, root, 0, 0);

    // All 8 symbols should have code length 3 (8 = 2^3).

    let mut i: u32 = 0;

    while i < 8 {

        assert s.codeLen[i] == 3;

        i += 1;

    }

    let msg: [u32; 8] = [0, 1, 2, 3, 4, 5, 6, 7];

    encode(s, &msg[..]);

    assert s.bitCount == 24;

    let mut decoded: [u32; 8] = [0; 8];

    let count: u32 = decode(s, root, 8, &mut decoded[..]);

    assert count == 8;

    for expected, i in msg {

        assert decoded[i] == expected;

    }

    return 0;

}

@default fn main() -> i32 {

    let mut nodes: [HNode; 63] = [HNode { freq: 0, kind: HNodeKind::Interior, left: 0xFFFFFFFF, right: 0xFFFFFFFF }; 63];

    let mut heap: [u32; 63] = [0; 63];

    let mut codeBits: [u32; 32] = [0; 32];

    let mut codeLen: [u32; 32] = [0; 32];

    let mut bitstream: [u8; 512] = [0; 512];

    let mut s: HuffState = HuffState {

        nodes: &mut nodes[..],

        nodeCount: 0,

        heap: &mut heap[..],

        heapSize: 0,

        codeBits: &mut codeBits[..],

        codeLen: &mut codeLen[..],

        bitstream: &mut bitstream[..],

        bitCount: 0,

    };

    let r1: i32 = testBasic(&mut s);

    if r1 != 0 {

        return 10 + r1;

    }

    let r2: i32 = testRoundTrip(&mut s);

    if r2 != 0 {

        return 20 + r2;

    }

    let r3: i32 = testSkewed(&mut s);

    if r3 != 0 {

        return 30 + r3;

    }

    let r4: i32 = testUniform(&mut s);

    if r4 != 0 {

        return 40 + r4;

    }

    return 0;

}