(define-module (scmvm vm forth)
  #:use-module (scmvm vm)
  #:use-module (scmvm util stack)
  #:use-module ((scheme base)
                #:select (read-u8 read-bytevector))
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-9)
  #:export ((make-forth-vm* . make-forth-vm)
            forth-instruction-set
            forth-vm-load-program!
            forth-vm-run!
            forth-vm-pc
            forth-vm-pc-set!
            forth-vm-memory-set!
            forth-vm-memory-ref
            forth-vm-data-stack))

;;; IO
(define *memory-size* 2048)

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

(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))))


;;; VM
(define-record-type <forth-vm>
  (make-forth-vm instruction-set registers)
  forth-vm?
  (instruction-set forth-vm-instruction-set)
  (registers forth-vm-registers forth-vm-registers-set!))

(define* (make-forth-vm* instruction-set #:key stack-size memory-size)
  "Create a fresh VM, with optional stack and memory sizes"
  (define registers (register-set '(*data-stack* *ret-stack* *pc* *ram* *vm-exit*)))
  ((hash-ref registers '*pc*) 1)
  ((hash-ref registers '*data-stack*) (if stack-size (make-stack stack-size) (make-stack)))
  ((hash-ref registers '*ret-stack*) (if stack-size (make-stack stack-size) (make-stack)))
  ((hash-ref registers '*ram*) (if memory-size (make-ram memory-size) (make-ram)))
  (make-forth-vm instruction-set registers))

(define *vm* (make-parameter #f))

(define-syntax-rule (with-forth-vm-registers (reg ...) expr ...)
  (let ([reg (hash-ref (forth-vm-registers (*vm*)) 'reg)] ...)
    expr ...))

(define (ram-word-ref k)
  (with-forth-vm-registers
   (*ram*)
   (if (< k 1)
       (error "null memory read")
       (bytevector-s32-native-ref (*ram*) (1- k)))))

(define (ram-byte-ref k)
  (with-forth-vm-registers
   (*ram*)
   (if (< k 1)
       (error "null memory read")
       (bytevector-u8-ref (*ram*) (1- k)))))

(define (ram-word-set! k v)
  (with-forth-vm-registers
   (*ram*)
   (if (< k 1)
       (error "null memory write")
       (bytevector-s32-native-set! (*ram*) (1- k) v))))

(define (jump! x)
  (with-forth-vm-registers
   (*pc*)
   (*pc* (logand #x2fffffff x))))

(define (fetch-byte!)
  (with-forth-vm-registers
   (*pc*)
   (let* ([byte (ram-byte-ref (*pc*))])
     (*pc* (+ (*pc*) 1))
     byte)))

(define (fetch-word!)
  (with-forth-vm-registers
   (*ram* *pc*)
   (let* ([word (ram-word-ref (*pc*))])
     (*pc* (+ (*pc*) 4))
     word)))

(define (forth-vm-pc vm)
  (parameterize ([*vm* vm])
    (with-forth-vm-registers
     (*pc*)
     (*pc*))))

(define (forth-vm-pc-set! vm k)
  (parameterize ([*vm* vm])
    (with-forth-vm-registers
     (*pc*)
     (*pc* k))))

(define (forth-vm-memory-set! vm k v)
  (parameterize ([*vm* vm])
    (ram-word-set! k v)))

(define (forth-vm-memory-ref vm k)
  (parameterize ([*vm* vm])
    (with-forth-vm-registers
     (*ram*)
     (ram-word-ref k))))

(define (forth-vm-data-stack vm)
  (parameterize ([*vm* vm])
    (with-forth-vm-registers
     (*data-stack*)
     (*data-stack*))))

 ;;; Instruction set
(define-instruction-set (forth-instruction-set *pc* *ram* *data-stack* *ret-stack* *vm-exit*)
  (define-instruction (push #x01)
    (stack-push (*data-stack*) (fetch-word!)))
  
  (define-instruction (! #x02)
    (let ([addr (stack-pop (*data-stack*))]
          [v (stack-pop (*data-stack*))])
      (ram-word-set! addr v)))
  
  (define-instruction (@ #x03)
    (let* ([addr (stack-pop (*data-stack*))]
           [v (ram-word-ref addr)])
      (stack-push (*data-stack*) v)))
  
  (define-instruction (+ #x04)
    (let ([v2 (stack-pop (*data-stack*))]
          [v1 (stack-pop (*data-stack*))])
      (stack-push (*data-stack*) (+ v1 v2))))
  
  (define-instruction (- #x05)
    (let ([v2 (stack-pop (*data-stack*))]
          [v1 (stack-pop (*data-stack*))])
      (stack-push (*data-stack*) (- v1 v2))))
  
  (define-instruction (and #x06)
    (let ([v2 (stack-pop (*data-stack*))]
          [v1 (stack-pop (*data-stack*))])
      (stack-push (*data-stack*) (logand v1 v2))))

  (define-instruction (or #x07)
    (let ([v2 (stack-pop (*data-stack*))]
          [v1 (stack-pop (*data-stack*))])
      (stack-push (*data-stack*) (logior v1 v2))))

  (define-instruction (nand #x08)
    (let ([v2 (stack-pop (*data-stack*))]
          [v1 (stack-pop (*data-stack*))])
      (stack-push (*data-stack*) (if (zero? (logand v1 v2)) 1 0))))

  (define-instruction (nor #x09)
    (let ([v2 (stack-pop (*data-stack*))]
          [v1 (stack-pop (*data-stack*))])
      (stack-push (*data-stack*) (if (zero? (logior v1 v2)) 1 0))))

  (define-instruction (xor #x0a)
    (let ([v2 (stack-pop (*data-stack*))]
          [v1 (stack-pop (*data-stack*))])
      (stack-push (*data-stack*) (logxor v1 v2))))

  (define-instruction (= #x0b)
    (let ([v2 (stack-pop (*data-stack*))]
          [v1 (stack-pop (*data-stack*))])
      (if (= v1 v2)
          (stack-push (*data-stack*) 1)
          (stack-push (*data-stack*) 0))))

  (define-instruction (> #x0c)
    (let ([v2 (stack-pop (*data-stack*))]
          [v1 (stack-pop (*data-stack*))])
      (if (> v1 v2)
          (stack-push (*data-stack*) 1)
          (stack-push (*data-stack*) 0))))

  (define-instruction (< #x0d)
    (let ([v2 (stack-pop (*data-stack*))]
          [v1 (stack-pop (*data-stack*))])
      (if (< v1 v2)
          (stack-push (*data-stack*) 1)
          (stack-push (*data-stack*) 0))))

  (define-instruction (jmp #x0e)
    (jump! (stack-pop (*data-stack*))))

  (define-instruction (branch #x0f)
    (let* ([addr (stack-pop (*data-stack*))]
           [test (stack-pop (*data-stack*))])
      (when (zero? test)
        (jump! addr))))

  (define-instruction (call #x10)
    (let ([addr (stack-pop (*data-stack*))])
      (stack-push (*ret-stack*) (*pc*))
      (jump! addr)))

  (define-instruction (return #x11)
    (jump! (stack-pop (*ret-stack*))))

  (define-instruction (>R #x12)
    (stack-push (*ret-stack*) (stack-pop (*data-stack*))))

  (define-instruction (R> #x13)
    (stack-push (*data-stack*) (stack-pop (*ret-stack*))))

  (define-instruction (drop #x14)
    (stack-pop (*data-stack*)))

  (define-instruction (nip #x15)
    (let ([v (stack-pop (*data-stack*))])
      (stack-pop (*data-stack*))
      (stack-push (*data-stack*) v)))

  (define-instruction (dup #x16)
    (stack-push (*data-stack*) (stack-peek (*data-stack*))))

  (define-instruction (swap #x17)
    (stack-swap (*data-stack*)))

  (define-instruction (rot #x18)
    (let* ([a (stack-pop (*data-stack*))]
           [b (stack-pop (*data-stack*))]
           [c (stack-pop (*data-stack*))])
      (stack-push (*data-stack*) a)
      (stack-push (*data-stack*) c)
      (stack-push (*data-stack*) b)))

  (define-instruction (over #x19)
    (let* ([a (stack-pop (*data-stack*))]
           [b (stack-pop (*data-stack*))])
      (stack-push (*data-stack*) b)
      (stack-push (*data-stack*) a)
      (stack-push (*data-stack*) b)))

  (define-instruction (not #x1a)
    (let ([a (stack-pop (*data-stack*))])
      (stack-push (*data-stack*) (if (zero? a) 1 0))))

  (define-instruction (set! #x1b)
    ;; use let* to induce an order of evaluation
    (let* ([idx (stack-pop (*data-stack*))]
           [obj (stack-pop (*data-stack*))])
      (stack-set! (*data-stack*) idx obj)))

  (define-instruction (bye #xff)
    (*vm-exit* #t)))

;;; Execution
(define* (forth-vm-run! vm #:optional debugger)
  "Begin execution at pc"
  (define caller (instruction-set-caller (forth-vm-instruction-set vm)
                                         (forth-vm-registers vm)))
  (parameterize ([*vm* vm])
    (with-forth-vm-registers
     (*vm-exit*)
     (let lp () 
       (when debugger
         (debugger))
       (define op (fetch-byte!))
       (caller op)
       (unless (*vm-exit*) (lp))))))

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