//! Stack-based virtual machine.

//! Implement a bytecode interpreter for a simple stack machine with

//! arithmetic, comparisons, jumps, local variables, and function calls.

//! Exercises: tagged unions, match, throw/try/catch, error handling,

//! records with slice fields, and complex interacting state.

const MAX_STACK: u32 = 64;

const MAX_CODE: u32 = 256;

const MAX_LOCALS: u32 = 16;

const MAX_FRAMES: u32 = 8;

/// Bytecode instructions.

union Op {

    /// Push an immediate value.

    Push(i32),

    /// Pop top two, push sum.

    Add,

    /// Pop top two, push difference (second - first).

    Sub,

    /// Pop top two, push product.

    Mul,

    /// Pop top two, push quotient.

    Div,

    /// Pop top two, push 1 if equal, 0 otherwise.

    Eq,

    /// Pop top two, push 1 if second < first, 0 otherwise.

    Lt,

    /// Pop top two, push 1 if second > first, 0 otherwise.

    Gt,

    /// Negate top of stack.

    Neg,

    /// Duplicate top of stack.

    Dup,

    /// Pop and discard top of stack.

    Pop,

    /// Store top of stack into local variable.

    Store(u32),

    /// Load local variable onto stack.

    Load(u32),

    /// Unconditional jump to instruction index.

    Jump(u32),

    /// Pop; if zero, jump to instruction index.

    JumpIfZero(u32),

    /// Call function at instruction index, saving return address.

    Call(u32),

    /// Return from function call.

    Ret,

    /// Halt execution, top of stack is result.

    Halt,

}

/// A call frame for function calls.

record Frame {

    returnAddr: u32,

    localBase: u32,

}

/// The VM state.

record VM {

    code: *[Op],

    codeLen: u32,

    stack: *mut [i32],

    sp: u32,

    locals: *mut [i32],

    frames: *mut [Frame],

    frameCount: u32,

    pc: u32,

}

/// VM error types.

union VmError {

    /// Stack underflow error.

    StackUnderflow,

    /// Stack overflow error.

    StackOverflow,

    /// Division by zero.

    DivByZero,

    /// Invalid instruction or PC.

    InvalidPC,

    /// Too many nested calls.

    CallOverflow,

}

/// Push a value onto the stack.

fn push(vm: *mut VM, value: i32) throws (VmError) {

    if vm.sp >= MAX_STACK {

        throw VmError::StackOverflow;

    }

    vm.stack[vm.sp] = value;

    vm.sp += 1;

}

/// Pop a value from the stack.

fn pop(vm: *mut VM) -> i32 throws (VmError) {

    if vm.sp == 0 {

        throw VmError::StackUnderflow;

    }

    vm.sp -= 1;

    return vm.stack[vm.sp];

}

/// Peek at the top of the stack without removing.

fn peek(vm: *VM) -> i32 throws (VmError) {

    if vm.sp == 0 {

        throw VmError::StackUnderflow;

    }

    return vm.stack[vm.sp - 1];

}

/// Execute the bytecode program.

