(* Unoptimized Uxn virtual machine. *)

open Effect

type stack = Stack of { data : bytes; mutable sp : int }
type mode = Mode of { short : bool; keep : bool; mutable temp : int }

let stack_create () =
  let data = Bytes.create 256 in
  Bytes.unsafe_fill data 0 256 '\x00';
  Stack { data; sp = 0 }

let peek (Mode { short; keep; temp }) (Stack { data; sp }) : int =
  let amt = if short then 2 else 1 in
  let sp = if keep then (temp - amt) land 0xff else (sp - amt) land 0xff in
  if short then Util.get_uint16_wrap data sp else Bytes.get_uint8 data sp
[@@inline]

let pop (Mode m as m' : mode) (Stack s' as s : stack) =
  let res = peek m' s in
  let amt = if m.short then 2 else 1 in
  if m.keep then m.temp <- (m.temp - amt) land 0xff
  else s'.sp <- (s'.sp - amt) land 0xff;
  res
[@@inline]

let push (Mode ({ short; keep; _ } as m) : mode) (Stack s : stack) (v : int) =
  if short then Util.set_uint16_wrap s.data s.sp (v land 0xffff)
  else Bytes.set_uint8 s.data s.sp (v land 0xff);
  let amt = if m.short then 2 else 1 in
  if keep then m.temp <- (m.temp + amt) land 0xff;
  s.sp <- (s.sp + amt) land 0xff
[@@inline]

let pushbyte (Mode m) s v =
  let m' = Mode { m with short = false } in
  push m' s v;
  let (Mode { temp; _ }) = m' in
  m.temp <- temp
[@@inline]

let popbyte (Mode m) s =
  let m' = Mode { m with short = false } in
  let r = pop m' s in
  let (Mode { temp; _ }) = m' in
  m.temp <- temp;
  r
[@@inline]

let popshort (Mode m) s =
  let m' = Mode { m with short = true } in
  let r = pop m' s in
  let (Mode { temp; _ }) = m' in
  m.temp <- temp;
  r
[@@inline]

let pop1 s = pop (Mode { short = false; keep = false; temp = 0 }) s [@@inline]

let push1 s v = push (Mode { short = false; keep = false; temp = 0 }) s v
[@@inline]

let push2 s v = push (Mode { short = true; keep = false; temp = 0 }) s v
[@@inline]

type machine =
  | Machine of { data : bytes; dev : bytes; stack : stack; callstack : stack }

type _ Effect.t +=
  | BRK : int Effect.t
  | DEI : int -> int Effect.t
  | DEI2 : int -> int Effect.t
  | DEO : (int * int) -> unit Effect.t

let ram (Machine { data; _ }) = data
let dev (Machine { dev; _ }) = dev
let wst (Machine { stack; _ }) = stack
let rst (Machine { callstack; _ }) = callstack

let create code =
  let data = Bytes.create 65536 in
  let dev = Bytes.create 256 in
  Bytes.unsafe_fill data 0 65536 '\x00';
  Bytes.unsafe_fill dev 0 256 '\x00';
  Bytes.blit_string code 0 data 0x100 (String.length code);
  Machine { data; dev; stack = stack_create (); callstack = stack_create () }

