(define-module (scmvm vm)
  #:use-module ((scheme base)
                #:select (read-u8 read-bytevector))
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-26)
  #:export (make-vm run-vm vm-memory-ref vm-memory-set!))

;;; Data Structures
(define *stack-size* 1000)
(define *memory-size* 2048)

(define* (make-stack #:optional (stack-size *stack-size*))
  "Make a new stack, optionally setting the size"
  (define the-stack (make-vector stack-size))
  (define top 0)
  (lambda (x)
    (case x
      [(push)
       (lambda (v)
         (if (>= top stack-size)
             (error "stack overflow")
             (begin
               (vector-set! the-stack top v)
               (set! top (1+ top)))))]
      [(pop)
       (lambda ()
         (if (zero? top)
             (error "pop empty stack")
             (begin
               (set! top (1- top))
               (vector-ref the-stack top))))]
      [(peek)
       (lambda ()
         (if (zero? top)
             (error "peek empty stack")
             (vector-ref the-stack (1- top))))]
      [(swap)
       (lambda ()
         (if (< (vector-length the-stack) 2)
             (error "no value to swap")
             (let ([a (vector-ref the-stack (- top 2))]
                   [b (vector-ref the-stack (- top 1))])
               (vector-set! the-stack (- top 2) b)
               (vector-set! the-stack (- top 1) a))))]
      [(ref)
       (lambda (k)
         (vector-ref the-stack k))]
      [else (error "stack dispatch unknown value")])))

(define (push stack v)
  ((stack 'push) v))

(define (pop stack)
  ((stack 'pop)))

(define (peek stack)
  ((stack 'peek)))

(define (swap stack)
  ((stack 'swap)))

(define (stack-ref stack k)
  ((stack 'ref) k))

(define* (make-ram #:optional (memory-size *memory-size*))
  (make-vector memory-size))

(define (ram-ref ram k)
  (vector-ref ram k))

(define (ram-set! ram k v)
  (vector-set! ram k v))

;;; IO
(define (read-word)
  "Read the next 32-bit value from (current-input-port)"
  (let ([bv (read-bytevector 4)])
    (bytevector-s32-ref bv 0 (native-endianness))))

;;; Program execution
(define *opcodes*
  '((#x01 . push)
    (#x02 . pop)
    (#x03 . store)
    (#x04 . load)
    (#x05 . +)
    (#x06 . -)
    (#x07 . and)
    (#x08 . or)
    (#x09 . nand)
    (#x0a . nor)
    (#x0b . xor)
    (#x0c . if)
    (#x0d . call)
    (#x0e . return)
    (#x0f . >R)
    (#x10 . R>)
    (#x11 . =)
    (#x12 . >)
    (#x13 . <)
    (#x14 . dup)
    (#x15 . swap)
    (#x16 . jmp)))

(define op-lookup (cute assq-ref *opcodes* <>))

(define (binop-lookup op)
  (case (op-lookup op)
    [(+) +]
    [(-) -]
    [(and) logand]
    [(or) logior]
    [(nand) (compose lognot logand)]
    [(nor) (compose lognot logior)]
    [(xor) logxor]))

(define (relop-lookup op)
  (case (op-lookup op)
    [(>) >]
    [(<) <]
    [(=) =]))

(define (jump addr)
  (seek (current-input-port) addr SEEK_SET))

(define (tell)
  (ftell (current-input-port)))

(define* (make-vm #:key stack-size memory-size)
  "Create a fresh VM, with optional stack and memory sizes"
  (define data-stack (if stack-size (make-stack stack-size) (make-stack)))
  (define ret-stack (if stack-size (make-stack stack-size) (make-stack)))
  (define ram (if memory-size (make-ram memory-size) (make-ram)))
  (define (fetch-and-execute)
    (define op (read-u8))
    (unless (eof-object? op)
      (case (op-lookup op)
        [(push)
         (push data-stack (read-word))]
        [(pop)
         (pop data-stack)]
        [(store)
         (let ([addr (read-word)]
               [v (pop data-stack)])
           (ram-set! ram addr v))]
        [(load)
         (let* ([addr (read-word)]
                [v (ram-ref ram addr)])
           (push data-stack v))]
        [(+ - and or nand nor xor)
         (let ([v2 (pop data-stack)]
               [v1 (pop data-stack)])
           (push data-stack ((binop-lookup op) v1 v2)))]
        [(= > <)
         (let ([v2 (pop data-stack)]
               [v1 (pop data-stack)])
           (if ((relop-lookup op) v1 v2)
               (push data-stack 1)
               (push data-stack 0)))]
        [(if)
         (let ([addr (read-word)])
           (when (zero? (peek data-stack))
             (jump addr)))]
        [(call)
         (let ([addr (read-word)])
           (push ret-stack (tell))
           (jump addr))]
        [(return)
         (jump (pop ret-stack))]
        [(>R)
         (push ret-stack (pop data-stack))]
        [(R>)
         (push data-stack (pop ret-stack))]
        [(dup)
         (push data-stack (peek data-stack))]
        [(swap)
         (swap data-stack)]
        [(jmp)
         (jump (read-word))])
      (fetch-and-execute)))
  (lambda (x)
    (case x
      [(run) fetch-and-execute]
      [(vm-memory-ref) (cute ram-ref ram <>)]
      [(vm-memory-set!) (lambda (k v) (ram-set! ram k v))])))

(define (vm-memory-ref vm k)
  "Externally access VM memory at k"
  ((vm 'vm-memory-ref) k))

(define (vm-memory-set! vm k v)
  "Externally set VM memory at k to v"
  ((vm 'vm-memory-set!) k v))

(define (run-vm vm port)
  "Read and execute instructions read from port on VM"
  (with-input-from-port port (vm 'run)))