fn execute(vm: *mut VM) -> i32 throws (VmError) {

    while vm.pc < vm.codeLen {

        let instr = vm.code[vm.pc];

        vm.pc += 1;

        match instr {

            case Op::Push(val) => {

                try push(vm, val);

            }

            case Op::Add => {

                let b = try pop(vm);

                let a = try pop(vm);

                try push(vm, a + b);

            }

            case Op::Sub => {

                let b = try pop(vm);

                let a = try pop(vm);

                try push(vm, a - b);

            }

            case Op::Mul => {

                let b = try pop(vm);

                let a = try pop(vm);

                try push(vm, a * b);

            }

            case Op::Div => {

                let b = try pop(vm);

                let a = try pop(vm);

                if b == 0 {

                    throw VmError::DivByZero;

                }

                try push(vm, a / b);

            }

            case Op::Eq => {

                let b = try pop(vm);

                let a = try pop(vm);

                let mut result: i32 = 0;

                if a == b {

                    result = 1;

                }

                try push(vm, result);

            }

            case Op::Lt => {

                let b = try pop(vm);

                let a = try pop(vm);

                let mut result: i32 = 0;

                if a < b {

                    result = 1;

                }

                try push(vm, result);

            }

            case Op::Gt => {

                let b = try pop(vm);

                let a = try pop(vm);

                let mut result: i32 = 0;

                if a > b {

                    result = 1;

                }

                try push(vm, result);

            }

            case Op::Neg => {

                let a = try pop(vm);

                try push(vm, 0 - a);

            }

            case Op::Dup => {

                let a = try peek(vm);

                try push(vm, a);

            }

            case Op::Pop => {

                try pop(vm);

            }

            case Op::Store(idx) => {

                let val = try pop(vm);

                let mut base: u32 = 0;

                if vm.frameCount > 0 {

                    base = vm.frames[vm.frameCount - 1].localBase;

                }

                vm.locals[base + idx] = val;

            }

            case Op::Load(idx) => {

                let mut base: u32 = 0;

                if vm.frameCount > 0 {

                    base = vm.frames[vm.frameCount - 1].localBase;

                }

                let val = vm.locals[base + idx];

                try push(vm, val);

            }

            case Op::Jump(target) => {

                vm.pc = target;

            }

            case Op::JumpIfZero(target) => {

                let cond = try pop(vm);

                if cond == 0 {

                    vm.pc = target;

                }

            }

            case Op::Call(target) => {

                if vm.frameCount >= MAX_FRAMES {

                    throw VmError::CallOverflow;

                }

                let mut base: u32 = 0;

                if vm.frameCount > 0 {

                    base = vm.frames[vm.frameCount - 1].localBase + MAX_LOCALS;

                }

                vm.frames[vm.frameCount] = Frame { returnAddr: vm.pc, localBase: base };

                vm.frameCount += 1;

                vm.pc = target;

            }

            case Op::Ret => {

                if vm.frameCount == 0 {

                    throw VmError::InvalidPC;

                }

                vm.frameCount -= 1;

                vm.pc = vm.frames[vm.frameCount].returnAddr;

            }

            case Op::Halt => {

                return try pop(vm);

            }

        }

        if vm.pc > vm.codeLen {

            throw VmError::InvalidPC;

        }

    }

    throw VmError::InvalidPC;

}

/// Helper to initialize VM and run a program.

fn runProgram(

    code: *[Op],

    codeLen: u32,

    stackBuf: *mut [i32],

    localsBuf: *mut [i32],

    framesBuf: *mut [Frame]

) -> i32 throws (VmError) {

    let mut vm = VM {

        code,

        codeLen,

        stack: stackBuf,

        sp: 0,

        locals: localsBuf,

        frames: framesBuf,

        frameCount: 0,

        pc: 0,

    };

    return try execute(&mut vm);

}

/// Test basic arithmetic: 3 + 4 * 2 = 11

fn testArith(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 8] = [Op::Halt; 8];

    code[0] = Op::Push(3);

    code[1] = Op::Push(4);

    code[2] = Op::Push(2);

    code[3] = Op::Mul;

    code[4] = Op::Add;

    code[5] = Op::Halt;

    let result: i32 = try! runProgram(&code[..], 6, stackBuf, localsBuf, framesBuf);

    assert result == 11;

    return 0;

}

/// Test local variables: x = 5, y = 7, push x + y.

fn testLocals(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 16] = [Op::Halt; 16];

    code[0] = Op::Push(5);

    code[1] = Op::Store(0);   // x = 5

    code[2] = Op::Push(7);

    code[3] = Op::Store(1);   // y = 7

    code[4] = Op::Load(0);    // push x

    code[5] = Op::Load(1);    // push y

    code[6] = Op::Add;        // x + y

    code[7] = Op::Halt;

    let result: i32 = try! runProgram(&code[..], 8, stackBuf, localsBuf, framesBuf);

    assert result == 12;

    return 0;

}

/// Test conditional jump: if 3 > 2 then push 42 else push 99.

fn testConditional(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 16] = [Op::Halt; 16];

    code[0] = Op::Push(3);

    code[1] = Op::Push(2);

    code[2] = Op::Gt;            // 3 > 2 => 1

    code[3] = Op::JumpIfZero(6); // if false, jump to 6

    code[4] = Op::Push(42);      // true branch

    code[5] = Op::Jump(7);       // skip false branch

    code[6] = Op::Push(99);      // false branch

    code[7] = Op::Halt;

    let result: i32 = try! runProgram(&code[..], 8, stackBuf, localsBuf, framesBuf);

    assert result == 42;

    return 0;

}

/// Test loop: sum 1..5 using jumps.

