(define-module (scmvm vm)
  #:use-module ((scheme base)
                #:select (read-u8 read-bytevector))
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-26)
  #:export ((make-vm* . make-vm) run-vm
	    vm-memory-ref vm-memory-byte-ref vm-memory-set! vm-memory vm-load-program!
	    vm-data-stack vm-ret-stack
	    vm-debugger vm-debugger-set!
            vm-pc vm-pc-set!
            *instruction-set* instruction-type instruction-code))

;;; Data Structures
(define *stack-size* 512)
(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-bytevector memory-size #x00))


;;; 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))))

;;; Instructions
(define *instruction-set*
  '((push   #x01)
    (!      #x02)
    (@      #x03)
    (+      #x04)
    (-      #x05)
    (and    #x06)
    (or     #x07)
    (nand   #x08)
    (nor    #x09)
    (xor    #x0a)
    (=      #x0b)
    (>      #x0c)
    (<      #x0d)
    (jmp    #x0e)
    (branch #x0f)
    (call   #x10)
    (return #x11)
    (>R     #x12)
    (R>     #x13)
    (drop   #x14)
    (nip    #x15)
    (dup    #x16)
    (swap   #x17)
    (rot    #x18)
    (over   #x19)
    (bye    #xff)))

(define instruction-code cadr)

(define (op-lookup code)
  (car (find (lambda (x) (= (instruction-code x) code)) *instruction-set*)))

(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-record-type <vm>
  (make-vm data-stack ret-stack memory pc debugger)
  vm?
  (data-stack vm-data-stack)
  (ret-stack vm-ret-stack)
  (memory vm-memory)
  (pc vm-pc vm-pc-set!)
  (debugger vm-debugger vm-debugger-set!))

(define* (make-vm* #:key stack-size memory-size debugger)
  "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)))
  (make-vm data-stack ret-stack ram 1 debugger))

;;; Execution
(define (run-vm vm)
  "Begin execution at pc"
  (define data-stack (vm-data-stack vm))
  (define ret-stack (vm-ret-stack vm))
  (define ram-word-ref (cute vm-memory-ref vm <>))
  (define ram-byte-ref (cute vm-memory-byte-ref vm <>))
  (define ram-word-set! (cute vm-memory-set! vm <> <>))
  (define debugger (vm-debugger vm))
  (define exit? #f)
  (define (jump x) (vm-pc-set! vm x))
  (define (fetch-byte)
    (let ([byte (ram-byte-ref (vm-pc vm))])
      (vm-pc-set! vm (+ (vm-pc vm) 1))
      byte))
  (define (fetch-word)
    (let ([word (ram-word-ref (vm-pc vm))])
      (vm-pc-set! vm (+ (vm-pc vm) 4))
      word))
  (when debugger
    (debugger))
  (define op (fetch-byte))
  (case (op-lookup op)
    [(push)
     (push data-stack (fetch-word))]
    [(!)
     (let ([addr (pop data-stack)]
           [v (pop data-stack)])
       (ram-word-set! addr v))]
    [(@)
     (let* ([addr (pop data-stack)]
            [v (ram-word-ref 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)))]
    [(jmp)
     (jump (pop data-stack))]
    [(branch)
     (let ([addr (pop data-stack)])
       (when (zero? (pop data-stack))
         (jump addr)))]
    [(call)
     (let ([addr (pop data-stack)])
       (push ret-stack (vm-pc vm))
       (jump addr))]
    [(return)
     (jump (pop ret-stack))]
    [(>R)
     (push ret-stack (pop data-stack))]
    [(R>)
     (push data-stack (pop ret-stack))]
    [(drop)
     (pop data-stack)]
    [(nip)
     (let ([v (pop data-stack)])
       (pop data-stack)
       (push data-stack v))]
    [(dup)
     (push data-stack (peek data-stack))]
    [(swap)
     (swap data-stack)]
    [(rot)
     (let* ([a (pop data-stack)]
            [b (pop data-stack)]
            [c (pop data-stack)])
       (push data-stack a)
       (push data-stack c)
       (push data-stack b))]
    [(over)
     (let* ([a (pop data-stack)]
            [b (pop data-stack)])
       (push data-stack b)
       (push data-stack a)
       (push data-stack b))]
    [(bye) (set! exit? #t)])
  (when (not exit?)
    (run-vm vm)))

(define (vm-memory-ref vm k)
  (if (< k 1)
      (error "null memory read")
      (bytevector-s32-native-ref (vm-memory vm) (1- k))))
(define (vm-memory-byte-ref vm k)
  (if (< k 1)
      (error "null memory read")
      (bytevector-u8-ref (vm-memory vm) (1- k))))
(define (vm-memory-set! vm k v)
  (if (< k 1)
      (error "null memory write")
      (bytevector-s32-native-set! (vm-memory vm) (1- k) v)))

(define (vm-load-program! vm prgm)
  "Loads the bytevector into the vm, starting at memory address 1"
  (let ([ram (vm-memory vm)])
    (bytevector-copy! prgm 0
                      ram 0
                      (bytevector-length prgm))))