//! returns: 0
//! NFA-based regex matcher.
//! Implement a simple regular expression engine using Thompson's NFA
//! construction. Supports: literal characters, '.', '*', '+', '?',
//! and concatenation.

const MAX_STATES: u32 = 128;
const MAX_TRANSITIONS: u32 = 256;
const NIL: u32 = 0xFFFFFFFF;

const TRANS_CHAR: u32 = 0;
const TRANS_EPSILON: u32 = 1;
const TRANS_DOT: u32 = 2;

record Trans {
    kind: u32,
    ch: u8,
    to: u32,
}

record Frag {
    start: u32,
    endState: u32,
}

record NfaState {
    trans: *mut [Trans],
    transCount: u32,
    stateFirst: *mut [u32],
    transNext: *mut [u32],
    stateCount: u32,
    acceptState: u32,
    current: *mut [u32],
    nextSet: *mut [u32],
    closure: *mut [u32],
    fragStack: *mut [Frag],
    fragTop: u32,
}

fn newState(nfa: *mut NfaState) -> u32 {
    let s: u32 = nfa.stateCount;
    nfa.stateFirst[s] = NIL;
    nfa.stateCount += 1;
    return s;
}

fn addTrans(nfa: *mut NfaState, from: u32, kind: u32, ch: u8, to: u32) {
    let idx: u32 = nfa.transCount;
    nfa.trans[idx] = Trans { kind, ch, to };
    nfa.transNext[idx] = nfa.stateFirst[from];
    nfa.stateFirst[from] = idx;
    nfa.transCount += 1;
}

fn pushFrag(nfa: *mut NfaState, f: Frag) {
    nfa.fragStack[nfa.fragTop] = f;
    nfa.fragTop += 1;
}

fn popFrag(nfa: *mut NfaState) -> Frag {
    nfa.fragTop -= 1;
    return nfa.fragStack[nfa.fragTop];
}

fn setEmpty(s: *mut [u32]) {
    s[0] = 0;
    s[1] = 0;
    s[2] = 0;
    s[3] = 0;
}

fn setAdd(s: *mut [u32], bit: u32) {
    let word: u32 = bit / 32;
    let pos: u32 = bit % 32;
    s[word] |= (1 << pos);
}

fn setHas(s: *[u32], bit: u32) -> bool {
    let word: u32 = bit / 32;
    let pos: u32 = bit % 32;
    return (s[word] >> pos) & 1 == 1;
}

fn setIsEmpty(s: *[u32]) -> bool {
    return s[0] == 0 and s[1] == 0 and s[2] == 0 and s[3] == 0;
}

fn setCopy(dst: *mut [u32], src: *[u32]) {
    dst[0] = src[0];
    dst[1] = src[1];
    dst[2] = src[2];
    dst[3] = src[3];
}

fn epsilonClosure(nfa: *mut NfaState, states: *mut [u32]) {
    setCopy(nfa.closure, states);

    let mut changed: bool = true;
    while changed {
        changed = false;
        let mut s: u32 = 0;
        while s < nfa.stateCount {
            if setHas(nfa.closure, s) {
                let mut t: u32 = nfa.stateFirst[s];
                while t != NIL {
                    if nfa.trans[t].kind == TRANS_EPSILON {
                        if not setHas(nfa.closure, nfa.trans[t].to) {
                            setAdd(nfa.closure, nfa.trans[t].to);
                            changed = true;
                        }
                    }
                    t = nfa.transNext[t];
                }
            }
            s += 1;
        }
    }

    setCopy(states, nfa.closure);
}

fn resetNFA(nfa: *mut NfaState) {
    nfa.stateCount = 0;
    nfa.transCount = 0;
    nfa.fragTop = 0;
    let mut i: u32 = 0;
    while i < MAX_STATES {
        nfa.stateFirst[i] = NIL;
        i += 1;
    }
    i = 0;
    while i < MAX_TRANSITIONS {
        nfa.transNext[i] = NIL;
        i += 1;
    }
}