/// local[0] = counter (starts at 1), local[1] = sum (starts at 0).

fn testLoop(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 32] = [Op::Halt; 32];

    code[0] = Op::Push(1);

    code[1] = Op::Store(0);      // counter = 1

    code[2] = Op::Push(0);

    code[3] = Op::Store(1);      // sum = 0

    // Loop start (pc = 4):

    code[4] = Op::Load(0);       // push counter

    code[5] = Op::Push(6);

    code[6] = Op::Lt;            // counter < 6

    code[7] = Op::JumpIfZero(17); // if false, exit loop

    code[8] = Op::Load(1);       // push sum

    code[9] = Op::Load(0);       // push counter

    code[10] = Op::Add;          // sum + counter

    code[11] = Op::Store(1);     // sum = sum + counter

    code[12] = Op::Load(0);      // push counter

    code[13] = Op::Push(1);

    code[14] = Op::Add;          // counter + 1

    code[15] = Op::Store(0);     // counter = counter + 1

    code[16] = Op::Jump(4);      // back to loop start

    code[17] = Op::Load(1);      // push sum

    code[18] = Op::Halt;

    let result: i32 = try! runProgram(&code[..], 19, stackBuf, localsBuf, framesBuf);

    // sum = 1+2+3+4+5 = 15

    assert result == 15;

    return 0;

}

/// Test function call: call a function that computes n*2+1 for n=10.

fn testCall(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 32] = [Op::Halt; 32];

    // Main: push argument on stack, call function, halt.

    code[0] = Op::Push(10);      // push argument

    code[1] = Op::Call(5);       // call function at 5

    // After return, result is on stack.

    code[2] = Op::Halt;

    code[3] = Op::Halt;          // padding

    code[4] = Op::Halt;          // padding

    // Function at pc 5:

    code[5] = Op::Store(0);      // pop arg into local[0]

    code[6] = Op::Load(0);       // push local[0]

    code[7] = Op::Push(2);

    code[8] = Op::Mul;           // n * 2

    code[9] = Op::Push(1);

    code[10] = Op::Add;          // n * 2 + 1

    code[11] = Op::Ret;          // return (result on stack)

    let result: i32 = try! runProgram(&code[..], 12, stackBuf, localsBuf, framesBuf);

    // 10 * 2 + 1 = 21

    assert result == 21;

    return 0;

}

/// Test division by zero detection using try...catch with error binding.

fn testDivByZero(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 8] = [Op::Halt; 8];

    code[0] = Op::Push(42);

    code[1] = Op::Push(0);

    code[2] = Op::Div;

    code[3] = Op::Halt;

    let mut caught: i32 = 0;

    try runProgram(&code[..], 4, stackBuf, localsBuf, framesBuf) catch e {

        if e == VmError::DivByZero {

            caught = 1;

        } else {

            caught = 2;

        }

    };

    assert caught == 1;

    return 0;

}

/// Test negation and equality.

fn testNegAndEq(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 16] = [Op::Halt; 16];

    code[0] = Op::Push(5);

    code[1] = Op::Neg;           // -5

    code[2] = Op::Push(-5);

    code[3] = Op::Eq;            // -5 == -5 => 1

    code[4] = Op::Halt;

    let result: i32 = try! runProgram(&code[..], 5, stackBuf, localsBuf, framesBuf);

    assert result == 1;

    return 0;

}

/// Test factorial using recursive calls: fact(6) = 720.

fn testFactorial(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 32] = [Op::Halt; 32];

    // Main: push 6, call fact, halt.

    code[0] = Op::Push(6);

    code[1] = Op::Call(4);       // call fact at 4

    code[2] = Op::Halt;

    code[3] = Op::Halt;          // padding

    // fact(n) at pc 4:

    //   Store arg to local[0]

    //   If n <= 1, push 1, return

    //   Else push n, push fact(n-1), multiply, return

    code[4] = Op::Store(0);      // local[0] = n

    code[5] = Op::Load(0);       // push n

    code[6] = Op::Push(2);

    code[7] = Op::Lt;            // n < 2

    code[8] = Op::JumpIfZero(11); // if n >= 2, skip

    code[9] = Op::Push(1);       // base case: push 1

    code[10] = Op::Ret;

    // Recursive case:

    code[11] = Op::Load(0);      // push n

    code[12] = Op::Push(1);

    code[13] = Op::Sub;          // n - 1

    code[14] = Op::Call(4);      // fact(n-1)

    code[15] = Op::Load(0);      // push n

    code[16] = Op::Mul;          // fact(n-1) * n

    code[17] = Op::Ret;

    let result: i32 = try! runProgram(&code[..], 18, stackBuf, localsBuf, framesBuf);

    assert result == 720;

    return 0;

}

