//! returns: 0
//! LZW compression and decompression.
//! Implement the Lempel-Ziv-Welch algorithm with a fixed-size dictionary.
//! Encode a byte buffer, decode it, and verify perfect round-trip.

const MAX_DICT: u32 = 512;
const INIT_DICT: u32 = 258;
const CLEAR_CODE: u32 = 256;
const EOI_CODE: u32 = 257;

record DictEntry {
    prefix: u32,
    suffix: u8,
}

record LzwState {
    encDict: *mut [DictEntry],
    encDictSize: u32,
    decDict: *mut [DictEntry],
    decDictSize: u32,
    encoded: *mut [u32],
    encLen: u32,
    decoded: *mut [u8],
    decLen: u32,
    temp: *mut [u8],
}

fn initEncDict(s: *mut LzwState) {
    s.encDictSize = INIT_DICT;
    let mut i: u32 = 0;
    while i < 256 {
        s.encDict[i] = DictEntry { prefix: 0xFFFF, suffix: i as u8 };
        i += 1;
    }
}

fn initDecDict(s: *mut LzwState) {
    s.decDictSize = INIT_DICT;
    let mut i: u32 = 0;
    while i < 256 {
        s.decDict[i] = DictEntry { prefix: 0xFFFF, suffix: i as u8 };
        i += 1;
    }
}

fn dictLookup(s: *LzwState, prefix: u32, suffix: u8) -> u32 {
    let mut i: u32 = 0;
    while i < s.encDictSize {
        if s.encDict[i].prefix == prefix and s.encDict[i].suffix == suffix {
            return i;
        }
        i += 1;
    }
    return 0xFFFFFFFF;
}

fn dictAdd(s: *mut LzwState, prefix: u32, suffix: u8) {
    if s.encDictSize < MAX_DICT {
        s.encDict[s.encDictSize] = DictEntry { prefix, suffix };
        s.encDictSize += 1;
    }
}

fn emitCode(s: *mut LzwState, code: u32) {
    s.encoded[s.encLen] = code;
    s.encLen += 1;
}

fn encode(s: *mut LzwState, data: *[u8]) {
    initEncDict(s);
    s.encLen = 0;

    if data.len == 0 {
        emitCode(s, EOI_CODE);
        return;
    }

    let mut w: u32 = data[0] as u32;
    let mut i: u32 = 1;

    while i < data.len {
        let c: u8 = data[i];
        let wc: u32 = dictLookup(s, w, c);
        if wc != 0xFFFFFFFF {
            w = wc;
        } else {
            emitCode(s, w);
            dictAdd(s, w, c);
            w = c as u32;
        }
        i += 1;
    }
    emitCode(s, w);
    emitCode(s, EOI_CODE);
}

fn decodeString(s: *mut LzwState, code: u32) -> u32 {
    let mut len: u32 = 0;
    let mut c: u32 = code;
    while c != 0xFFFF and c < s.decDictSize {
        s.temp[len] = s.decDict[c].suffix;
        len += 1;
        c = s.decDict[c].prefix;
    }
    let mut a: u32 = 0;
    let mut b: u32 = len - 1;
    while a < b {
        let tmp: u8 = s.temp[a];
        s.temp[a] = s.temp[b];
        s.temp[b] = tmp;
        a += 1;
        b -= 1;
    }
    return len;
}

fn firstByte(s: *LzwState, code: u32) -> u8 {
    let mut c: u32 = code;
    while s.decDict[c].prefix != 0xFFFF and s.decDict[c].prefix < s.decDictSize {
        c = s.decDict[c].prefix;
    }
    return s.decDict[c].suffix;
}

fn decDictAdd(s: *mut LzwState, prefix: u32, suffix: u8) {
    if s.decDictSize < MAX_DICT {
        s.decDict[s.decDictSize] = DictEntry { prefix, suffix };
        s.decDictSize += 1;
    }
}

fn outputByte(s: *mut LzwState, b: u8) {
    s.decoded[s.decLen] = b;
    s.decLen += 1;
}