fn compile(nfa: *mut NfaState, pattern: *[u8]) -> u32 {
    resetNFA(nfa);

    let mut i: u32 = 0;
    while i < pattern.len {
        let ch: u8 = pattern[i];

        if i + 1 < pattern.len {
            let nextCh: u8 = pattern[i + 1];

            if nextCh == 42 {
                let s: u32 = newState(nfa);
                let e: u32 = newState(nfa);
                let body: u32 = newState(nfa);

                if ch == 46 {
                    addTrans(nfa, body, TRANS_DOT, 0, body);
                } else {
                    addTrans(nfa, body, TRANS_CHAR, ch, body);
                }
                addTrans(nfa, s, TRANS_EPSILON, 0, body);
                addTrans(nfa, s, TRANS_EPSILON, 0, e);
                addTrans(nfa, body, TRANS_EPSILON, 0, e);

                pushFrag(nfa, Frag { start: s, endState: e });
                i += 2;
                continue;
            }
            if nextCh == 43 {
                let s: u32 = newState(nfa);
                let m: u32 = newState(nfa);
                let e: u32 = newState(nfa);

                if ch == 46 {
                    addTrans(nfa, s, TRANS_DOT, 0, m);
                    addTrans(nfa, m, TRANS_DOT, 0, m);
                } else {
                    addTrans(nfa, s, TRANS_CHAR, ch, m);
                    addTrans(nfa, m, TRANS_CHAR, ch, m);
                }
                addTrans(nfa, m, TRANS_EPSILON, 0, e);

                pushFrag(nfa, Frag { start: s, endState: e });
                i += 2;
                continue;
            }
            if nextCh == 63 {
                let s: u32 = newState(nfa);
                let m: u32 = newState(nfa);
                let e: u32 = newState(nfa);

                if ch == 46 {
                    addTrans(nfa, s, TRANS_DOT, 0, m);
                } else {
                    addTrans(nfa, s, TRANS_CHAR, ch, m);
                }
                addTrans(nfa, s, TRANS_EPSILON, 0, e);
                addTrans(nfa, m, TRANS_EPSILON, 0, e);

                pushFrag(nfa, Frag { start: s, endState: e });
                i += 2;
                continue;
            }
        }

        let s: u32 = newState(nfa);
        let e: u32 = newState(nfa);
        if ch == 46 {
            addTrans(nfa, s, TRANS_DOT, 0, e);
        } else {
            addTrans(nfa, s, TRANS_CHAR, ch, e);
        }
        pushFrag(nfa, Frag { start: s, endState: e });
        i += 1;
    }

    if nfa.fragTop == 0 {
        let s: u32 = newState(nfa);
        nfa.acceptState = s;
        return s;
    }

    let numFrags: u32 = nfa.fragTop;
    let mut frags: [Frag; 64] = [Frag { start: 0, endState: 0 }; 64];
    let mut k: u32 = 0;
    while k < numFrags {
        frags[numFrags - 1 - k] = popFrag(nfa);
        k += 1;
    }

    let mut j: u32 = 0;
    while j < numFrags - 1 {
        addTrans(nfa, frags[j].endState, TRANS_EPSILON, 0, frags[j + 1].start);
        j += 1;
    }

    nfa.acceptState = frags[numFrags - 1].endState;
    return frags[0].start;
}

fn nfaMatches(nfa: *mut NfaState, start: u32, input: *[u8]) -> bool {
    setEmpty(nfa.current);
    setAdd(nfa.current, start);
    epsilonClosure(nfa, nfa.current);

    let mut i: u32 = 0;
    while i < input.len {
        let ch: u8 = input[i];
        setEmpty(nfa.nextSet);

        let mut s: u32 = 0;
        while s < nfa.stateCount {
            if setHas(nfa.current, s) {
                let mut t: u32 = nfa.stateFirst[s];
                while t != NIL {
                    if nfa.trans[t].kind == TRANS_CHAR and nfa.trans[t].ch == ch {
                        setAdd(nfa.nextSet, nfa.trans[t].to);
                    } else if nfa.trans[t].kind == TRANS_DOT {
                        setAdd(nfa.nextSet, nfa.trans[t].to);
                    }
                    t = nfa.transNext[t];
                }
            }
            s += 1;
        }

        epsilonClosure(nfa, nfa.nextSet);
        setCopy(nfa.current, nfa.nextSet);

        if setIsEmpty(nfa.current) {
            return false;
        }
        i += 1;
    }

    return setHas(nfa.current, nfa.acceptState);
}