/// Test dup instruction.

fn testDup(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 8] = [Op::Halt; 8];

    code[0] = Op::Push(7);

    code[1] = Op::Dup;

    code[2] = Op::Add;           // 7 + 7 = 14

    code[3] = Op::Halt;

    let result: i32 = try! runProgram(&code[..], 4, stackBuf, localsBuf, framesBuf);

    assert result == 14;

    return 0;

}

/// Test stack underflow detection using try...catch with error binding.

fn testStackUnderflow(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 4] = [Op::Halt; 4];

    code[0] = Op::Add;  // nothing on stack

    code[1] = Op::Halt;

    let mut caught: i32 = 0;

    try runProgram(&code[..], 2, stackBuf, localsBuf, framesBuf) catch e {

        if e == VmError::StackUnderflow {

            caught = 1;

        } else {

            caught = 2;

        }

    };

    assert caught == 1;

    return 0;

}

/// Test that try...catch on success path does not execute catch block.

fn testSuccessNoCatch(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 4] = [Op::Halt; 4];

    code[0] = Op::Push(99);

    code[1] = Op::Halt;

    let mut caught: i32 = 0;

    let result: i32 = try runProgram(&code[..], 2, stackBuf, localsBuf, framesBuf) catch e {

        caught = 1;

        return 1;

    };

    // Catch block should not have run.

    assert caught == 0;

    // Should have the success value.

    assert result == 99;

    return 0;

}

/// Test call overflow detection by exhausting frames.

fn testCallOverflow(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 4] = [Op::Halt; 4];

    // Infinite recursion: function calls itself.

    code[0] = Op::Call(0);

    code[1] = Op::Halt;

    let mut caught: i32 = 0;

    try runProgram(&code[..], 2, stackBuf, localsBuf, framesBuf) catch e {

        if e == VmError::CallOverflow {

            caught = 1;

        } else {

            caught = 2;

        }

    };

    assert caught == 1;

    return 0;

}

/// Test that catch with no binding works (discard the error).

fn testCatchNoBinding(stackBuf: *mut [i32], localsBuf: *mut [i32], framesBuf: *mut [Frame]) -> i32 {

    let mut code: [Op; 4] = [Op::Halt; 4];

    code[0] = Op::Pop;  // underflow

    code[1] = Op::Halt;

    // Catch without binding - just swallow the error.

    try runProgram(&code[..], 2, stackBuf, localsBuf, framesBuf) catch {};

    return 0;

}

@default fn main() -> i32 {

    let mut stackBuf: [i32; 64] = [0; 64];

    let mut localsBuf: [i32; 128] = [0; 128];

    let mut framesBuf: [Frame; 8] = [Frame { returnAddr: 0, localBase: 0 }; 8];

    let r1 = testArith(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r1 != 0 {

        return 10 + r1;

    }

    let r2 = testLocals(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r2 != 0 {

        return 20 + r2;

    }

    let r3 = testConditional(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r3 != 0 {

        return 30 + r3;

    }

    let r4 = testLoop(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r4 != 0 {

        return 40 + r4;

    }

    let r5 = testCall(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r5 != 0 {

        return 50 + r5;

    }

    let r6 = testDivByZero(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r6 != 0 {

        return 60 + r6;

    }

    let r7 = testNegAndEq(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r7 != 0 {

        return 70 + r7;

    }

    let r8 = testFactorial(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r8 != 0 {

        return 80 + r8;

    }

    let r9 = testDup(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r9 != 0 {

        return 90 + r9;

    }

    let r10 = testStackUnderflow(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r10 != 0 {

        return 100 + r10;

    }

    let r11 = testSuccessNoCatch(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r11 != 0 {

        return 110 + r11;

    }

    let r12 = testCallOverflow(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r12 != 0 {

        return 120 + r12;

    }

    let r13 = testCatchNoBinding(&mut stackBuf[..], &mut localsBuf[..], &mut framesBuf[..]);

    if r13 != 0 {

        return 130 + r13;

    }

    return 0;

}