(cs441-load "sets.scm")
(cs441-load "parser.scm")

;;;;;;;;;;;;;;;;
;; Heap allocator

;; pointers
(define-record pointer (contents))
(define ptr make-pointer)
(define follow pointer->contents)

(define semi-space-size 1000)
(define heap (make-vector (* 2 semi-space-size) #f))
(define allocp 0)
(define from-space 0)
(define space-remaining
  (lambda ()
    (- semi-space-size (- allocp from-space))))

(define heap-alloc
  (lambda args
    (let ((n (+ (length args) 1)))
      (if (< (space-remaining) n)
          (error "heap exhausted"))
      (let ((p allocp))
        (set! allocp (+ allocp n))
        (range-copy! (list->vector (cons n args)) 0 heap p n)
        (ptr p)))))
      
(define heap-ref
  (lambda (p offset)
    (vector-ref heap (+ (follow p) offset))))

(define heap-set!
  (lambda (p offset v)
    (vector-set! heap (+ (follow p) offset) v)))

(define range-copy!
  (lambda (from from-index to to-index length)
    (if (not (zero? length))
        (begin
          (vector-set! to to-index (vector-ref from from-index))
          (range-copy!
            from (+ from-index 1)
            to (+ to-index 1)
            (- length 1))))))

;;;;;;;;;;;;;;;;
;; record types for heap objects

(load "heap-record.scm")

(define-heap-record Empty-ff ())
(define-heap-record Extend-ff (f d r))

(define-heap-record Prim (name))
(define-heap-record Closure (body formals env))
(define-heap-record Kont (k))

(define-heap-record Kinitial ())
(define-heap-record Kfun (args env k))
(define-heap-record Karg (vfun args vargs env k))
(define-heap-record Kif (then else env k))
(define-heap-record Kset! (env name k))
(define-heap-record Kbegin (exprs env k))
(define-heap-record Kletrec (bindings name env body k))

(define-heap-record Pear (a b))

(define-heap-record box (contents))
(define unbox box->contents)
(define set-box!
  (lambda (b v)
    (heap-set! b 2 v)))

(define Pears->pairs
  (lambda (p)
    (if (Pear? p)
        (cons (Pear->a p) (Pears->pairs (Pear->b p)))
        p)))

;;;;;;;;;;;;;;;;

(define make-empty
  (lambda ()
    (make-Empty-ff)))

(define extend make-Extend-ff)

(define lookup
  (lambda (f x)
    (variant-case f
      (Extend-ff (f d r)
        (if (eq? d x)
            r
            (lookup f x))))))

(define extend-list
  (lambda (env doms rans)
    (if (null? doms)
        env
        (extend-list
          (extend env (car doms) (car rans))
          (cdr doms) 
          (cdr rans)))))

(define firsts
  (lambda (lst)
    (map car lst)))

(define seconds
  (lambda (lst)
    (map cadr lst)))

(define last 
  (lambda (lst)
    (cond ((null? lst) '())
          ((null? (cdr lst)) (car lst))
          (else (last (cdr lst))))))

(define primitives
  (list
    (list 'not     not)
    (list 'symbol? symbol?)
    (list 'number? number?)
    (list 'null?   null?)
    (list 'pair?   Pear?)
    (list 'car     Pear->a)
    (list 'cdr     Pear->b)
    (list 'zero?   zero?)
    (list 'eq?     eq?)
    (list 'equal?  equal?)
    (list '=       =)
    (list 'cons    make-Pear)
    (list '+       +)
    (list '*       *)
    (list '-       -)
    (list '/       /)))

(define make-initial-env
  (lambda ()
    (extend-list (make-empty)
      (firsts primitives)
      (map (lambda (x) (make-box (make-Prim x))) (firsts primitives)))))

(define prim-func
  (lambda (name)
    (letrec ((find (lambda (prims)
                     (cond ((null? prims)
                            (error "unknown primitive"))
                           ((eq? name (car (car prims)))
                            (cadr (car prims)))
                           (else
                             (find (cdr prims)))))))
      (find primitives))))

(define make-initial-cont make-Kinitial)

;; machine registers

(define *e* #f)
(define *env* #f)
(define *k* #f)
(define *v* #f)

(define interp
  (lambda ()
    (if (< (space-remaining) 20) (collect))
    (variant-case *e*
      (Const (value)
        (set! *v* value)
        (apply-k))
      (Var (name)
        (set! *v* (unbox (lookup *env* name)))
        (apply-k))
      (Lam (formals body) 
        (set! *v* (make-Closure body formals *env*))
	(apply-k))
      (Ap (fun args)
        (set! *k* (make-Kfun args *env* *k*))
        (set! *e* fun)
        (interp))
      (Beg (exprs)
        (set! *k* (make-Kbegin (cdr exprs) *env* *k*))
        (set! *e* (car exprs))
        (interp))
      (Set! (name body)
        (set! *k* (make-Kset! *env* name *k*))
        (set! *e* body)
        (interp))
      (If (test then else)
        (set! *k* (make-Kif then else *env* *k*))
        (set! *e* test)
        (interp))
      (Let (bindings body)
        (set! *e* (make-Ap
                    (make-Lam
                      (firsts bindings)
                      body)
                    (seconds bindings)))
        (interp))
      (Let* (bindings body)
        (if (null? bindings)
            (set! *e* body)
            (let ((x (car (car bindings)))
                  (e (cadr (car bindings)))
                  (rest (cdr bindings)))
              (set! *e* (make-Ap
                          (make-Lam (list x)
                            (make-Let* rest body))
                          (list e)))))
        (interp))
      (Letrec (bindings body)
        (if (null? bindings)
            (begin
              (set! *e* body)
              (interp))
            (let* ((boxes (map make-box (firsts bindings)))
                   (newenv (extend-list *env* (firsts bindings) boxes))
                   (x (car (car bindings)))
                   (e (cadr (car bindings))))
              (set! *k* (make-Kletrec (cdr bindings) x newenv body *k*))
              (set! *e* e)
              (set! *env* newenv)
              (interp))))
      (Abort (body)
        (set! *k* (make-initial-cont))
        (set! *e* body)
        (interp))
      (Let/cc (name body)
        (set! *env* (extend *env* name (make-box (make-Kont *k*))))
        (set! *e* body)
        (interp)))))

(define apply-k
  (lambda ()
    (if (< (space-remaining) 20) (collect))
    (variant-case *k*
      (Kinitial ()
        *v*)
      (Kfun (args env k)
        (if (null? args)
            (apply-proc *v* '() k)
            (begin
              (set! *e* (car args))
              (set! *env* env)
              (set! *k* (make-Karg *v* (cdr args) '() env k))
              (interp))))
      (Karg (vfun args vargs env k)
        (if (null? args)
            (apply-proc vfun (reverse (Pears->pairs (make-Pear *v* vargs))) k)
            (begin
              (set! *e* (car args))
              (set! *env* env)
              (set! *k* (make-Karg vfun (cdr args) (make-Pear *v* vargs) env k))
              (interp))))
      (Kif (then else env k)
        (set! *e* (if *v* then else))
        (set! *env* env)
        (set! *k* k)
        (interp))
      (Kset! (env name k)
        (set! *k* k)
        (set! *v* (set-box! (lookup env name) *v*))
	(apply-k))
      (Kbegin (exprs env k)
        (if (null? exprs)
            (apply-k)
            (begin
              (set! *e* (car exprs))
              (set! *env* env)
              (set! *k* (make-Kbegin (cdr exprs) env k))
              (interp))))
      (Kletrec (bindings name env body k)
        (set-box! (lookup env name) *v*)
        (set! *env* env)
        (if (null? bindings)
            (begin
              (set! *e* body)
              (set! *k* k)
              (interp))
            (let* ((x (car (car bindings)))
                   (e (cadr (car bindings))))
              (set! *e* e)
              (set! *k* (make-Kletrec (cdr bindings) x env body k))
              (interp)))))))

(define apply-proc
  (lambda (vfun vargs k)
    (variant-case vfun
      (Closure (body formals env)
        (set! *e* body)
        (set! *env* (extend-list env formals (map make-box vargs)))
        (set! *k* k)
        (interp))
      (Prim (name)
        (set! *k* k)
        (set! *v* (apply (prim-func name) vargs))
        (apply-k))
      (Kont (k)
        (set! *k* k)
        (set! *v* (car vargs))
        (apply-k)))))

(define eval
  (lambda (sexp)
    (set! *e* (parse sexp))
    (set! *env* (make-initial-env))
    (set! *k* (make-initial-cont))
    (set! *v* #f)
    (interp)))

(define collect
  (lambda ()
    (let* ((to-space (if (zero? from-space)
                         semi-space-size
                         0))
           (to-allocp to-space)
           (copy-object
             (lambda (p)
               (cond ((not (pointer? p))
                      p)
                     ((pointer? (heap-ref p 0))
                      (heap-ref p 0))
                     (else
                       (let ((size (heap-ref p 0))
                             (new (ptr to-allocp)))
                         (range-copy!
                           heap (follow p)
                           heap to-allocp
                           size)
                         (set! to-allocp (+ to-allocp size))
                         (heap-set! p 0 new)
                         new))))))
      ;; Copy the roots
      (set! *env* (copy-object *env*))
      (set! *k* (copy-object *k*))
      (set! *v* (copy-object *v*))
      ;; Copy everything reachable from the roots
      (letrec ((loop (lambda (copyp)
                       (if (< copyp to-allocp)
                           (let ((size (vector-ref heap copyp)))
                             (for 2 size
                               (lambda (i)
                                 (let ((p (vector-ref heap (+ copyp i))))
                                   (if (pointer? p)
                                       (vector-set! heap (+ copyp i)
                                         (copy-object p))
                                       #f))))
                             (loop (+ copyp size)))
                           #f))))
        (loop to-space))
      (printf "Collected ~a cells~%"
        (- (- allocp from-space) (- to-allocp to-space)))
      ;; Flip semi-spaces
      (set! from-space to-space)
      (set! allocp to-allocp))))

(define for
  (lambda (low high proc)
    (if (< low high)
        (begin
          (proc low)
          (for (+ low 1) high proc)))))