fn testLiteral(nfa: *mut NfaState) -> i32 {
    let start: u32 = compile(nfa, "abc");

    assert nfaMatches(nfa, start, "abc");
    if nfaMatches(nfa, start, "ab") { return 2; }
    if nfaMatches(nfa, start, "abcd") { return 3; }
    if nfaMatches(nfa, start, "abd") { return 4; }

    return 0;
}

fn testStar(nfa: *mut NfaState) -> i32 {
    let start: u32 = compile(nfa, "a*");

    assert nfaMatches(nfa, start, "");
    assert nfaMatches(nfa, start, "a");
    assert nfaMatches(nfa, start, "aaa");
    if nfaMatches(nfa, start, "b") { return 4; }

    return 0;
}

fn testPlus(nfa: *mut NfaState) -> i32 {
    let start: u32 = compile(nfa, "a+");

    if nfaMatches(nfa, start, "") { return 1; }
    assert nfaMatches(nfa, start, "a");
    assert nfaMatches(nfa, start, "aaaaa");

    return 0;
}

fn testQuestion(nfa: *mut NfaState) -> i32 {
    let start: u32 = compile(nfa, "a?");

    assert nfaMatches(nfa, start, "");
    assert nfaMatches(nfa, start, "a");
    if nfaMatches(nfa, start, "aa") { return 3; }

    return 0;
}

fn testDot(nfa: *mut NfaState) -> i32 {
    let start: u32 = compile(nfa, "..");

    assert nfaMatches(nfa, start, "ab");
    assert nfaMatches(nfa, start, "zz");
    if nfaMatches(nfa, start, "a") { return 3; }
    if nfaMatches(nfa, start, "abc") { return 4; }

    return 0;
}

fn testComplex(nfa: *mut NfaState) -> i32 {
    let start: u32 = compile(nfa, "ab*c");

    assert nfaMatches(nfa, start, "ac");
    assert nfaMatches(nfa, start, "abc");
    assert nfaMatches(nfa, start, "abbc");
    assert nfaMatches(nfa, start, "abbbc");
    if nfaMatches(nfa, start, "a") { return 5; }
    if nfaMatches(nfa, start, "adc") { return 6; }

    return 0;
}

fn testDotStar(nfa: *mut NfaState) -> i32 {
    let start: u32 = compile(nfa, ".*");

    assert nfaMatches(nfa, start, "");
    assert nfaMatches(nfa, start, "hello");
    assert nfaMatches(nfa, start, "x");

    return 0;
}

@default fn main() -> i32 {
    let mut trans: [Trans; 256] = [Trans { kind: 0, ch: 0, to: 0 }; 256];
    let mut stateFirst: [u32; 128] = [0xFFFFFFFF; 128];
    let mut transNext: [u32; 256] = [0xFFFFFFFF; 256];
    let mut current: [u32; 4] = [0; 4];
    let mut nextSet: [u32; 4] = [0; 4];
    let mut closureBuf: [u32; 4] = [0; 4];
    let mut fragStack: [Frag; 64] = [Frag { start: 0, endState: 0 }; 64];

    let mut nfa: NfaState = NfaState {
        trans: &mut trans[..],
        transCount: 0,
        stateFirst: &mut stateFirst[..],
        transNext: &mut transNext[..],
        stateCount: 0,
        acceptState: 0,
        current: &mut current[..],
        nextSet: &mut nextSet[..],
        closure: &mut closureBuf[..],
        fragStack: &mut fragStack[..],
        fragTop: 0,
    };

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

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

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

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

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

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

    let r7: i32 = testDotStar(&mut nfa);
    if r7 != 0 { return 70 + r7; }

    return 0;
}