let dispatch (Machine m) (pc : int) =
  let pc = ref pc in

  while true do
    let op = Bytes.get_uint8 m.data (!pc land 0xffff) in
    pc := (!pc + 1) land 0xffff;

    let short = op land 0x20 <> 0 in
    let keep = op land 0x80 <> 0 in
    let return = op land 0x40 <> 0 in
    let opcode = op land 0x1f in

    match op with
    | 0x00 -> pc := perform BRK
    | 0x20 (* JCI *) ->
        let cond = pop1 m.stack in
        let addr = Util.get_int16_wrap ~wrap:0xffff m.data !pc in
        if cond != 0 then pc := !pc + addr + 2 else pc := !pc + 2
    | 0x40 (* JMI *) ->
        let addr = Util.get_int16_wrap ~wrap:0xffff m.data !pc in
        pc := !pc + addr + 2
    | 0x60 (* JSI *) ->
        let addr = Util.get_int16_wrap ~wrap:0xffff m.data !pc in
        push2 m.callstack (!pc + 2);
        pc := !pc + addr + 2
    | 0x80 (* LIT *) ->
        let lit = Bytes.get_uint8 m.data !pc in
        push1 m.stack lit;
        pc := !pc + 1
    | 0xa0 (* LIT2 *) ->
        let lit = Util.get_uint16_wrap ~wrap:0xffff m.data !pc in
        push2 m.stack lit;
        pc := !pc + 2
    | 0xc0 (* LITr *) ->
        let lit = Bytes.get_uint8 m.data !pc in
        push1 m.callstack lit;
        pc := !pc + 1
    | 0xe0 (* LIT2r *) ->
        let lit = Util.get_uint16_wrap ~wrap:0xffff m.data !pc in
        push2 m.callstack lit;
        pc := !pc + 2
    | _ -> begin
        let stk = if return then m.callstack else m.stack in
        let stk' = if return then m.stack else m.callstack in
        let mode =
          Mode { short; keep; temp = (match stk with Stack { sp; _ } -> sp) }
        in
        match[@warning "-8"] opcode with
        | 0x01 (* INC *) ->
            let r = pop mode stk in
            push mode stk (r + 1)
        | 0x02 (* POP *) -> ignore (pop mode stk)
        | 0x03 (* NIP *) ->
            let b = pop mode stk in
            ignore (pop mode stk);
            push mode stk b
        | 0x04 (* SWP *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            push mode stk b;
            push mode stk a
        | 0x05 (* ROT *) ->
            let c = pop mode stk in
            let b = pop mode stk in
            let a = pop mode stk in
            push mode stk b;
            push mode stk c;
            push mode stk a
        | 0x06 (* DUP *) ->
            let a = pop mode stk in
            push mode stk a;
            push mode stk a
        | 0x07 (* OVR *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            push mode stk a;
            push mode stk b;
            push mode stk a
        | 0x08 (* EQU *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            pushbyte mode stk (if a = b then 1 else 0)
        | 0x09 (* NEQ *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            pushbyte mode stk (if a != b then 1 else 0)
        | 0x0a (* GTH *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            pushbyte mode stk (if a > b then 1 else 0)
        | 0x0b (* LTH *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            pushbyte mode stk (if a < b then 1 else 0)
        | 0x0c (* JMP *) ->
            let addr = pop mode stk in
            if short then pc := addr else pc := !pc + Util.uint8_to_int8 addr
        | 0x0d (* JCN *) ->
            let addr = pop mode stk in
            let cond = popbyte mode stk in
            if cond != 0 then
              if short then pc := addr else pc := !pc + Util.uint8_to_int8 addr
        | 0x0e (* JSR *) ->
            push2 m.callstack !pc;
            let addr = pop mode stk in
            if short then pc := addr else pc := !pc + Util.uint8_to_int8 addr
        | 0x0f (* STH *) ->
            let a = pop mode stk in
            push mode stk' a
        | 0x10 (* LDZ *) ->
            let addr = popbyte mode stk in
            push mode stk
              (if short then Util.get_uint16_wrap m.data addr
               else Bytes.get_uint8 m.data addr)
        | 0x11 (* STZ *) ->
            let addr = popbyte mode stk in
            let v = pop mode stk in
            if short then Util.set_uint16_wrap m.data addr v
            else Bytes.set_uint8 m.data addr v
        | 0x12 (* LDR *) ->
            let addr = !pc + Util.uint8_to_int8 (popbyte mode stk) in
            push mode stk
              (if short then Util.get_uint16_wrap ~wrap:0xffff m.data addr
               else Bytes.get_uint8 m.data addr)
        | 0x13 (* STR *) ->
            let addr = !pc + Util.uint8_to_int8 (popbyte mode stk) in
            let v = pop mode stk in
            if short then Util.set_uint16_wrap ~wrap:0xffff m.data addr v
            else Bytes.set_uint8 m.data addr v
        | 0x14 (* LDA *) ->
            let addr = popshort mode stk in
            push mode stk
              (if short then Util.get_uint16_wrap ~wrap:0xffff m.data addr
               else Bytes.get_uint8 m.data addr)
        | 0x15 (* STA *) ->
            let addr = popshort mode stk in
            let v = pop mode stk in
            if short then Util.set_uint16_wrap ~wrap:0xffff m.data addr v
            else Bytes.set_uint8 m.data addr v
        | 0x16 (* DEI *) ->
            let port = popbyte mode stk in
            push mode stk
              (if short then perform (DEI2 port) else perform (DEI port))
        | 0x17 (* DEO *) ->
            let port = popbyte mode stk in
            let value = pop mode stk in
            if short then Util.set_uint16_wrap m.dev port value
            else Bytes.set_uint8 m.dev port value;
            perform (DEO (port, value))
        | 0x18 (* ADD *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            push mode stk (a + b)
        | 0x19 (* SUB *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            push mode stk (a - b)
        | 0x1a (* MUL *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            push mode stk (a * b)
        | 0x1b (* DIV *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            push mode stk (if b = 0 then 0 else a / b)
        | 0x1c (* AND *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            push mode stk (a land b)
        | 0x1d (* ORA *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            push mode stk (a lor b)
        | 0x1e (* EOR *) ->
            let b = pop mode stk in
            let a = pop mode stk in
            push mode stk (a lxor b)
        | 0x1f (* SFT *) ->
            let sft = popbyte mode stk in
            let value = pop mode stk in
            push mode stk ((value lsr (sft land 0xf)) lsl sft lsr 4)
      end
  done