fn decode(s: *mut LzwState) {
    initDecDict(s);
    s.decLen = 0;

    if s.encLen == 0 {
        return;
    }

    let mut pos: u32 = 0;
    let mut code: u32 = s.encoded[pos];
    pos += 1;

    if code == EOI_CODE {
        return;
    }

    outputByte(s, code as u8);
    let mut prevCode: u32 = code;

    while pos < s.encLen {
        code = s.encoded[pos];
        pos += 1;

        if code == EOI_CODE {
            return;
        }

        if code < s.decDictSize {
            let len: u32 = decodeString(s, code);
            let mut i: u32 = 0;
            while i < len {
                outputByte(s, s.temp[i]);
                i += 1;
            }
            decDictAdd(s, prevCode, s.temp[0]);
        } else {
            let fb: u8 = firstByte(s, prevCode);
            let len: u32 = decodeString(s, prevCode);
            let mut i: u32 = 0;
            while i < len {
                outputByte(s, s.temp[i]);
                i += 1;
            }
            outputByte(s, fb);
            decDictAdd(s, prevCode, fb);
        }
        prevCode = code;
    }
}

fn testSimple(s: *mut LzwState) -> i32 {
    let data: *[u8] = "ABABABABAB";
    encode(s, data);

    assert s.encoded[s.encLen - 1] == EOI_CODE;

    decode(s);
    assert s.decLen == 10;
    let mut i: u32 = 0;
    while i < 10 {
        assert s.decoded[i] == data[i];
        i += 1;
    }
    return 0;
}

fn testDistinct(s: *mut LzwState) -> i32 {
    let data: [u8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];
    encode(s, &data[..]);

    decode(s);
    assert s.decLen == 16;
    let mut i: u32 = 0;
    while i < 16 {
        assert s.decoded[i] == data[i];
        i += 1;
    }
    return 0;
}

fn testRepetitive(s: *mut LzwState) -> i32 {
    let mut data: [u8; 64] = [65; 64];
    encode(s, &data[..]);

    assert s.encLen < 32;

    decode(s);
    assert s.decLen == 64;
    let mut i: u32 = 0;
    while i < 64 {
        assert s.decoded[i] == 65;
        i += 1;
    }
    return 0;
}

fn testMixed(s: *mut LzwState) -> i32 {
    let data: *[u8] = "TOBEORNOTTOBEORTOBEORNOT";
    encode(s, data);

    decode(s);
    assert s.decLen == 24;
    let mut i: u32 = 0;
    while i < 24 {
        assert s.decoded[i] == data[i];
        i += 1;
    }

    assert s.encLen - 1 < 24;

    return 0;
}

fn testEmpty(s: *mut LzwState) -> i32 {
    encode(s, &[]);
    assert s.encLen == 1;
    assert s.encoded[0] == EOI_CODE;

    decode(s);
    assert s.decLen == 0;
    return 0;
}

fn testSingle(s: *mut LzwState) -> i32 {
    let data: *[u8] = "*";
    encode(s, data);

    decode(s);
    assert s.decLen == 1;
    assert s.decoded[0] == 42;
    return 0;
}

@default fn main() -> i32 {
    let mut encDict: [DictEntry; 512] = [DictEntry { prefix: 0xFFFF, suffix: 0 }; 512];
    let mut decDict: [DictEntry; 512] = [DictEntry { prefix: 0xFFFF, suffix: 0 }; 512];
    let mut encoded: [u32; 512] = [0; 512];
    let mut decoded: [u8; 256] = [0; 256];
    let mut temp: [u8; 256] = [0; 256];

    let mut s: LzwState = LzwState {
        encDict: &mut encDict[..],
        encDictSize: 0,
        decDict: &mut decDict[..],
        decDictSize: 0,
        encoded: &mut encoded[..],
        encLen: 0,
        decoded: &mut decoded[..],
        decLen: 0,
        temp: &mut temp[..],
    };

    let r1: i32 = testSimple(&mut s);
    if r1 != 0 { return 10 + r1; }

    let r2: i32 = testDistinct(&mut s);
    if r2 != 0 { return 20 + r2; }

    let r3: i32 = testRepetitive(&mut s);
    if r3 != 0 { return 30 + r3; }

    let r4: i32 = testMixed(&mut s);
    if r4 != 0 { return 40 + r4; }

    let r5: i32 = testEmpty(&mut s);
    if r5 != 0 { return 50 + r5; }

    let r6: i32 = testSingle(&mut s);
    if r6 != 0 { return 60 + r6; }

    return 0;
}