; From (stack.srfi-9 . module-file)
(define (refer-local-op var next)
  (vector 'refer-local var next))


(define (refer-local-op? v)
  (and (vector? v)
       (eq? 'refer-local (vector-ref v 0))))


(define (refer-local-var v) (vector-ref v 1))


(define (refer-local-next v) (vector-ref v 2))


(define (refer-free-op var next)
  (vector 'refer-free var next))


(define (refer-free-op? v)
  (and (vector? v)
       (eq? 'refer-free (vector-ref v 0))))


(define (refer-free-var v) (vector-ref v 1))


(define (refer-free-next v) (vector-ref v 2))


(define (indirect-op next)
  (vector 'indirect next))


(define (indirect-op? v)
  (and (vector? v)
       (eq? 'indirect (vector-ref v 0))))


(define (indirect-next v) (vector-ref v 1))


(define (constant-op obj next)
  (vector 'constant obj next))


(define (constant-op? v)
  (and (vector? v)
       (eq? 'constant (vector-ref v 0))))


(define (constant-obj v) (vector-ref v 1))


(define (constant-next v) (vector-ref v 2))


(define (close-op n body next)
  (vector 'close n body next))


(define (close-op? v)
  (and (vector? v) (eq? 'close (vector-ref v 0))))


(define (close-body v) (vector-ref v 2))


(define (close-n v) (vector-ref v 1))


(define (close-next v) (vector-ref v 3))


(define (box-op n next) (vector 'box n next))


(define (box-op? v)
  (and (vector? v) (eq? 'box (vector-ref v 0))))


(define (box-n v) (vector-ref v 1))


(define (box-next v) (vector-ref v 2))


(define (test-op if-true if-false)
  (vector 'test if-true if-false))


(define (test-op? v)
  (and (vector? v) (eq? 'test (vector-ref v 0))))


(define (test-then v) (vector-ref v 1))


(define (test-else v) (vector-ref v 2))


(define (assign-local-op var next)
  (vector 'assign-local var next))


(define (assign-local-op? v)
  (and (vector? v)
       (eq? 'assign-local (vector-ref v 0))))


(define (assign-local-var v) (vector-ref v 1))


(define (assign-local-next v) (vector-ref v 2))


(define (assign-free-op var next)
  (vector 'assign-free var next))


(define (assign-free-op? v)
  (and (vector? v)
       (eq? 'assign-free (vector-ref v 0))))


(define (assign-free-var v) (vector-ref v 1))


(define (assign-free-next v) (vector-ref v 2))


(define (conti-op conti next)
  (vector 'conti conti next))


(define (conti-op? v)
  (and (vector? v) (eq? 'conti (vector-ref v 0))))


(define (conti-next v) (vector-ref v 2))


(define (nuate-op stack next)
  (vector 'nuate stack next))


(define (nuate-op? v)
  (and (vector? v) (eq? 'nuate (vector-ref v 0))))


(define (nuate-stack v) (vector-ref v 1))


(define (nuate-next v) (vector-ref v 2))


(define (frame-op ret next)
  (vector 'frame ret next))


(define (frame-op? v)
  (and (vector? v) (eq? 'frame (vector-ref v 0))))


(define (frame-ret v) (vector-ref v 1))


(define (frame-next v) (vector-ref v 2))


(define (argument-op next)
  (vector 'argument next))


(define (argument-op? v)
  (and (vector? v)
       (eq? 'argument (vector-ref v 0))))


(define (argument-next v) (vector-ref v 1))


(define (shift-op n m next)
  (vector 'shift n m next))


(define (shift-op? v)
  (and (vector? v) (eq? 'shift (vector-ref v 0))))


(define (shift-n v) (vector-ref v 1))


(define (shift-m v) (vector-ref v 2))


(define (shift-next v) (vector-ref v 3))


(define (apply-op n) (vector 'apply n))


(define (apply-op? v)
  (and (vector? v) (eq? 'apply (vector-ref v 0))))


(define (apply-n v) (vector-ref v 1))


(define (return-op n) (vector 'return n))


(define (return-op? v)
  (and (vector? v) (eq? 'return (vector-ref v 0))))


(define (return-n v) (vector-ref v 1))


(define (halt-op) (vector 'halt))


(define (halt-op? v)
  (and (vector? v) (eq? 'halt (vector-ref v 0))))


(define (make-set . args) args)


(define (list->set list) list)


(define (set-member? x s)
  (cond ((null? s) #f)
        ((eq? x (car s)) #t)
        (else (set-member? x (cdr s)))))


(define (set-add-element x s)
  (if (set-member? x s) s (cons x s)))


(define (set-union s1 s2)
  (if (null? s1)
    s2
    (set-union
      (cdr s1)
      (set-add-element (car s1) s2))))


(define (set-minus s1 s2)
  (if (null? s1)
    (make-set)
    (let ((e (car s1)))
      (if (set-member? e s2)
        (set-minus (cdr s1) s2)
        (cons (car s1) (set-minus (cdr s1) s2))))))


(define (set-intersect s1 s2)
  (if (null? s1)
    (make-set)
    (let ((e (car s1)) (rest (cdr s1)))
      (if (set-member? e s2)
        (cons e (set-intersect rest s2))
        (set-intersect rest s2)))))


(define (compile-lookup x e return-local return-free)
  (let nxtlocal ((locals (car e)) (n 0))
    (if (null? locals)
      (let nxtfree ((free (cdr e)) (n 0))
        (if (eq? (car free) x)
          (return-free n)
          (nxtfree (cdr free) (+ n 1))))
      (if (eq? (car locals) x)
        (return-local n)
        (nxtlocal (cdr locals) (+ n 1))))))


(define (compile-refer x e next)
  (compile-lookup
    x
    e
    (lambda (n) (refer-local-op n next))
    (lambda (n) (refer-free-op n next))))


(define (collect-free vars e next)
  (if (null? vars)
    next
    (collect-free
      (cdr vars)
      e
      (compile-refer (car vars) e (argument-op next)))))


(define (tail? next) (return-op? next))


(define (make-boxes sets vars next)
  (let f ((vars vars) (n 0))
    (if (null? vars)
      next
      (if (set-member? (car vars) sets)
        (box-op n (f (cdr vars) (+ n 1)))
        (f (cdr vars) (+ n 1))))))


(define (find-sets x v)
  (cond ((symbol? x) (make-set))
        ((pair? x)
         (let ((vx (list->vector x)))
           (case (vector-ref vx 0)
             ((quote) (make-set))
             ((lambda)
              (find-sets
                (vector-ref vx 2)
                (set-minus v (list->set (vector-ref vx 1)))))
             ((if)
              (set-union
                (find-sets (vector-ref vx 1) v)
                (set-union
                  (find-sets (vector-ref vx 2) v)
                  (find-sets (vector-ref vx 3) v))))
             ((set!)
              (let ((var (vector-ref vx 1)))
                (set-union
                  (if (set-member? var v)
                    (make-set var)
                    (make-set))
                  (find-sets (vector-ref vx 2) v))))
             ((call/cc) (find-sets (vector-ref vx 1) v))
             (else
              (let next ((x x))
                (if (null? x)
                  (make-set)
                  (set-union (find-sets (car x) v) (next (cdr x)))))))))
        (else (make-set))))


(define (find-free x b)
  (cond ((symbol? x)
         (if (set-member? x b) (make-set) (make-set x)))
        ((pair? x)
         (let ((v (list->vector x)))
           (case (car x)
             ((quote) (make-set))
             ((lambda)
              (find-free
                (vector-ref v 2)
                (set-union (vector-ref v 1) b)))
             ((if)
              (set-union
                (find-free (vector-ref v 1) b)
                (set-union
                  (find-free (vector-ref v 2) b)
                  (find-free (vector-ref v 3) b))))
             ((set!)
              (let ((var (vector-ref v 1)))
                (set-union
                  (if (set-member? var b)
                    (make-set)
                    (make-set var))
                  (find-free (vector-ref v 2) b))))
             ((call/cc) (find-free (vector-ref v 1) b))
             (else
              (let next ((x x))
                (if (null? x)
                  (make-set)
                  (set-union (find-free (car x) b) (next (cdr x)))))))))
        (else (make-set))))


(define (compile x e s next k-fail)
  (cond ((symbol? x)
         (compile-refer
           x
           e
           (if (set-member? x s) (indirect-op next) next)))
        ((list? x)
         (let* ((v (list->vector x))
                (x-length (vector-length v)))
           (case (vector-ref v 0)
             ((quote)
              (if (= x-length 2)
                (constant-op (vector-ref v 1) next)
                (k-fail x)))
             ((lambda)
              (if (= x-length 3)
                (let* ((vars (vector-ref v 1))
                       (body (vector-ref v 2))
                       (free (find-free body vars))
                       (sets (find-sets body vars)))
                  (collect-free
                    free
                    e
                    (close-op
                      (length free)
                      (make-boxes
                        sets
                        vars
                        (compile
                          body
                          (cons vars free)
                          (set-union sets (set-intersect s free))
                          (return-op (length vars))
                          k-fail))
                      next)))
                (k-fail x)))
             ((if)
              (if (= x-length 4)
                (let ((thenc (compile (vector-ref v 2) e s next k-fail))
                      (elsec (compile (vector-ref v 3) e s next k-fail)))
                  (compile
                    (vector-ref v 1)
                    e
                    s
                    (test-op thenc elsec)
                    k-fail))
                (k-fail x)))
             ((set!)
              (if (= x-length 3)
                (let ((expr (vector-ref v 2)))
                  (compile-lookup
                    (vector-ref v 1)
                    e
                    (lambda (n)
                      (compile
                        expr
                        e
                        s
                        (assign-local-op n next)
                        k-fail))
                    (lambda (n)
                      (compile expr e s (assign-free-op n next) k-fail))))
                (k-fail x)))
             ((call/cc)
              (if (= x-length 2)
                (let ((c (conti-op
                           (argument-op
                             (compile
                               (vector-ref v 1)
                               e
                               s
                               (if (tail? next)
                                 (shift-op 1 (return-n next) (apply-op 0))
                                 (apply-op 0))
                               k-fail)))))
                  (if (tail? next) c (frame-op next c)))
                (k-fail x)))
             (else
              (let loop ((args (cdr x))
                         (c (compile
                              (car x)
                              e
                              s
                              (let ((n-args (- (vector-length v) 1)))
                                (if (tail? next)
                                  (shift-op
                                    n-args
                                    (return-n next)
                                    (apply-op n-args))
                                  (apply-op n-args)))
                              k-fail)))
                (if (null? args)
                  (if (tail? next) c (frame-op next c))
                  (loop (cdr args)
                        (compile (car args) e s (argument-op c) k-fail))))))))
        (else (constant-op x next))))


(define stack (make-vector 10000))


(define (push x s)
  (vector-set! stack s x)
  (+ s 1))


(define (index s i)
  (vector-ref stack (- (- s i) 1)))


(define (index-set! s i v)
  (vector-set! stack (- (- s i) 1) v))


(define (shift-args n m s)
  (let nxtarg ((i (- n 1)))
    (if (< i 0)
      (- s m)
      (begin
        (index-set! s (+ i m) (index s i))
        (nxtarg (- i 1))))))


(define (closure body n s)
  (let ((v (make-vector (+ n 1))))
    (vector-set! v 0 body)
    (let f ((i 0))
      (if (= i n)
        v
        (begin
          (vector-set! v (+ i 1) (index s i))
          (f (+ i 1)))))))


(define (closure-body c) (vector-ref c 0))


(define (index-closure c n)
  (vector-ref c (+ n 1)))


(define (extract-prim-args n s)
  (let loop ((n (- n 1)) (args '()))
    (if (>= n 0)
      (let ((arg (index s n)))
        (loop (- n 1) (cons arg args)))
      args)))


(define (continuation s)
  (closure
    (refer-local-op
      0
      (nuate-op (save-stack s) (return-op 1)))
    0
    s))


(define (save-stack s)
  (let ((v (make-vector s)))
    (let copy ((i 0))
      (if (= i s)
        v
        (begin
          (vector-set! v i (vector-ref stack i))
          (copy (+ i 1)))))))


(define (restore-stack v)
  (let ((s (vector-length v)))
    (let copy ((i 0))
      (if (= i s)
        s
        (begin
          (vector-set! stack i (vector-ref v i))
          (copy (+ i 1)))))))


(define (make-box v) (vector 'box v))


(define (box? v)
  (and (vector? v) (eq? 'box (vector-ref v 0))))


(define (unbox v) (vector-ref v 1))


(define (set-box! v x) (vector-set! v 1 x))


(define (VM-state a x f c s)
  (for-each
    display
    `("VM executing: "
      ,(vector-ref x 0)
      #\newline
      "    accum: "
      ,a
      #\newline
      "    frame: "
      ,f
      #\newline
      "    stack: "
      ,s
      #\newline
      ,@(apply append
               (map (lambda (e) (list "           " e #\newline))
                    (vector->list (save-stack s))))
      #\newline)))


(define (VM k-fail a x f c s)
  (cond ((halt-op? x) a)
        ((refer-local-op? x)
         (VM k-fail
             (index f (refer-local-var x))
             (refer-local-next x)
             f
             c
             s))
        ((refer-free-op? x)
         (VM k-fail
             (index-closure c (refer-free-var x))
             (refer-local-next x)
             f
             c
             s))
        ((indirect-op? x)
         (VM k-fail (unbox a) (indirect-next x) f c s))
        ((constant-op? x)
         (VM k-fail
             (constant-obj x)
             (constant-next x)
             f
             c
             s))
        ((close-op? x)
         (let ((n (close-n x)))
           (VM k-fail
               (closure (close-body x) n s)
               (close-next x)
               f
               c
               (- s n))))
        ((box-op? x)
         (let ((n (box-n x)))
           (index-set! s n (make-box (index s n)))
           (VM k-fail a (box-next x) f c s)))
        ((test-op? x)
         (VM k-fail
             a
             (if a (test-then x) (test-else x))
             f
             c
             s))
        ((assign-local-op? x)
         (set-box! (index f (assign-local-var x)) a)
         (VM k-fail a (assign-local-next x) f c s))
        ((assign-free-op? x)
         (set-box!
           (index-closure f (assign-free-var x))
           a)
         (VM k-fail a (assign-free-next x) f c s))
        ((conti-op? x)
         (VM k-fail (continuation s) (conti-next x) f c s))
        ((nuate-op? x)
         (VM k-fail
             a
             (nuate-next x)
             f
             c
             (restore-stack (nuate-stack x))))
        ((frame-op? x)
         (VM k-fail
             a
             (frame-next x)
             f
             c
             (push (frame-ret x) (push f (push c s)))))
        ((argument-op? x)
         (VM k-fail a (argument-next x) f c (push a s)))
        ((shift-op? x)
         (VM k-fail
             a
             (shift-next x)
             f
             c
             (shift-args (shift-n x) (shift-m x) s)))
        ((apply-op? x)
         (if (procedure? a)
           (let ((n-args (apply-n x)))
             (VM k-fail
                 (apply a (extract-prim-args n-args s))
                 (return-op n-args)
                 f
                 c
                 s))
           (VM k-fail a (closure-body a) s a s)))
        ((return-op? x)
         (let ((s (- s (return-n x))))
           (VM k-fail
               a
               (index s 0)
               (index s 1)
               (index s 2)
               (- s 3))))
        (else (k-fail a x f c s))))


(define r5rs-names
  '(()
    cons
    pair?
    car
    cdr
    number?
    +
    -
    *
    /
    <
    <=
    =
    >=
    >))


(define r5rs-values
  (vector
    cons
    pair?
    car
    cdr
    number?
    +
    -
    *
    /
    <
    <=
    =
    >=
    >))


(define (compile-error x)
  (for-each
    display
    `("Compiler error at: " ,x #\newline)))


(define (VM-error a x f c s)
  (for-each
    display
    `("VM Error! registers: "
      #\newline
      "a: "
      ,a
      #\newline
      "x: "
      ,x
      #\newline
      "f: "
      ,f
      #\newline
      "c: "
      ,c
      #\newline
      "s: "
      ,s
      #\newline)))


(define (stack-eval x)
  (let* ((code (compile
                 x
                 r5rs-names
                 '()
                 (halt-op)
                 compile-error))
         (closure
           (let init ((i (- (vector-length r5rs-values) 1)) (s 0))
             (if (<= 0 i)
               (init (- i 1)
                     (push (vector-ref r5rs-values i) s))
               (closure code (vector-length r5rs-values) s)))))
    (VM VM-error closure code 0 closure 0)))


(define (read-prompted)
  (display "dybvig-stack> ")
  (read))


(define (stack-repl)
  (let repl ((expr (read-prompted)))
    (if (not (eof-object? expr))
      (let ((val (stack-eval expr)))
        (write val)
        (newline)
        (repl (read-prompted))))))


(stack-repl)