#lang racket
(require [only-in plait test test/exn error print-only-errors])

(module call-dyn plait
  (define-type DynExp
    [numE (val : Number)]
    [plusE (l : DynExp) (r : DynExp)]
    [varE (name : Symbol)]
    [let1E (id : Symbol) (named-expr : DynExp) (bound-body : DynExp)]
    [lamE (param : Symbol) (body : DynExp)]
    [appDynE (fun : DynExp) (val : DynExp)]
    [appE (fun : DynExp) (val : DynExp)])
  (define-type Value
    [numV (n : Number)]
    [lamV (arg : Symbol) (body : DynExp) (env : Env)])
  
  (define-type-alias Env (Hashof Symbol Value))
  (define mt-env (hash empty)) ;; "empty environment"
  
  (define (lookup (s : Symbol) (n : Env))
    (type-case (Optionof Value) (hash-ref n s)
      [(none) (error s "not bound")]
      [(some v) v]))
  
  (define (extend old-env new-name value)
    (hash-set old-env new-name value))
  (define (interp expr env)
    (type-case DynExp expr
      [(numE n) (numV n)]
      [(plusE l r)
       (numV (+ (numV-n (interp l env)) (numV-n (interp r env))))]
      [(let1E bound-id named-expr bound-body)
       (interp bound-body (extend env bound-id (interp named-expr env)))]
      [(varE name) (lookup name env)]
      [(lamE bound-id bound-body) (lamV bound-id bound-body env)]
      [(appDynE fun-expr arg-expr)
       (let ([fval (interp fun-expr env)])
         (type-case Value fval
           [(lamV bound-id bound-body f-env)
            (interp bound-body (extend env bound-id (interp arg-expr f-env)))]
           [else (error 'interp "non-function")]))]
      [(appE fun-expr arg-expr)
       (let ([fval (interp fun-expr env)])
         (type-case Value fval
           [(lamV bound-id bound-body f-env)
            (interp bound-body (extend f-env bound-id (interp arg-expr env)))]
           [else (error 'interp "non-function")]))]))
  )

(require 'call-dyn)
(module+ test
  (define (example body)
    (let1E 'x (numE 3)
           (let1E 'f (lamE 'y (plusE (varE 'x) (varE 'y)))
                  (let1E 'x (numE 5)
                         body))))

  (test (interp (example (appE (varE 'f) (numE 4))) mt-env)
        (numV 7))

  (test (interp
         (example (appDynE (varE 'f) (numE 4))) mt-env)
        (numV 9))

  (test (interp
         (example
          (let1E 'f (lamE 'x (varE 'x))
                 (appDynE (varE 'f) (numE 4))))
         mt-env)
        (numV 4))

  (test (interp
         (example
          (let1E 'f (lamE 'y (varE 'x))
                 (appDynE (varE 'f) (numE 4))))
         mt-env)
        (numV 5))

  (test (interp
         (example
          (let1E 'f (lamE 'y (varE 'x))
                 (let1E 'x (numE 3)
                        (appDynE (varE 'f) (numE 4)))))
         mt-env)
        (numV 3))

  (test/exn (interp (appE (varE 'g) (numE 4)) mt-env) "not bound")
  (test/exn (interp (example (appE (numE 4) (varE 'f))) mt-env) "function")
  (test/exn (interp (example (appDynE (numE 4) (varE 'f))) mt-env) "function")
  )

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; This part provides a small part of the plai/gc2/collector language
(define current-heap (make-parameter (make-vector 0 #f)))
(define (heap-set! index val) (vector-set!  (current-heap) index val))
(define (heap-ref index) (vector-ref (current-heap) index))
(define (heap-size) (vector-length (current-heap)))
(define-syntax-rule (with-heap vec expr ...)
  (parameterize
      ([current-heap vec])
    (begin
      expr ...)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; This part is the partial implementation of a collector API For
;; simplicity, use our very first allocator without any data
;; structure.

(define (init-allocator)
  (vector-fill! (current-heap) 'free))

(define (alloc/header tag . vals)
  (define loc (malloc (+ 1 (length vals))))
  (heap-set! loc tag)
  (for ([i (in-range (length vals))]
        [v (in-list vals)])
    (heap-set! (+ loc i 1) v))
  loc)

(define (gc:alloc-flat val) (alloc/header 'flat val))

(define (gc:cons val1 val2) (alloc/header 'cons val1 val2))


;; Linear time allocator, based on Lecture 20

(define (malloc size)
  (define ptr (find-free-space size))
  (unless ptr  (error 'alloc "out of memory"))
  ptr)

(module+ test
  (with-heap (make-vector 10 #f)
    (init-allocator)
    (test/exn (malloc 100) "out of memory")))

(define (find-free-space n)
  (define (n-free-blocks? start n)
    (for/fold ([ok #t])
              ([i (in-range start (+ start n))])
      (and ok (< i (heap-size)) (equal? (heap-ref i) 'free))))

  (define (loop start)
    (and
     (< start (heap-size))
     (case (heap-ref start)
       [(flat) (loop (+ start 2))]
       [(cons) (loop (+ start 3))]
       [(free) (if (n-free-blocks? start n)
                   start
                   (loop (+ start 1)))]
       [else (error 'find-free-space
                    "unexpected tag ~a" start)])))
  (loop 0))

;; Here is the function you need to write for Q2
(define (free/mark-white!)
  (for ([current-cell (current-heap)]
        [i (in-range 0 (heap-size))])
    (case current-cell
      [(flat) (heap-set! i 'white-flat)]
      [(cons) (heap-set! i 'white-cons)]
      [(white-flat) (heap-set! i 'free) (heap-set! (+ i 1) 'free)]
      [(white-cons) (heap-set! i 'free) (heap-set! (+ i 1) 'free) (heap-set! (+ i 2) 'free)]
      [else (void)])))

(module+ test
  (with-heap (make-vector 7 '?)
    (init-allocator)
    (test (current-heap) (make-vector 7 'free))
    (gc:alloc-flat 'first)
    (test (current-heap) #(flat first free free free free free))
    (gc:alloc-flat 'rest)
    (test (current-heap) #(flat first flat rest free free free))
    (gc:cons 0 2)
    (test (current-heap) #(flat first flat rest cons 0 2))
    (free/mark-white!)
    (test (current-heap) #(white-flat first white-flat rest white-cons 0 2))
    (free/mark-white!)
    (test (current-heap)  (make-vector 7 'free))
    (gc:alloc-flat 'first)
    (gc:alloc-flat 'rest)
    (gc:cons 0 2)
    (free/mark-white!)
    (test (current-heap) #(white-flat first white-flat rest white-cons 0 2))
    )

  (with-heap  (vector 'flat 'first 'flat 'rest 'white-cons 0 2)
    (free/mark-white!)
    (test (current-heap) #(white-flat first white-flat rest free free free))
    (heap-set! 0 'flat)
    (free/mark-white!)
    (test (current-heap) #(white-flat first free free free free free)))

  (with-heap (make-vector 5 #f)
    (init-allocator)
    (gc:cons 0 0)
    (test (current-heap) #(cons 0 0 free free))
    (malloc 2)
    (test/exn (malloc 100) "out of memory")
    (heap-set! 0 'fail)
    (test/exn (malloc 2) "unexpected tag"))
  )

